Source-linked AI summary
Diffusion Models for Implicit Image Segmentation Ensembles
Julia Wolleb, Robin Sandkühler, Florentin Bieder, Philippe Valmaggia, Philippe C. Cattin
TL;DR
Medical lesion segmentation needs image-specific predictions with measurable uncertainty. The paper modifies DDPM training and sampling by conditioning on the image, then exploits stochastic sampling to produce mask ensembles and uncertainty maps; on BRATS2020, ensembling improves segmentation performance while providing detailed uncertainty maps.
Problem
Medical image segmentation requires uncertainty estimates for predictions, especially when results inform treatments such as radiation therapy.
Method
The method trains a DDPM on ground-truth segmentation masks while using the brain MR image as a prior during training and stochastic sampling.
Results
On BRATS2020, ensembles of sampled masks increase segmentation performance across reported scores and produce pixel-wise variance maps showing uncertainty.
Takeaways & Limitations
Stochastic diffusion sampling enables implicit segmentation ensembles and uncertainty maps without training a new model.
Abstract
from arXiv · showhide
Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.
1. Introduction
Medical image segmentation is clinically relevant, but uncertainty in lesion predictions is especially important for treatment decisions. The paper proposes using DDPMs to generate image-specific segmentation masks, uncertainty maps, and implicit ensembles.
- BRATS2020 provides four MR sequences and pixel-wise ground-truth segmentations for brain tumor analysis.The sequences are T1-weighted, T2-weighted, FLAIR, and contrast-enhanced T1-weighted images.
- The proposed DDPM-based method generates uncertainty maps for produced segmentation masks.
- Stochastic DDPM sampling generates different segmentation masks for the same input image and pretrained model.
- Averaging sampled masks forms an implicit ensemble that boosts segmentation performance.
- The work compares its approach with state-of-the-art segmentation algorithms and visually compares variance maps with common uncertainty maps.
2. Method
The method adapts DDPMs to generate image-specific segmentation masks by conditioning noisy masks on brain MR images throughout training and sampling. Stochastic sampling produces multiple masks that support uncertainty estimation and implicit ensembling.
- Diffusion model: DDPM training adds noise across timesteps, then trains a U-Net with MSE loss to predict the preceding less-noisy image.The reverse process learns to predict x_t−1 from x_t, while sampling begins from Gaussian noise.
- Image conditioning: The method concatenates the brain MR image b with the noisy segmentation mask x_b,t so anatomical information conditions the model.Only the ground-truth segmentation receives noise; the image remains an input prior during the modified process.
- Sampling: During inference, the conditioned reverse process outputs a slightly denoised segmentation mask x_b,t−1 with dimension (1, h, w).The sampling procedure takes the original brain MRI as input and returns a predicted segmentation-mask sample.
- Sampling: Because each sampling step is stochastic, repeated sampling for the same brain MR image yields different segmentation masks from one pretrained model.The random component z in the reverse process creates this stochasticity.
- Ensembling and uncertainty: The resulting implicit ensemble supports pixel-wise variance maps for uncertainty visualization and mean-map ensembling to boost segmentation performance.No new model needs to be trained to generate the ensemble.
3. Dataset and Training Details
The evaluation uses preprocessed axial slices from the BRATS2020 brain-tumor dataset, with four-channel MR inputs and specified training and noise-schedule settings.
- Dataset preprocessing: BRATS2020 provides four MR sequences per patient, stacked into four channels and sliced axially for evaluation.The preprocessing excludes the lowest 80 and uppermost 26 slices, normalizes intensities by trimming one percentile at each end, and crops images to (4, 224, 224).
- Dataset split: The training set contains 16,298 images from 332 patients, while the test set contains 1,082 non-empty-segmentation images from 37 patients.No data augmentation is applied.
4. Results and Discussion
The method is evaluated using segmentation metrics, sampled-mask ensembles, and uncertainty maps. Five-sample ensembles improve performance, while variance maps identify uncertain regions.
- Segmentation performance: The evaluation reports Dice, Jaccard index, and 95 percentile Hausdorff Distance (HD95), comparing the method with nnU-Net and SegNet baselines.Five segmentation masks are sampled per test image, averaged, and thresholded at 0.5 to form the ensemble prediction.
- Segmentation performance: An ensemble of five sampled masks increases the method’s performance across the reported segmentation scores.The evaluation also counts empty segmentation masks and reports scores with those cases disregarded in square brackets.
- Uncertainty estimation: Using 100 sampled masks, pixel-wise variance maps visualize segmentation uncertainty, while thresholded mean maps produce ensembled masks.The examples include the input MR image, ground truth, sampled masks, mean map, and variance map.
- Uncertainty estimation: For examples b1, b2, and b3, the ensemble improves performance over the average scores of the 100 individual samples.Table 2 reports segmentation scores for both the ensemble mask and the individual sampled masks.
- Ensemble size and uncertainty comparisons: The Dice score improves with ensemble size for the three examples, with gains appearing by five samples before the curve flattens.The comparison includes variance maps from the proposed method, Bayesian SegNet with Monte Carlo dropout, and SegNet aleatoric uncertainty maps.
5. Conclusion
The paper presents DDPM-based biomedical segmentation that implicitly ensembles stochastic mask samples and derives uncertainty maps from their variance. It reports improved segmentation scores and identifies clinical uncertainty measurement as a key use.
- Conclusion: The proposed DDPM approach implicitly ensembles different segmentation masks for the same brain MR image without training a new model.The ensemble arises from the stochastic sampling process.
- Conclusion: Ensembling the sampled masks increases performance across different segmentation scores.The method also computes uncertainty maps from the variance of the sampled segmentation masks.
- Conclusion: Variance-based uncertainty maps are presented as useful for measuring model-decision uncertainty in clinical applications.Future work will investigate segmentation of the different tumor classes in BRATS2020.
Appendix A. Exemplary Image of BRATS2020
The appendix presents an exemplary BRATS2020 image containing four different MR sequences and the corresponding ground truth segmentation.
- Appendix A. Exemplary Image of BRATS2020: The example includes four different MR sequences alongside the ground truth segmentation.
Appendix B. Implementation Details
The appendix specifies training configurations for SegNet, Bayesian SegNet, nnU-Net, and aleatoric uncertainty estimation.
- Overview: The appendix states that implementation details are provided for the comparing methods.
- SegNet: SegNet uses Adam with learning rate 10^-4, batch size 20, binary cross-entropy loss, and 100 training epochs.
- Bayesian SegNet: Bayesian SegNet adds dropout layers with probability p = 0.5 while retaining the SegNet training schedule.
- nnU-Net: The nnU-Net implementation uses the hyperparameter settings from its official implementation.
- Aleatoric Uncertainty Estimation: Aleatoric uncertainty estimation doubles SegNet’s output channels to produce both a prediction and a variance map.It follows the cited aleatoric loss implementation while retaining SegNet training settings.
Appendix C. Further Examples
The appendix provides additional examples of mean and variance maps generated from 100 sampling runs for three test-set images.
- Mean and variance maps are provided for three additional exemplary test-set images: b4, b5, and b6.
- Figure 7 shows additional produced mean and variance maps from 100 sampling runs.