Medical Image Segmentation Isn't Like Regular CV
If you've trained segmentation models on COCO or Cityscapes and think medical imaging works the same way, you're in for a rough time. The data is different. The stakes are different. And the regulatory environment adds constraints that don't exist anywhere else in computer vision.
I've spent four years building segmentation pipelines for radiology departments — lung nodule detection on CT, cardiac chamber segmentation on MRI, and retinal vessel segmentation on fundoscopy. Here's what I wish someone had told me before I started.
U-Net Is Still the Baseline
U-Net was published in 2015. It's 2025, and it's still the architecture most medical imaging teams start with. Not because nothing better exists, but because the encoder-decoder structure with skip connections handles the fundamental challenge of medical segmentation: precise boundary delineation at full resolution.
The vanilla U-Net uses a contracting path (conv-pool repeated) and an expanding path (upconv-concat repeated). Skip connections pass high-resolution features from the encoder to the decoder. For binary segmentation tasks (tumor vs. background), this works remarkably well.
import torch
import torch.nn as nn
class DoubleConv(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.conv(x)When to Upgrade from Vanilla U-Net
Attention U-Net adds gating mechanisms to the skip connections. Instead of blindly concatenating encoder features, attention gates learn which spatial regions are relevant. For segmenting small structures — like pancreatic tumors that occupy maybe 2% of the image — this makes a measurable difference. In our pancreas segmentation project, attention gates improved Dice score from 0.71 to 0.76.
U-Net++ densifies the skip connections, creating nested sub-networks between encoder and decoder levels. It's more expensive to train but consistently outperforms vanilla U-Net on multi-organ segmentation tasks where you need accurate boundaries for structures at multiple scales.
nnU-Net deserves special mention. It's not a new architecture — it's a self-configuring framework that automatically selects preprocessing, architecture, and training parameters based on your dataset properties. For teams without deep expertise in medical image segmentation, nnU-Net often outperforms manually tuned architectures because its heuristics encode years of practical experience.
For a related perspective, see Batch Normalization vs Layer Normalization in Production Tra.
Training Data: The Expensive Part
Getting labeled medical images is brutally expensive. A radiologist charges $200-400 per hour. Annotating a single 3D CT volume with voxel-level segmentation labels can take 2-4 hours. A training dataset of 200 volumes costs $80,000-$320,000 in annotation alone. This isn't an abstract concern — it's a budget line item that determines your project scope.
Semi-supervised approaches help stretch limited labels. Our current workflow uses a small fully-labeled set (50-100 volumes) for initial training, then applies the model to unlabeled data and has radiologists correct the predictions rather than annotate from scratch. Correction is 3-5x faster than de novo annotation.
Augmentation for Medical Images
Standard augmentations — flips, rotations, scaling — transfer directly. But medical imaging has its own effective augmentations:
- Elastic deformations simulate anatomical variability between patients. A sigma of 10-15 pixels works well for most abdominal CT tasks.
- Intensity windowing variations account for differences in scanner protocols. CT images use Hounsfield units, and the display window varies by clinical question.
- Bias field simulation mimics the intensity inhomogeneity common in MRI data. Without this, models trained on one scanner often fail on another.
- Cutout and random erasing force the model to use contextual information rather than relying on a single distinctive feature.
One thing that doesn't work: aggressive color jittering. Medical images have physically meaningful intensity values. A CT pixel value of 30 HU means soft tissue. Randomly shifting intensities breaks that physical meaning.
Loss Functions for Imbalanced Tasks
Binary cross-entropy alone won't cut it for most medical segmentation. The target structure is often tiny — a lung nodule might occupy 0.1% of the image. BCE treats every pixel equally, so the model quickly learns to predict "background everywhere" and achieves 99.9% accuracy while being useless.
Dice loss directly optimizes the overlap metric and handles class imbalance naturally. Combining it with BCE — typically as 0.5 * bce + 0.5 * dice_loss — gives you both pixel-level accuracy and region-level overlap optimization.
For a related perspective, see Streaming Feature Computation with Apache Flink for Real-Tim.
def dice_loss(pred, target, smooth=1e-5):
pred = torch.sigmoid(pred)
intersection = (pred * target).sum(dim=(2, 3))
union = pred.sum(dim=(2, 3)) + target.sum(dim=(2, 3))
dice = (2.0 * intersection + smooth) / (union + smooth)
return 1.0 - dice.mean()
def combined_loss(pred, target):
bce = nn.functional.binary_cross_entropy_with_logits(pred, target)
dice = dice_loss(pred, target)
return 0.5 * bce + 0.5 * diceClinical Integration and Validation
Model metrics alone don't determine clinical utility. A Dice score of 0.85 sounds great until a radiologist tells you the model consistently overestimates tumor boundaries by 2mm, which changes the radiation treatment plan. Clinician feedback during development isn't optional — it's the only way to catch errors that metrics miss.
We run structured feedback sessions every two weeks during development. Five radiologists review 20 randomly selected model predictions, rating each on a 5-point scale for clinical acceptability. Predictions rated 1-2 get detailed written feedback explaining the failure mode. This qualitative data drives architecture and training decisions more effectively than chasing Dice score improvements.
Regulatory requirements (FDA 510(k) in the US, CE marking in Europe) add 6-18 months to deployment timelines. Plan for this from the start. Document everything — training data provenance, model architecture decisions, validation methodology, failure mode analysis. The regulatory submission is fundamentally a documentation exercise.
3D Segmentation Challenges
Most medical images are volumetric — CT and MRI scans produce 3D volumes, not 2D images. Processing them adds significant complexity. A standard CT abdomen scan might be 512x512x300 voxels. Loading the full volume into GPU memory at float32 requires 300MB, and the intermediate activations during training consume several times that.
Patch-based training is the standard approach: extract random 3D patches (typically 128x128x64 or 96x96x96 voxels) from the volume, train on patches, and use sliding window inference with overlap at test time. The overlap regions get averaged to produce smooth predictions without boundary artifacts.
from monai.inferers import sliding_window_inference
def predict_volume(model, volume, roi_size=(128, 128, 64), overlap=0.5):
model.eval()
with torch.no_grad():
prediction = sliding_window_inference(
inputs=volume.unsqueeze(0),
roi_size=roi_size,
sw_batch_size=4,
predictor=model,
overlap=overlap,
mode="gaussian",
)
return prediction.squeeze(0)The MONAI library (Medical Open Network for AI) provides these utilities out of the box, along with medical-specific transforms, loss functions, and network architectures. It's built on PyTorch and has become the de facto standard for medical imaging research. I'd recommend it over building custom pipelines unless you have specific requirements MONAI can't handle.
We covered a related topic in Model Monitoring in Production: Detecting Data Drift and Per.
Cross-Validation in Small Medical Datasets
With 50-200 training volumes, a random 80/20 train-test split introduces high variance in evaluation metrics. Five-fold cross-validation gives more reliable performance estimates. For medical imaging, we stratify folds by institution (if multi-site data is available) and pathology severity to ensure each fold represents the full data distribution.
The computational cost is real — 5x training time. For nnU-Net with a 200-volume dataset, that's 5x three days = fifteen days of GPU time. But it's the difference between reporting "Dice = 0.82" and "Dice = 0.82 +/- 0.04 across 5 folds," and that variance estimate matters for clinical decision-making about whether the model is ready for deployment.
External validation on a held-out dataset from a different institution is even more important than cross-validation. A model that performs well on data from three hospitals in the same country may fail on data from a hospital with different scanners, protocols, or patient demographics. We always budget for at least one external validation site, even if getting data sharing agreements takes months.
Deployment Infrastructure for Medical AI
Deploying a medical segmentation model into clinical workflows isn't like deploying a web recommendation model. There's no A/B testing — you can't show different segmentation results to different radiologists and measure which one leads to better outcomes. The deployment is binary: the model either integrates into the clinical workflow or it doesn't.
DICOM integration is the first hurdle. Medical images are stored and transmitted in DICOM format, not JPEG or PNG. Your inference pipeline needs to read DICOM files (pydicom library), extract the pixel data, apply the correct windowing and normalization based on the DICOM metadata, run inference, and write results back as DICOM structured reports or DICOM segmentation objects (DICOM SEG). The DICOM format has hundreds of header fields, and different scanner manufacturers populate them inconsistently.
We deploy medical AI models as DICOM-native services using the DICOMweb protocol. The model runs as a container that receives studies via DICOMweb STOW-RS, processes them, and stores results via DICOMweb STOW-RS back to the PACS (Picture Archiving and Communication System). This avoids custom integration with each hospital's specific IT infrastructure.
Performance requirements are surprisingly relaxed for most radiology applications. A radiologist queues studies for review — they don't need results in milliseconds. Processing a CT volume in 30-60 seconds is perfectly acceptable. The queue-based architecture means you can use a single GPU server to handle hundreds of cases per day without real-time constraints. Turnaround time SLAs are typically "within 15 minutes of study receipt," which is generous by ML infrastructure standards.