Uncertainty estimation using a 3D probabilistic U-Net for segmentation with small radiotherapy clinical trial
Phillip Chlap1, Hang Min2, Jason Dowling2
1University of New South Wales, South Western Sydney Clinical School, Sydney, Australia; Ingham Institute for Applied Medical Research, Sydney, Australia; Liverpool and Macarthur Cancer Therapy Centres, Department of Radiation Oncology, Sydney, Australia.
Summary
This study developed a 3D probabilistic U-Net model to estimate uncertainty in medical image segmentation, even with limited data. The model accurately quantifies segmentation confidence, crucial for determining clinical utility in new cases.
Area of Science:
- Medical image analysis
- Machine learning in healthcare
- Radiotherapy planning
Background:
- Medical image segmentation models aim for precise anatomical structure prediction.
- Obtaining perfect ground truth is impossible due to inherent image variability and observer differences.
- Estimating uncertainty (aleatoric and epistemic) is vital for understanding model reliability, especially with limited data.
Purpose of the Study:
- To develop a 3D probabilistic U-Net capable of estimating aleatoric and epistemic uncertainty in medical image segmentation.
- To address the challenge of limited data in training segmentation models.
- To improve the reliability and interpretability of automated segmentation in clinical settings.
Main Methods:
- Utilized a 3D probabilistic U-Net architecture for segmentation.
- Employed an expanded Generalised Evidence Lower Bound (ELBO) with Constrained Optimisation (GECO) and contour loss for focused training on areas of observer disagreement.
- Applied Ensemble and Monte-Carlo Dropout (MCDO) methods for uncertainty quantification during inference.
- Validated the approach on two radiotherapy datasets (TOPGEAR and RAVES) with limited training cases (n=10).
Main Results:
- Achieved Dice Similarity Coefficient (DSC) of 0.7 and Surface DSC (sDSC) of 0.43 for TOPGEAR, and 0.75 (DSC) and 0.71 (sDSC) for RAVES.
- Demonstrated accurate estimation of model confidence by both Ensemble and MCDO methods (p < 0.001).
- Segmentation quality varied across cases, highlighting the importance of uncertainty estimation.
Conclusions:
- Successfully trained auto-segmentation models that can estimate both aleatoric and epistemic uncertainty using limited datasets.
- Model-estimated prediction confidence is essential for assessing the applicability of the model to unseen cases.
- This approach enhances the trustworthiness of automated segmentation in clinical decision-making.
