SpineSegDiff — Diffusion Models for Spine MRI Segmentation

Diffusion-based framework for automated segmentation of lumbar spine MRI, developed at ETH Zürich’s Biomedical Data Science Lab. It produces multiclass segmentations of both T1- and T2-weighted scans — vertebral bodies, intervertebral discs and the spinal canal — with a focus on imaging in low-back-pain patients, trained and evaluated on the multi-centre SPIDER dataset.
Problem
Manual segmentation of spine MRI is time-consuming and operator-dependent, while reliable, modality-agnostic automated tools are needed to support diagnosis and management of degenerative spine conditions. SpineSegDiff targets multi-modal (T1w and T2w) lumbar spine segmentation and aims to provide not just accurate masks but also a measure of prediction reliability for clinical decision-making.
Method
The approach builds on the DiffUNet framework: a U-shaped denoising network with a dedicated multi-scale image encoder processes the MRI jointly with a partially-noised segmentation mask and learns to denoise toward the correct segmentation. The central contribution is a presegmentation training strategy — a pre-trained nnU-Net generates an initial mask which is then partially noised rather than starting from pure noise. This reduces the required diffusion timesteps from 1000 down to roughly 30 while maintaining accuracy, substantially improving inference efficiency. The diffusion sampling process also enables uncertainty quantification and ensemble prediction.
Key features
- Multi-modal MRI segmentation across T1-weighted and T2-weighted scans.
- Multiclass output covering lumbar spine anatomy (vertebral bodies, intervertebral discs, spinal canal).
- Uncertainty quantification via diffusion sampling, producing heatmaps that highlight anatomical boundaries.
- Ensemble prediction and configurable cross-validation folds.
- Efficient inference using a small number of sampling timesteps (enabled by the presegmentation strategy).
Data & results
Trained and evaluated on the SPIDER dataset — a multi-centre lumbar spine MRI collection (Radboudumc, Jeroen Bosch Hospital, Rijnstate Hospital, Sint Maartenskliniek). SpineSegDiff achieves segmentation performance comparable to the state-of-the-art nnU-Net, with particular strength on degenerated intervertebral discs, while additionally producing uncertainty heatmaps that give clinicians a reliability signal.
Implementation
Implemented in Python with PyTorch and the MONAI medical-imaging framework; training and inference run via command-line scripts with configurable data directories, class counts, timesteps and folds. Published at MIDL 2025 (Monzón, Iff, Konukoglu, Jutzeler — Monzón & Iff co-first). Code released under the Apache 2.0 license, building on the SPIDER dataset, MONAI, and the DiffUNet architecture.
