Segmentation Confidence for Arbitrary CNNs

Introduction

Deep learning models, and in particular CNNs, are widely used in medical image segmentation [1, 2, 3, 4, 5, 6, 7]. However, these models often provide little insight into their decision process [8, 9, 10], and they lack an inherent measure of confidence. The produced probability maps provide an indirect signal of confidence per voxel. However, there is no obvious way to produce a single measure of confidence for the entire image. Since CNNs are highly successful at segmentation, the probabilities are usually higher inside and around the ROI, and mostly 0 further from the ROI. This results in a sharp transition between the foreground and background, rather than a gradual increase, limiting the ability to use volume-wide statistics as a measure of confidence. While classification models output probabilities per class, which can estimate confidence, this is not available for segmentation.

There has been research to overcome this shortcoming [11], often with Bayesian frameworks. Gal and Ghahramani [12] have shown that CNNs trained with dropout can approximate Bayesian inference by keeping dropout active during test-time and performing multiple forward passes for classification. Ghoshal et al. [13] propose Monte Carlo drop weights, and decompose predictive probabilities into aleatoric and

epistemic uncertainties. Teye et al. [14] use batch normalization layers instead of dropout. Wang et al. [15] use test-time augmentation and Monte Carlo simulation. There have also been non-Bayesian approaches. Some methods [16, 17] train an ensemble of models to estimate classification uncertainty. Kohl et al. [18] propose a probabilistic U-Net. Yao et al. [19] combine an aleatoric uncertainty-based correction strategy with semi-automated prompt augmentation. Most of these approaches require dedicated architectures [20] or training an ensemble of models [21] to accomplish segmentation and confidence/uncertainty estimation simultaneously. Many of them produce an uncertainty map, but do not deduce a single score for the entire image. This limits their usability in typical medical image analysis workflows. Notably, Wang et al. [15] use the volume variation coefficient (VVC) of multiple forward passes to estimate structure-wise uncertainty.

While CNNs perform well overall, a poor prediction for a particular image may lead to mistakes. Flagging uncertain predictions could provide alternatives, such as obtaining a repeat scan if there is an image quality issue. Another use case is flagging uncertain outputs in large-scale studies for visual quality control, manual correction, or exclusion from downstream analysis. Confidence scores can also be useful in active learning [22, 23] for prioritizing low-confidence images for annotation, as these are likely not well represented in the training dataset. Finally, a model’s ability to report its own confidence is a significant step towards explainable AI.

In this study, we use two alternatives (§ 2.2) to produce multiple predictions from a single input image, which can be readily implemented as an add-on to existing segmentation pipelines. We then explore two approaches to estimate a voxel-based uncertainty map (§ 2.3), and propose a novel method to estimate a single, image-wide confidence score (§ 2.4). We evaluate our method on placenta (ultrasound) and hippocampus (MRI) datasets to test its generalizability.

Methods

Data and Segmentation Models

Placenta 3D ultrasound (3DUS) dataset: This dataset contains $N = 210$ images from 121 unique subjects in the first trimester, and $N = 49$ images from 41 unique subjects taken during the second trimester. Placentas were manually segmented using ITK-SNAP (www.itksnap.org). Images were down-sampled to $128 \times128 \times128$ voxels before training.

Automated quantification of the placenta from 3DUS data is challenging [24], and quick identification of sub-optimal segmentations during clinical screening can allow the option of a repeat scan while the patient is still in the clinic.

Hippocampus MRI dataset: We used the hippocampus task from the Medical Segmentation Decathlon [25]. We used the $N = 260$ training set with anterior and posterior hippocampus labels. Images were up-sampled to $64^3$ voxels.

Segmentation: We employed typical 2D and 3D U-Net models with 2 levels of down / up-sampling, and skip-connections. All convolutional layers have $3^2$ ($3^3$) pixel (voxel) kernels. The placenta dataset has a single ROI. The hippocampus dataset has 2 ROIs (anterior/posterior), and we trained a separate model per ROI. We used binary cross-entropy loss.

