Improving Uncertainty-Error Correspondence in Deep Bayesian Medical Image Segmentation

Increased usage of automated tools like deep learning in medical image segmentation has alleviated the bottleneck of manual contouring. This has shifted manual labour to quality assessment (QA) of automated contours which involves detecting errors and correcting them. A potential solution to semi-automated QA is to use deep Bayesian uncertainty to recommend potentially erroneous regions, thus reducing time spent on error detection. Previous work has investigated the correspondence between uncertainty and error, however, no work has been done on improving the “utility” of Bayesian uncertainty maps such that it is only present in inaccurate regions and not in the accurate ones. Our work trains the FlipOut model with the Accuracy-vsUncertainty (AvU) loss which promotes uncertainty to be present only in inaccurate regions. We apply this method on datasets of two radiotherapy body sites, namely CT scans of head and neck and prostate MRs. Uncertainty heatmaps (i.e. predictive entropy) are evaluated using ROC and Precision-Recall curves. Numerical results show that when compared to the Bayesian baseline the proposed method successfully suppresses uncertainty for accurate voxels, with similar presence of uncertainty for inaccurate voxels.

Results

Visual Results

Head and Neck CTs	Prostate MRs

To reproduce results

1. This repo allows us to evaluate uncertainty-error correspondence on two radiotherapy body-sites (head and neck CTs and prostate MRs). The data is not included in this repo, but can be downloaded from the following links:

   • Head and neck CTs:
     • RTOG trial (Train/Test)
       • MICCAI2015 Head and Neck Segmentation Challenge
       • DeepMindTCIA (remove scans also present in MICCAI2015)
     • StructSeg (Test)
   • Prostate MRs:
     • ProstateX (Train)
     • Medical Decathlon - Task 5 (Test)
     • PROMISE12 (Test)

2. Use train-hn.py and train-pros.py

3. Evaluation functions can be found in eval-hn.py and eval-pros.py

Closing thoughts

Given that deep learning models have shown the capability of reaching near expert-level performance in medical image segmentation, their next evolution is improved robustness and clinical utility. Our work shows progress on this front for two radiotherapy datasets using a uncertainty-driven loss in a Bayesian setting. We hope the community is inspired by our positive results to further contribute to human-centric approaches to deep learning-based modeling.