Tutorial

Brain Tumor Segmentation on MRI with MONAI and UNet

Difficulty: Intermediate Time: 20 min read

Introduction

Medical imaging data is fundamentally different from standard RGB photos. MRI and CT scans are 3-dimensional volumes stored in specialized formats like NIfTI (.nii.gz) or DICOM. To perform deep learning on these volumes, we use MONAI (Medical Open Network for AI), a PyTorch-based framework specifically designed for healthcare imaging. In this tutorial, we will build a 3D U-Net to automatically segment brain tumors from multi-modal MRI scans using the BraTS (Brain Tumor Segmentation) dataset.

Architecture Overview


graph TD
    Input[4-Channel 3D MRI Volume] --> Enc1[Encoder Block 1]
    
    Enc1 --> Down1((Downsample)) --> Enc2[Encoder Block 2]
    Enc2 --> Down2((Downsample)) --> Enc3[Encoder Block 3]
    Enc3 --> Down3((Downsample)) --> Bottleneck[Bottleneck]
    
    Bottleneck --> Up1((Upsample)) --> Dec1[Decoder Block 1]
    Enc3 -. Skip Connection .-> Dec1
    
    Dec1 --> Up2((Upsample)) --> Dec2[Decoder Block 2]
    Enc2 -. Skip Connection .-> Dec2
    
    Dec2 --> Up3((Upsample)) --> Dec3[Decoder Block 3]
    Enc1 -. Skip Connection .-> Dec3
    
    Dec3 --> Output[4-Channel 3D Segmentation Mask]

Prerequisites

  • Python 3.9+
  • PyTorch, MONAI, nibabel
  • A GPU with at least 12GB VRAM (3D convolutions are memory-intensive)
  • The Medical Decathlon Task01_BrainTumour dataset

Step 1: Install Dependencies

pip install torch torchvision monai nibabel matplotlib

Step 2: Define the MONAI Transform Pipeline

MONAI provides a suite of dictionary-based transforms (ending in d) that handle loading NIfTI files, aligning voxel orientations, normalizing intensities, and applying data augmentation in 3D.

from monai.transforms import (
    Compose, LoadImaged, EnsureChannelFirstd, Spacingd, 
    Orientationd, ScaleIntensityRanged, CropForegroundd, 
    RandCropByPosNegLabeld, RandAffined
)

# Define the training transforms
train_transforms = Compose([
    # Load image and label from NIfTI files
    LoadImaged(keys=["image", "label"]),
    # Ensure shape is [Channels, Depth, Height, Width]
    EnsureChannelFirstd(keys=["image", "label"]),
    # Resample all images to 1x1x1 mm isotropic voxel spacing
    Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
    # Reorient to Right-Anterior-Superior (RAS) axis
    Orientationd(keys=["image", "label"], axcodes="RAS"),
    # Crop away zero-intensity background to save compute
    CropForegroundd(keys=["image", "label"], source_key="image"),
    # Extract 96x96x96 random patches ensuring we hit tumor regions
    RandCropByPosNegLabeld(
        keys=["image", "label"],
        label_key="label",
        spatial_size=(96, 96, 96),
        pos=1, neg=1, num_samples=4
    ),
    # Random 3D rotations for augmentation
    RandAffined(keys=['image', 'label'], prob=0.5, rotate_range=(0.1, 0.1, 0.1))
])

Step 3: Setup Datasets and Dataloaders

We use MONAI's CacheDataset to load everything into memory (if RAM allows) to drastically speed up training, or standard Dataset for lazy loading.

from monai.data import CacheDataset, DataLoader
import os

# Assume we have a list of dictionaries with paths
# data_dicts = [{"image": "img1.nii.gz", "label": "seg1.nii.gz"}, ...]
# train_files, val_files = data_dicts[:80], data_dicts[80:]

train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)

# (Validation transforms and dataloader setup omitted for brevity)

Step 4: The 3D U-Net Model

MONAI provides a heavily optimized 3D U-Net. Because MRI scans often have multiple channels (e.g., T1, T1Gd, T2, FLAIR), the in_channels might be 4. The out_channels represents the number of segmentation classes (e.g., Background, Edema, Enhancing Tumor, Non-Enhancing Core).

from monai.networks.nets import UNet
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = UNet(
    spatial_dims=3,          # 3D Network
    in_channels=4,           # 4 MRI Modalities
    out_channels=4,          # 4 Tumor Classes
    channels=(16, 32, 64, 128, 256), # Number of filters per layer
    strides=(2, 2, 2, 2),    # Downsampling strides
    num_res_units=2          # Use residual blocks
).to(device)

Step 5: Training Loop with Dice Loss

For medical segmentation, standard Cross-Entropy is often inadequate due to massive class imbalance (the tumor is tiny compared to the background). We use DiceLoss instead, which optimizes the intersection over union.

from monai.losses import DiceLoss

loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

max_epochs = 50
for epoch in range(max_epochs):
    model.train()
    epoch_loss = 0
    for batch_data in train_loader:
        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item()
    print(f"Epoch {epoch+1}/{max_epochs}, Loss: {epoch_loss/len(train_loader):.4f}")

Troubleshooting: CUDA Out of Memory (OOM)

Training 3D Convolutional Neural Networks requires massive amounts of VRAM. If you encounter a CUDA out of memory error, try the following steps in order:

  1. Reduce Batch Size: Change batch_size=2 to batch_size=1 in your DataLoader.
  2. Enable Mixed Precision: Use PyTorch's torch.cuda.amp.autocast() to compute gradients in 16-bit precision, cutting memory usage in half.
  3. Reduce Spatial Size: Change the spatial_size in RandCropByPosNegLabeld from (96, 96, 96) to (64, 64, 64).
  4. Gradient Accumulation: If you must use a batch size of 1, accumulate gradients over multiple steps before calling optimizer.step() to simulate a larger batch size.

Conclusion

MONAI abstracts away the immense complexity of handling 3D medical images. By using robust DictionaryTransforms to handle spacing and orientation issues natively within the dataloader, and DiceLoss to handle class imbalance, developers can focus on architecture and deployment rather than writing tedious, error-prone voxel-manipulation code. Once trained, this model can be easily exported via TorchScript and deployed using the Triton Inference Server.