Brain MRI Segmentation with U-Net (MobileNetV2 Encoder)

Lightweight medical image segmentation pipeline (PyTorch + SMP)

  • Developed a lightweight medical image segmentation pipeline on Kaggle’s MRI dataset.
  • Converted 3D NIfTI volumes to 2D slices, applied z-score normalization, and binarized masks (nearest-neighbor).
  • Trained a U-Net with a MobileNetV2 encoder (PyTorch/SMP) using BCE+Dice, mixed precision (AMP), gradient clipping, cosine LR, and early stopping.
  • Evaluated with Dice/IoU/Precision/Recall/F1 and performed a threshold sweep to optimize inference.

Repo: GitHub Dataset: Kaggle

Left: Input Image. Middle: Ground Truth. Right: Segmentation with U-Net (MobileNetV2) prediction.

Highlights

  • Frameworks: PyTorch, segmentation-models-pytorch (SMP)
  • Losses: BCE + Dice
  • Training: AMP, gradient clipping, cosine schedule, early stopping
  • Metrics: Dice, IoU, Precision, Recall, F1

References