Hierarchical Uncertainty Estimation for Medical Image Segmentation Networks

Xinyu Bai,Wenjia Bai
2023-08-17
Abstract:Learning a medical image segmentation model is an inherently ambiguous task, as uncertainties exist in both images (noise) and manual annotations (human errors and bias) used for model training. To build a trustworthy image segmentation model, it is important to not just evaluate its performance but also estimate the uncertainty of the model prediction. Most state-of-the-art image segmentation networks adopt a hierarchical encoder architecture, extracting image features at multiple resolution levels from fine to coarse. In this work, we leverage this hierarchical image representation and propose a simple yet effective method for estimating uncertainties at multiple levels. The multi-level uncertainties are modelled via the skip-connection module and then sampled to generate an uncertainty map for the predicted image segmentation. We demonstrate that a deep learning segmentation network such as U-net, when implemented with such hierarchical uncertainty estimation module, can achieve a high segmentation performance, while at the same time provide meaningful uncertainty maps that can be used for out-of-distribution detection.
Image and Video Processing,Computer Vision and Pattern Recognition,Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to effectively estimate the uncertainty of model predictions in medical image segmentation tasks. Since medical images may themselves contain noise, and there may be human errors and biases in the manual annotation process, these uncertainties make the learning of medical image segmentation models an inherently ambiguous task. In order to build a reliable image segmentation model, in addition to evaluating its performance, it is also necessary to be able to estimate the uncertainty of model predictions. This paper proposes a simple and effective method to estimate uncertainty at multiple levels by using multi - level image representations. Specifically, the authors propose a framework named VAE - U - net, which provides a meaningful uncertainty map while generating predicted image segmentation by integrating uncertainty modeling in the skip - connection module in the U - net architecture. This method not only improves the segmentation performance of the model but also can provide additional safety guarantees when facing abnormal inputs.