Conformal Prediction for Image Segmentation Using Morphological Prediction Sets

TL;DR

This work addresses uncertainty quantification in binary image segmentation by using conformal prediction to create prediction sets with finite-sample guarantees at a user-defined confidence level $1-\alpha$, relying only on predicted masks (no access to predictor internals). It introduces a novel method that builds nested, dilation-based margins around the predicted mask, using a nonconformity score computed on calibration data to determine the dilation amount $\hat{\lambda}$ and the conformal margin $\mu^{\lambda}$. The approach is predictor-agnostic and applicable to black-box segmentation models, demonstrated on medical-imaging datasets where empirical coverage meets or exceeds the nominal level and margins reflect uncertainty. The results highlight the practical utility of conformal prediction for providing reliable uncertainty bounds in clinical workflows without requiring access to model internals, with future work aiming at multiclass/instance segmentation and adaptive margins.

Abstract

Image segmentation is a challenging task influenced by multiple sources of uncertainty, such as the data labeling process or the sampling of training data. In this paper we focus on binary segmentation and address these challenges using conformal prediction, a family of model- and data-agnostic methods for uncertainty quantification that provide finite-sample theoretical guarantees and applicable to any pretrained predictor. Our approach involves computing nonconformity scores, a type of prediction residual, on held-out calibration data not used during training. We use dilation, one of the fundamental operations in mathematical morphology, to construct a margin added to the borders of predicted segmentation masks. At inference, the predicted set formed by the mask and its margin contains the ground-truth mask with high probability, at a confidence level specified by the user. The size of the margin serves as an indicator of predictive uncertainty for a given model and dataset. We work in a regime of minimal information as we do not require any feedback from the predictor: only the predicted masks are needed for computing the prediction sets. Hence, our method is applicable to any segmentation model, including those based on deep learning; we evaluate our approach on several medical imaging applications.

Figures (2)

• Figure 1: Example: White Blood Cell (WBC) dataset Zheng_2018_WBC, prediction (nucleus) with UniverSeg Butoi_2023_UniverSeg. (a) Input image $X$. (b) Sigmoid scores $\hat{f}(X)$, assumed to be unavailable. (c) ground-truth mask $Y$. (d) predicted mask $\widehat{Y}\xspace$. (e) intersection of $Y$ and $\widehat{Y}\xspace$ in purple (true positives). (f) prediction set $\mathcal{C}_{\lambda}(\widehat{Y}\xspace)$: adding a margin via $\lambda = 6$ dilations of $\widehat{Y}\xspace$ by structuring element $B=$, the missing pixels (e, in red) are covered, as per nonconformity score in Eq. \ref{['eq:score-smallest-margin']}. Colors : true positives; : dilation margin; : false negatives recovered.
• Figure 2: Example: polyps dataset (see Sec. \ref{['sec:experiments']}), prediction with PraNet Fan_2020_pranet. For $\alpha=0.10$ and $\tau = 0.99$: (a) input image, (b) ground-truth mask, (c) predicted mask, (d) prediction set via dilation (Eq. \ref{['eq:set-iter-dilations']}), (e) prediction set via thresholding on sigmoid (as in Angelopoulos_2022_CRC). Pixels in light blue ( ) are the margin: in (d) it contains only pixels contiguous to the prediction (c) while in (e), it does not necessarily do so because of the underlying sigmoid scores (not shown). As shown in Tab. \ref{['tab:dil_vs_threshold_pranet']}, for this model configuration the latter has much larger stretch (Eq. \ref{['eq:stretch']}, lower is better). White pixels represent the background.

Theorems & Definitions (2)

• Theorem 1
• proof