Medical Image Segmentation

Deep learning pipeline for 3D thoracic organ segmentation from CT scans, comparing ENet, CustomNet, and nnU-Net architectures.

September 1, 2024

Medical ImagingDeep LearningPyTorch3D Segmentation

Overview

Group project for AI for Medical Imaging (UvA, part of VU M.Sc. AI program). Evaluated 2D and 3D segmentation architectures for multi-organ medical image analysis on the SegTHOR challenge — segmenting the heart, esophagus, trachea, and aorta from 3D CT scans.

Data Preprocessing

A key challenge was misaligned heart annotations in the ground truth data. For patient 27, two ground truths existed — one correctly placed, one misaligned. We attempted correction using Elastix (rigid and affine transformations) before developing a custom alignment approach for the full dataset.

Architectures

ENet (Baseline): 2D slice-wise segmentation with lightweight encoder-decoder architecture
CustomNet (2D): Custom architecture building on ENet with additional convolutional blocks and skip connections
nnU-Net (3D): State-of-the-art self-configuring framework that automatically adapts architecture, preprocessing, and training to the dataset. Operated on full 3D CT volumes rather than 2D slices.

Training

Combined cross-entropy and Dice loss for handling severe class imbalance between organs
3D data augmentation pipeline with elastic deformations
Trained on HPC cluster with SLURM job scheduling
Evaluated using Dice score, Hausdorff distance, and ASSD per organ class

3D Organ Segmentation by nnU-Net

nnU-Net Segmentation Results

Technologies

Python, PyTorch, nnU-Net, ENet, NIfTI, Elastix, HPC (SLURM)