For our 2D models for both the placenta and the hippocampus, we observed that a single model can be successfully trained on sagittal and coronal slices, since the general shape and position of the ROI are roughly similar along these axes. Slices along the axial plane differ considerably from the other 2 planes, and inclusion in training significantly lowered performance. Therefore, our final 2D models were trained on sagittal and coronal slices. The placenta model has 364K trainable parameters, while the two 2D hippocampus models each have 162K. We also trained 3D models for each hippocampus ROI (1,938K trainable parameters each). We used 4-fold cross-validation. Since our focus is not on segmentation accuracy, but rather on the performance of confidence scoring on the test set, we used 1 partition for training and 3 for testing within each fold to provide more test cases, rather than the usual 3:1 train/test split. The random partitioning for placentas was based on the unique subjects to avoid leakage.

Test-time augmentation strategies

Test-time flip augmentations: We used test-time augmentation with horizontal/vertical flips for the 2D models. We thus applied a combination of 4 augmentations per 2D slice. As our models were trained on both the sagittal and coronal planes, we produced 8 full-volume 3D outputs per input.

This approach was viable for our 2D models, since the random online augmentations during training included the horizontal/vertical flips. Specifically, for training the placenta model, we used horizontal/vertical flips, in-plane rotations, cropping, blurring, and gamma adjustment. For the hippocampus models, we only applied horizontal/vertical flips. It is common practice to utilize random augmentations like these during training, and modifying the test-time loop to include the flip augmentations and produce multiple outputs per input is a straightforward approach to apply our confidence scoring method to existing segmentation pipelines.

Test-time noise augmentations: Some models may perform worse with test-time flips, e.g., if augmentations were not used during training, or if the ROI is highly asymmetric. To test our method’s applicability in such cases, we trained two 3D models for the two hippocampus ROIs without any augmentations. At test-time, we added random Gaussian noise ($\mathit{torch.randn} \times0.5$) to a random 20% subset of the voxels. We blurred ($\sigma= 1$) the noise before adding to the original image for a more realistic look. For each input, we produced 15 noisy versions, resulting in 16 outputs. While other noise distributions may be more suitable for various modalities, this general approach is applicable to any existing 2D or 3D models, whether augmentations were used during training or not.

Uncertainty estimation strategies

Standard deviation (stdev): First, we simply calculate the voxel-wise stddev of the multiple probability maps. We normalize to $[0, 1]$ via min-max scaling. This gives us an approximate uncertainty map $u(x)$, in which values close to 0.0 denote voxels where the multiple predictions mostly agree, and values close to 1.0 denote the most disagreement. As expected, we usually observed highest uncertainty along the boundary of the ROI (Figure 1c), as well as in regions with artifacts or other noise. We used this approach in our 2D models.

Entropy: When the multiple predictions all have a probability close to 0.5 for a given voxel, since there is agreement, the low stdev registers as a low uncertainty voxel, even though the 0.5 probabilities suggest high uncertainty. To remedy this, we also calculated the binary entropy $H(p) = -p \ln p - (1 - p) \ln(1 - p)$ for all the outputs and took their voxel-wise mean. In our 3D models, we min-max scaled this mean entropy map, and averaged it with the stdev map $u(x)$ (Figure 2).

Confidence score estimation

We use the voxel-wise mean of the multiple output probability maps to create a final probability map $p(x)$ and threshold it to create a binary mask $M_0(x) = \{p(x) \ge0.50\}$ (Figure 1d). Next, we calculated a base mask volume $V_0 = |M_0(x)|$. This volume is independent of the uncertainty $u(x)$, since it includes all voxels that have $p(x) \ge0.50$ regardless of uncertainty. If we introduce an uncertainty constraint such that only voxels under a given uncertainty are allowed, e.g., $u(x) \le0.9$, then the segmentation would be more conservative and thus the resultant mask volume would be lower than $V_0$. We observe that if the overall model prediction for a given image $I(x)$ was highly confident, there would be very

