A Coding Implementation on MONAI for End-to-End 3D Spleen Segmentation Using UNet on Medical CT Volumes

MONAI, an open-source medical imaging framework, has released a tutorial demonstrating an end-to-end 3D spleen segmentation pipeline using a UNet model on CT volumes from the Medical Segmentation Decathlon Task09 dataset. The implementation applies medical imaging transformations including orientation alignment, voxel-spacing normalization, and patch-based sampling, then trains a 3D UNet with mixed precision and DiceCE loss for binary organ segmentation. The pipeline provides a complete workflow from raw medical volumes through training and validation to qualitative visualization of model predictions against ground-truth masks.

In this tutorial, we build an end-to-end 3D medical image segmentation pipeline using MONAI https://github.com/project-monai/monai to segment the spleen on the Medical Segmentation Decathlon Task09 dataset. We work with volumetric CT scans, apply medical imaging transformations such as orientation alignment, voxel-spacing normalization, intensity windowing, foreground cropping, and patch-based sampling, and then train a 3D UNet model for binary organ segmentation. We also use mixed precision training, DiceCE loss, sliding-window inference, Dice-based validation, and qualitative visualization to understand how the model learns and how its predictions compare with the ground-truth masks. Also, we move from raw medical volumes to a complete train–validate–visualize segmentation system. python pip install -q "monai nibabel,tqdm,matplotlib ==1.5.2" 2 /dev/null import os, time, glob, tempfile, warnings import numpy as np import torch import matplotlib.pyplot as plt from torch.amp import autocast, GradScaler from monai.apps import DecathlonDataset from monai.data import DataLoader, decollate batch from monai.networks.nets import UNet from monai.networks.layers import Norm from monai.losses import DiceCELoss from monai.metrics import DiceMetric from monai.inferers import sliding window inference from monai.utils import set determinism from monai.transforms import Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, Orientationd, Spacingd, ScaleIntensityRanged, CropForegroundd, RandCropByPosNegLabeld, RandFlipd, RandRotate90d, RandShiftIntensityd, AsDiscrete, warnings.filterwarnings "ignore" We start by installing MONAI with the required medical imaging and visualization dependencies. We then import PyTorch, NumPy, Matplotlib, and the main MONAI modules needed for datasets, transforms, model training, metrics, and inference. We also suppress warnings to keep the notebook output clean while we focus on the segmentation workflow. QUICK RUN = True device = torch.device "cuda" if torch.cuda.is available else "cpu" root dir = tempfile.mkdtemp roi size = 96, 96, 96 num samples = 4 batch size = 2 max epochs = 15 if QUICK RUN else 200 val every = 3 train cache = 8 if QUICK RUN else 24 val cache = 2 if QUICK RUN else 6 set determinism seed=0 print f"Device: {device} | epochs: {max epochs} | data dir: {root dir}" train transforms = Compose common + image key="image", image threshold=0 , RandFlipd keys= "image", "label" , prob=0.2, spatial axis=0 , RandFlipd keys= "image", "label" , prob=0.2, spatial axis=1 , RandFlipd keys= "image", "label" , prob=0.2, spatial axis=2 , RandRotate90d keys= "image", "label" , prob=0.2, max k=3 , RandShiftIntensityd keys= "image" , offsets=0.10, prob=0.5 , EnsureTyped keys= "image", "label" , val transforms = Compose common + EnsureTyped keys= "image", "label" We define the main configuration for the tutorial, including the device, dataset directory, patch size, batch size, number of epochs, and cache settings. We then create the preprocessing pipeline for CT volumes by loading images, aligning orientation, resampling voxel spacing, scaling intensities, and cropping the foreground. We also define the training and validation transforms, with the training pipeline including random crops, flips, rotations, and intensity shifts. train ds = DecathlonDataset root dir=root dir, task="Task09 Spleen", section="training", transform=train transforms, download=True, val frac=0.2, cache num=train cache, num workers=2, seed=0 val ds = DecathlonDataset root dir=root dir, task="Task09 Spleen", section="validation", transform=val transforms, download=False, val frac=0.2, cache num=val cache, num workers=2, seed=0 train loader = DataLoader train ds, batch size=batch size, shuffle=True, num workers=2, pin memory=torch.cuda.is available val loader = DataLoader val ds, batch size=1, shuffle=False, num workers=1, pin memory=torch.cuda.is available print f"Train volumes: {len train ds } | Val volumes: {len val ds }" loss fn = DiceCELoss to onehot y=True, softmax=True optimizer = torch.optim.AdamW model.parameters , lr=1e-4, weight decay=1e-5 scheduler = torch.optim.lr scheduler.CosineAnnealingLR optimizer, T max=max epochs scaler = GradScaler "cuda", enabled=torch.cuda.is available dice metric = DiceMetric include background=False, reduction="mean" post pred = Compose AsDiscrete argmax=True, to onehot=2 post label = Compose AsDiscrete to onehot=2 We load the Medical Segmentation Decathlon Task09 Spleen dataset using MONAI’s DecathlonDataset. We split the data into training and validation sections, apply the appropriate transforms, and wrap both datasets with PyTorch-style data loaders. We then create a 3D UNet model, define the DiceCE loss, set up the AdamW optimizer, learning-rate scheduler, mixed-precision scaler, Dice metric, and post-processing steps. best dice, best epoch = -1.0, -1 loss hist, dice hist, dice epochs = , , best path = os.path.join root dir, "best spleen unet.pth" for epoch in range 1, max epochs + 1 : model.train ; epoch loss, t0 = 0.0, time.time for batch in train loader: x, y = batch "image" .to device , batch "label" .to device optimizer.zero grad set to none=True with autocast "cuda", enabled=torch.cuda.is available : logits = model x loss = loss fn logits, y scaler.scale loss .backward scaler.step optimizer ; scaler.update epoch loss += loss.item scheduler.step epoch loss /= len train loader ; loss hist.append epoch loss print f" {epoch:3d}/{max epochs} loss={epoch loss:.4f} " f"lr={scheduler.get last lr 0 :.2e} {time.time -t0:.0f}s " if epoch % val every == 0 or epoch == max epochs: model.eval ; dice metric.reset with torch.no grad : for vb in val loader: vx, vy = vb "image" .to device , vb "label" .to device with autocast "cuda", enabled=torch.cuda.is available : vout = sliding window inference vx, roi size, 4, model, overlap=0.5 vout = post pred o for o in decollate batch vout vlab = post label o for o in decollate batch vy dice metric y pred=vout, y=vlab d = dice metric.aggregate .item dice hist.append d ; dice epochs.append epoch if d best dice: best dice, best epoch = d, epoch torch.save model.state dict , best path print f" val Dice={d:.4f} best={best dice:.4f} @ {best epoch} " print f"\nDone. Best mean Dice {best dice:.4f} at epoch {best epoch}." We run the full training loop, where each epoch trains the 3D UNet on cropped volumetric patches from the spleen dataset. We use automatic mixed precision to reduce memory usage and speed up training when a GPU is available. We also validate the model at regular intervals using sliding-window inference, track the Dice score, and save the best-performing checkpoint. fig, ax = plt.subplots 1, 2, figsize= 12, 4 ax 0 .plot range 1, len loss hist +1 , loss hist, "-o", ms=3 ax 0 .set title="Training loss", xlabel="epoch", ylabel="DiceCE loss" ax 1 .plot dice epochs, dice hist, "-o", color="seagreen", ms=4 ax 1 .set title="Validation mean Dice", xlabel="epoch", ylabel="Dice" ; ax 1 .set ylim 0, 1 plt.tight layout ; plt.show model.load state dict torch.load best path, map location=device ; model.eval with torch.no grad : sample = next iter val loader img = sample "image" .to device with autocast "cuda", enabled=torch.cuda.is available : pred = sliding window inference img, roi size, 4, model, overlap=0.5 pred = torch.argmax pred, dim=1 .cpu .numpy 0 img np, lab np = img.cpu .numpy 0, 0 , sample "label" .numpy 0, 0 z = int np.argmax lab np.sum axis= 0, 1 fig, ax = plt.subplots 1, 3, figsize= 13, 5 ax 0 .imshow img np :, :, z , cmap="gray" ; ax 0 .set title "CT slice" ax 1 .imshow lab np :, :, z , cmap="viridis" ; ax 1 .set title "Ground truth" ax 2 .imshow pred :, :, z , cmap="viridis" ; ax 2 .set title "Prediction" for a in ax: a.axis "off" plt.tight layout ; plt.show We first plot the training loss and validation Dice score to see how the model improves over time. We then reload the best-saved model checkpoint and run inference on a single validation volume using sliding-window prediction. We visualize the CT slice, ground-truth mask, and predicted segmentation side by side to inspect the model’s qualitative performance. In conclusion, we finished a practical MONAI-based workflow for 3D spleen segmentation using a 3D UNet model. We prepared the Medical Segmentation Decathlon dataset, transformed and augmented the CT volumes, trained the model with DiceCE loss, validated it using sliding-window inference, and tracked both loss and Dice score over time. We also inspected the final prediction visually by comparing the CT slice, ground-truth label, and model output side by side. Now, we have a clear understanding of how MONAI supports medical segmentation tasks from data loading and preprocessing to model training, evaluation, checkpointing, and qualitative analysis. Check out the Full Codes with Notebook . Also, feel free to follow us on and don’t forget to join our Twitter https://x.com/intent/follow?screen name=marktechpost and Subscribe to 150k+ ML SubReddit https://www.reddit.com/r/machinelearningnews/ . Wait are you on telegram? our Newsletter https://www.aidevsignals.com/ now you can join us on telegram as well. https://t.me/machinelearningresearchnews Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us https://forms.gle/wbash1wF6efRj8G58 Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions. - Sana Hassan - Sana Hassan - Sana Hassan - Sana Hassan