Pixel-wise confidence estimation for segmentation in Bayesian Convolutional Neural Networks

Abstract

Bayesian convolutional neural networks represent an emerging state-of-the-art computer vision framework. This particular version of CNN allows the modeling of two types of uncertainties: epistemic uncertainty and aleatoric uncertainty. These uncertainties permit to evaluate the degree of confidence in predictions yielded by a model. Nevertheless, few studies automatically integrate uncertainty in an end-to-end prediction pipeline. Our research proposes a novel way to assess the degree of confidence in yielded predictions. To that end, a model is trained using the cross-entropy loss function. Thereafter, the learning of this model is resumed with a pixel weighting dynamic calculation designed to reduce the uncertainty of well-classified pixels and penalize wrong classifications. From the epistemic uncertainty measures provided by these two trainings, a histogram is calculated. A final neural network model is used to determine an interval of confidence for the predictions. This interval defines the pixels to be considered with caution at test time, depending on the uncertainty they yield. Validation is performed and shows that the uncertainty yielded by our sample weighting provides a better confidence interval than the regular, unweighted, cross-entropy loss function. Furthermore, our expected calibration error averaged over all datasets (0.07) is lower than available methods (0.1 for the state-of-the-art, 0.18 without calibration). Furthermore, the proposed uncertainty-based thresholding provides better accuracy than baseline uncertainty thresholding, while also minimizing the number of confident errors.