Automated Segmentation of Lung Nodules with LUNA16 and MONAI

Introduction

Lung cancer is the leading cause of cancer death worldwide. Early detection of pulmonary nodules in CT scans drastically improves survival rates. In this tutorial, we will use the Medical Open Network for AI (MONAI), a PyTorch-based framework, to train a 3D U-Net capable of automatically segmenting lung nodules from CT scans using the public LUNA16 dataset.

Architecture Overview

graph TD
    Raw(Raw CT Scan - NIfTI) --> Transform(MONAI Dictionary Transforms)
    Transform --> HU[Hounsfield Unit Scaling]
    Transform --> Crop[Random 3D Patch Cropping]
    HU --> Dataset(CacheDataset)
    Crop --> Dataset
    
    Dataset --> Loader(DataLoader)
    Loader --> UNet(3D U-Net Model)
    
    UNet --> Pred(Predicted Nodule Mask)
    Pred --> Dice(Dice Loss Calculation)
    Dice --> Backprop(Backpropagation & Optimizer)
    
    UNet -.-> |Evaluation| Slide(Sliding Window Inference)
    Slide -.-> |Stitch Patches| FinalOutput(Full 3D Segmentation)
        

Prerequisites

• Python 3.9+
• PyTorch, MONAI
• A subset of the LUNA16 dataset (CT scans and corresponding segmentation masks)
• A CUDA-capable GPU is highly recommended

Step 1: Install Dependencies

pip install monai torch torchvision nibabel pydicom

Step 2: Defining the Data Transforms

Medical images are notoriously massive and noisy. We use MONAI's dictionary-based transforms to load the NIfTI files, normalize the Hounsfield Units (HU) specific to lung tissue, and randomly crop 3D patches (since we cannot fit a whole CT scan into GPU memory).

import monai
from monai.transforms import (
    Compose, LoadImaged, AddChanneld, Spacingd, 
    Orientationd, ScaleIntensityRanged, RandCropByPosNegLabeld, ToTensord
)

# Define the data pipeline for training
train_transforms = Compose([
    LoadImaged(keys=["image", "label"]),
    AddChanneld(keys=["image", "label"]),
    Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
    Orientationd(keys=["image", "label"], axcodes="RAS"),
    # Clip Hounsfield Units to the lung window (-1000 to 400) and normalize to [0, 1]
    ScaleIntensityRanged(keys=["image"], a_min=-1000.0, a_max=400.0, b_min=0.0, b_max=1.0, clip=True),
    # Randomly crop 96x96x96 patches centered around nodules
    RandCropByPosNegLabeld(keys=["image", "label"], label_key="label", 
                           spatial_size=(96, 96, 96), pos=1, neg=1, 
                           num_samples=4, image_key="image", image_threshold=0),
    ToTensord(keys=["image", "label"]),
])

Step 3: Creating the Dataloaders

We load our dictionaries of file paths into a CacheDataset, which caches the deterministic transforms (loading, spacing, orientation) in RAM for massively faster training epochs.

from monai.data import CacheDataset, DataLoader

# Example list of file paths
train_files = [{"image": "scan1.nii.gz", "label": "mask1.nii.gz"}, {"image": "scan2.nii.gz", "label": "mask2.nii.gz"}]

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

Step 4: Model, Loss, and Training Loop

We instantiate a 3D U-Net architecture. Because nodules are tiny compared to the whole lung (class imbalance), standard Cross Entropy Loss fails. We use the Dice Loss to optimize for intersection over union.

import torch
from monai.networks.nets import UNet
from monai.losses import DiceLoss

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

# Initialize a 3D U-Net
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2, # Background and Nodule
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)

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

# Standard PyTorch Training Loop
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} average loss: {epoch_loss / len(train_loader):.4f}")

Step 5: Sliding Window Inference

Because we trained on small 96x96x96 patches, we cannot simply pass a full 512x512x256 CT scan into the model for inference; we would run out of GPU memory immediately. MONAI provides a sliding_window_inference function that seamlessly sweeps the model across the full image, generates predictions patch-by-patch, and stitches them back together into a final 3D volume, using Gaussian blending at the overlapping seams.

from monai.inferers import sliding_window_inference

model.eval()
with torch.no_grad():
    # Load a full validation image (not randomly cropped)
    val_data = next(iter(val_loader))
    val_inputs = val_data["image"].to(device)
    
    # Run sliding window inference using the same 96x96x96 patch size used during training
    roi_size = (96, 96, 96)
    sw_batch_size = 4
    val_outputs = sliding_window_inference(
        val_inputs, roi_size, sw_batch_size, model, overlap=0.5
    )
    
    # val_outputs now contains the full 3D segmentation mask
    print(f"Final output shape: {val_outputs.shape}")

Conclusion

MONAI radically simplifies the boilerplate associated with 3D medical image deep learning. By providing domain-specific transforms (like Hounsfield Unit scaling and targeted spatial cropping), advanced losses (Dice), and brilliant utilities like sliding window inference, it allows researchers to focus entirely on model architecture and clinical outcomes rather than battling out-of-memory errors and data parsing logic.