UBAD_MRI

UBAD_MRI is an unsupervised brain MRI anomaly detection framework based on latent diffusion models and random-walk mask injection. The method is designed to localize brain abnormalities (e.g. tumors) without using any real anomaly annotations during training.

The pipeline operates on 2D PNG slices extracted from 3D MRI volumes (BraTS-style datasets).

---

1. Method Overview

UBAD-MRI consists of:

• A pretrained medical VAE for mapping 2D MRI slices into a compact latent space
• A latent diffusion UNet trained to reconstruct healthy anatomy
• A Random-Walk Mask Injection strategy to simulate synthetic anomalies
• An anomaly score derived from reconstruction discrepancy

Key characteristics:

• Fully unsupervised (no tumor masks used in training)
• Diffusion performed in latent space for efficiency
• Supports multiple MRI modalities (T1, T2, FLAIR, T1CE)
• Evaluated using Dice, AUROC, and Average Precision (AP)

---

2. Dataset Preparation

2.1 Data Format

The dataset used by UBAD_MRI is stored as PNG images (uint8, grayscale).

Each MRI volume is converted into 2D axial slices and saved as PNG files. The dataloader reads images using PIL.Image.

2.2 Expected Directory Structure

Your --data-root directory must follow this structure:

data_root/
├── train/
│   ├── t2/                     # modality images (PNG)
│   │   ├── <case>-slice_<z>-t2.png
│   ├── brainmask/
│   │   ├── <case>-slice_<z>-brainmask.png
│   └── segmentation/           # empty or ignored during training
│
├── val/
│   ├── t2/
│   ├── brainmask/
│   └── segmentation/
│
└── test/
    ├── t2/
    ├── brainmask/
    └── segmentation/           # ground-truth tumor masks (PNG)

Notes:

• <case> refers to the subject ID.
• train and val contain healthy slices only.
• test contains tumor slices with corresponding segmentation masks.
• Segmentation masks are never used during training, only for evaluation.

---

3. Training

Example training command:

torchrun train-UBAD.py \
  --modality t2 \
  --model UNet_L \
  --image-size 256 \
  --augmentation True \
  --data-root /content/data_brats_t2

Important Arguments --modality: MRI modality (t1, t2, flair, t1ce)

--model: UNet backbone (UNet_XS, UNet_S, UNet_M, UNet_L, UNet_XL)

--image-size: input image resolution (default: 256)

--augmentation: enable data augmentation (Albumentations)

--data-root: root directory of the prepared dataset

Training logs and checkpoints are automatically saved to the experiment folder.

4. Evaluation Example evaluation command:

torchrun evaluate-UBAD.py \
  --model-path /content/UBAD_MRI/UBAD_t2_UNet_L/001-UNet_L/checkpoints/last.pt \
  --data-root /content/data_brats_t2

During evaluation, the pipeline:

Encodes test images into latent space using the VAE

Applies DDIM sampling to reconstruct healthy anatomy

Decodes reconstructed latents back to image space

Computes anomaly maps from reconstruction error

Evaluates performance using:

AUROC

Average Precision (AP)

Dice score

Visualization results (input / reconstruction / ground truth / anomaly map) are saved automatically.

5. MRI Modalities The framework supports training and evaluation on individual modalities:

T1

T2

FLAIR

Each modality should be trained independently by specifying --modality and pointing --data-root to the corresponding dataset.