Uncertainty learning improves segmentation accuracy at organ boundaries
Contrastive uncertainty learning for robust multi-structure segmentation in lower-pelvic MR
Computer Vision and Pattern RecognitionMachine Learning
Summary
Segmenting different organs in prostate MRI scans can be tricky because boundaries between tissues are often blurry or unclear, causing mistakes. This paper presents a new method that uses the model’s own uncertainty to help it learn better features, especially around these difficult edges. By encouraging the model to focus on clear examples and separate uncertain or wrong predictions in its internal representation, it reduces errors near boundaries. The method was tested on almost 600 scans, showing similar overall accuracy but much better precision at organ surfaces, reducing risky mistakes.
What this means in practice
- •For medical imaging developers: Improve automated organ delineation in MRI systems by integrating uncertainty-guided feature learning to reduce boundary errors in prostate imaging.
- •For robotic surgery teams: Use enhanced MRI segmentation with uncertainty learning to provide more reliable organ boundary maps for guidance during minimally invasive procedures.
Authors
Bilal Sidiqi, Yipei Wang, Wen Yan, Shaheer U. Saeed, Daniel C. Alexander, Yipeng Hu
Abstract
Boundary delineation of anatomical structures in prostate MRI is often compromised by regional ambiguity. Automated segmentation methods relying on global overlap loss functions tend to be insensitive to surface deviations, leading to significant errors in ambiguous regions. Most existing uncertainty-aware methods treat uncertainty as a post-hoc quality indicator or prediction filter. We propose Contrastive Uncertainty Learning (CUL), a supervised framework that uses predictive uncertainty to directly guide representation learning. Using Monte Carlo Dropout, CUL identifies voxel-wise uncertainty and structures the embedding geometry accordingly: low-uncertainty, correctly predicted features form stable voxels while high-uncertainty or misclassified features are actively repelled, preventing feature collapse at difficult boundaries. We evaluate CUL on a 589-scan 3D prostate MRI multi-structure segmentation task (8 anatomical structures) using a ViT with UperNet decoder architecture. Average Dice ($83.38\%$) remains comparable to the baseline ($83.35\%$), while surface distance decreases significantly ($p<0.001$) across 6 of 8 anatomical structures. A hard-case analysis reveals a $22{\times}$ asymmetry between significantly improved and degraded cases ($18.8\%$ vs $0.9\%$). Ablation studies show that removing uncertainty-guided sample selection eliminates improvements for the transition zone, while replacing the contrastive objective with a boundary loss additionally removes bladder improvements, demonstrating that both uncertainty-guided sampling and the contrastive objective are essential for the observed gains. These findings demonstrate that the proposed uncertainty-guided representation learning is a useful mechanism for reducing the risk of extreme errors in safety-critical medical segmentation