few voxels with a high uncertainty $u(x)$. Conversely, in a low-confidence prediction, there would be high uncertainty in substantially more voxels. This leads to the intuition that if the mask volume decreases rapidly as we tighten the uncertainty constraint, then the confidence for the prediction as a whole is likely low. In contrast, a very gradual decrease in volume suggests low uncertainty and high confidence (Figure 3a, 3b).

To quantify this intuitive concept, we introduced $N_u$ uncertainty constraints decreased by uniform steps $\Delta_u = 1/N_u$, and calculated $N_u - 1$ more masks similarly to $M_0$. In our experiments, we empirically used $N_u = 10$ and therefore $\Delta_u = 0.1$. Thus, for example, $M_1$ is the set of voxels in the segmentation that have a corresponding uncertainty of less than 0.90 ($p(x) \ge0.50 \land u(x) \le0.90$), and $M_2$ only includes voxels in the mask that have an uncertainty less than 0.80. In general, for $i \in[0, N_u - 1]$, we have $M_i(x) = \{p(x) \ge0.5 \land u(x) \le(1 - i * \Delta_u)\}$, with volume $V_i = |M_i(x)|$.

As mentioned above, $M_0$ is the largest mask and the volumes decrease through $V_0$ to $V_{N_u-1}$. How these volumes change is informative (Figure 3a vs. 3b). To capture this, we calculated the $N_u - 1$ ratios of each volume to the base volume, $V_i/V_0$, for $i \in[1, N_u - 1]$. We also calculated the $N_u - 2$ ratios of consecutive steps, $V_{i+1}/V_i$, for $i \in[1, N_u - 2]$. Since the volumes are monotonically decreasing from $V_0$ through $V_{N_u-1}$, all of these ratios are in the $[0, 1]$ range.

$V_i$ is a measure of overall certainty (Figure 3c). Dividing by $V_0$ normalizes for object size. We note that simply averaging $u(x)$ across $V_0$ would be sensitive to object size, without a straightforward way to compare different subjects. Our formulation transforms the uncertainty map $u(x)$ to a series of volumetric measures $V_i$, allowing us to compare different subjects by simply expressing the volumes as the ratio $V_i/V_0$. The $\Delta_u$ determines the granularity of this transformation, and we found 0.1 to be adequate for our purposes. In contrast, $V_{i+1}/V_i$ is a measure of speed of uncertainty change (Figure 3d). The blue curve locally has slower change than the red curve, thus a smaller $V_{i+1}/V_i$. To obtain a single confidence score $C$ from these ratios, we simply took their mean, as a balanced and robust indicator:

Correlation with Dice and outlier detection

We examined the correlation of our confidence $C$ with the segmentation Dice. We also used these scores to identify poor segmentations. To do this in unsupervised manner, we applied an elliptic envelope model to the confidence scores (without taking Dice as input), fitting an ellipsoid around the central cluster of the data points, to separate inliers from outliers. We used the scikit-learn implementation of the model, with a contamination value of 0.1. Since we were interested only in the abnormally low confidence cases, we filtered out cases with above mean confidence (i.e., the highest confidence scores) from the outlier set. We then performed a two-tailed t-test to assess Dice differences between inliers and outliers.

Using confidence to select annotation candidates

We also explored how confidence scores can be used to select images to annotate and add to the training set from a pool of unlabeled candidates. We trained a base 2D model (Sec. 2.1) on 21 random images from the $N = 121$ placenta dataset. We split the remaining images into $N = 60$ test set, and $N = 40$ annotation candidates. Using the base model, we produced segmentations and confidence scores for the $N = 40$ candidates. We compared 5 strategies for selecting candidates: best Dice, worst Dice, best confidence, worst confidence, and random. We note that the best/worst Dice strategies were included as oracle baselines, and would not be available in practice, while confidence can be computed with the base model. We added the images selected by each strategy to the $N = 21$ base set in increments of 5 images, and trained new models. Since the number of images in the training set were limited, we trained 3 models for each scenario with different random seeds. We compared performance on the $N = 60$ test set.

Results and discussion

Compared methods: We compare our confidence score to the volume variation coefficient (VVC) method [15] for estimating a volume-wide confidence. Other methods noted in the Introduction do not produce a volume-wide score.

