Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks

Abstract

Despite the state-of-the-art performance for medical image segmentation, deep convolutional neural networks (CNNs) have rarely provided uncertainty estimations regarding their segmentation outputs, e.g., model (epistemic) and image-based (aleatoric) uncertainties. In this work, we analyze these different types of uncertainties for CNN-based 2D and 3D medical image segmentation tasks at both pixel level and structure level. We additionally propose a test-time augmentation-based aleatoric uncertainty to analyze the effect of different transformations of the input image on the segmentation output. Test-time augmentation has been previously used to improve segmentation accuracy, yet not been formulated in a consistent mathematical framework. Hence, we also propose a theoretical formulation of test-time augmentation, where a distribution of the prediction is estimated by Monte Carlo simulation with prior distributions of parameters in an image acquisition model that involves image transformations and noise. We compare and combine our proposed aleatoric uncertainty with model uncertainty. Experiments with segmentation of fetal brains and brain tumors from 2D and 3D Magnetic Resonance Images (MRI) showed that 1) the test-time augmentation-based aleatoric uncertainty provides a better uncertainty estimation than calculating the test-time dropout-based model uncertainty alone and helps to reduce overconfident incorrect predictions, and 2) our test-time augmentation outperforms a single-prediction baseline and dropout-based multiple predictions.

Automatic Brain Tumor Segmentation Using Convolutional Neural Networks with Test-Time Augmentation

“…In this paper, we extend the work of [28] and [27], and apply test-time augmentation to automatic multi-class brain tumor segmentation. For a given input image, instead of obtaining a single inference, we augment the input image with different transformation parameters to obtain multiple predictions from the input, with the same network and associated trained weights.…”

“…In [14], test images were augmented by mirroring for brain tumor segmentation. In [27], a mathematical formulation was proposed for test-time augmentation, where a distribution of the prediction was estimated by Monte Carlo simulation with prior distributions of parameters in an image acquisition model. That work also proposed a test-time augmentation-based aleatoric uncertainty estimation method that can help to reduce overconfident predictions.…”

“…That work also proposed a test-time augmentation-based aleatoric uncertainty estimation method that can help to reduce overconfident predictions. The framework in [27] has been validated with binary segmentation tasks, while its application to multi-class segmentation has yet to be demonstrated.…”

Accuracy, uncertainty, and adaptability of automatic myocardial ASL segmentation using deep CNN

“…The control/labeled pulse and the bSSFP image acquisition were cardiac triggered to occur at mid-diastole in consecutive heartbeats. bSSFP parameters: TR/TE = 3.2/1.5 ms, flip angle = 50 0 , slice thickness = 10 mm, matrix size = 96x96, and parallel acceleration factor of 2 for SENSE (36) MC dropout has been applied to evaluate model uncertainty in semantic segmentation tasks in both computer vison and medical imaging applications (20,40). In these studies, the typical output of the model uncertainty is a standard deviation pixelby-pixel map that provides spatial information detailing where, within the image, the model is uncertain.…”

As Easy as 1, 2...4? Uncertainty in Counting Tasks for Medical Imaging

“…However, these parametric methods are restricted to unimodal, symmetric distributions, which are not necessarily realistic. Test-time augmentation has been used to perturb the data and thus infer the uncertainty from the differences in predictions [6,7]. In these approaches, the estimated uncertainty will depend wholly on the model's lack of invariance to the chosen augmentations: this may suggest that the models are undertrained or lacking capacity.…”