Estimating Predictive Uncertainty in Gastrointestinal Image Segmentation

Abstract

Deep learning models are known to achieve state-of-the-art performance in numerous applications. Applying these models to the medical field can relieve demanding workloads and decrease the number of observational oversights in diagnostics. Despite this, deep learning models are considered to exhibit a “black box” nature due to their complex structure and lack of transparency, interpretability, and explainability to how they arrived at a decision. These attributes are crucial to establish trust and reliability between the users and model. We investigated the use of computer-aided detection of polyps in the gastrointestinal tract using segmentation models. To increase explainability behind the model prediction, we adopted two existing uncertainty estimation methods, Monte Carlo (MC) dropout and deep ensembles. We further explored these using two state-of-the-art deep learning architectures, U-Net and ResUNet++, trained with two different loss metrics, binary cross-entropy and dice similarity coefficient (DSC). The uncertainty estimates were visualized as heatmaps, showing spatial uncertainties for the predicted segmentation mask. Our results show that deep ensembles provide informative uncertainty representations that connect large uncertainties to misclassified pixels. Additionally, we established a correlation between the large uncertainty estimates and the corresponding pixels that are likely to be misclassified. MC dropout was insufficient at providing such information to its uncertainty representations. Out of all the combinations tested, our results show that using the U-Net based deep ensemble trained with the DSC metric gave the overall highest score in terms of the mean DSC during test time. Deep ensembles were also able to increase this test score when increasing the ensemble size, giving a DSC score equal to 0.8172 using an ensemble size of 16. We conclude that uncertainty estimation, using deep ensembles, can improve the understanding of deep learning models that can aid users to make informed decisions on whether to trust a model’s predictions.