2D models: For the 2D models, we used flip augmentations during test-time, and estimated uncertainty with standard deviation (Figure 4). For the placenta, the mean Dice was between 0.81-0.82 across the 4 folds. Using confidence $C$, the mean

Dice difference between inliers-outliers were in the range 0.14-0.23. We note that the outlier detection only uses $C$, and the t-test is between the associated Dice scores rather than $C$ (which would be lower in the outlier group, by construction). The inlier-outlier Dice difference was 0.07-0.10 with VVC. For anterior hippocampus, mean Dice was between 0.84-0.85. Inlier-outlier Dice difference was 0.09-0.11 for $C$ and 0.04-0.08 for VVC. For posterior hippocampus, the mean Dice was between 0.82-0.84. Inlier-outlier Dice difference was 0.08-0.21 for $C$ and 0.05-0.21 for VVC.

3D models: Here, we used Gaussian noise augmentation, and estimated uncertainty with stdev and entropy. For anterior hippocampus (not shown for brevity), the mean Dice was between 0.83-0.84. Inlier-outlier Dice difference was 0.05-0.07 with confidence, and 0.05-0.07 with VVC. For posterior hippocampus (Figure 5), the mean Dice was between 0.79-0.81. Inlier-outlier Dice difference was 0.08-0.14 with confidence, and 0.07-0.14 with VVC. The inlier-outlier differences were significant for each fold ($p < 0.05$) for both $C$ and VVC, in both 2D and 3D models.

Selection by confidence: Figure 6 shows the mean test performance for each selection strategy, where $N = 0$ is the base model. Adding images with the worst Dice/confidence scores improves test performance faster than other strategies. The Dice scores of the “worst confidence” strategy were significantly better than the “random” strategy for the smaller $N$

($p < 0.01$ for $N \in\{5, \ldots, 20\}$, and $p < 0.03$ for $N = 25$ in paired t-tests). The performances converge, as expected, as we approach the entire candidate pool size, with the random strategy performing better at $N = 30$ ($p < 0.03$).

Discussion: We proposed a novel method to estimate a single confidence score for CNNs used in medical image segmentation, by examining how rapidly the segmentation volume decreases as we gradually allow less uncertainty.

Our confidence score $C$ is highly correlated with the segmentation Dice. We observe that the upper left quadrant of the plots in Figure 4 and Figure 5 is typically empty, i.e., $C$ is rarely high for low-Dice predictions in our experiments. This suggests high sensitivity, which is useful for practical workflows. $C$ values are usually in the range of 0.60-0.90, and most cases with low Dice receive a $C$ below 0.75. The unsupervised outlier detection experiment confirms this finding. For each ROI in Figure 4, VVC fails to flag several low-Dice ($< 0.6$) cases that our method successfully captures (red arrows). Compared to VVC, our method has more consistent performance for a given task: e.g., for the anterior hippocampus (Figure 4, middle row), the VVC correlations fluctuate from -0.247 to -0.668, whereas our correlations are stable in the [0.699, 0.711] range. Our correlations and inlier-outlier differences substantially surpass VVC counterparts in most experiments.

In our experiments we used test-time flips and Gaussian noise. These or other suitable augmentations may be used for new tasks. More extensive evaluation remains as future work.

Our proposed confidence score $C$ shows a strong correlation with Dice in our extensive experiments, reliably indicating segmentation quality. Our method’s compatibility with existing segmentation pipelines and its applicability without requiring specialized architectures make it a versatile tool. This technique can help in many scenarios such as identifying poor-quality inputs or for selecting images for manual segmentation. Overall, our approach contributes a valuable layer of uncertainty estimation and quality control, enhancing the potential for increased clinical adoption of CNNs. Application in clinic remains as future work. Our code is available at github.com/MedICL-VV/segmentation-confidence.

Acknowledgments: This work was supported, in part, by NIH R01HD109739, U01-HD087180, R03-HD-069742, and Penn Presbyterian George L. and Emily McMichael Harrison Fund for Research in Obstetrics and Gynecology.