Conformal Prediction for Image Segmentation Using Morphological Prediction Sets

2025-03-07

AI Summary

In image segmentation—particularly binary segmentation—multiple sources of uncertainty (e.g., annotation ambiguity and sampling variability) necessitate reliable, pixel-wise confidence quantification. To address this, we propose a morphological conformal prediction framework: for the first time, we incorporate morphological dilation into conformal prediction, constructing pixel-level prediction sets with interpretable margins directly from pretrained model outputs (i.e., raw segmentation masks). Our method estimates nonconformity scores and calibrates margins using a calibration set, requiring no model retraining, fine-tuning, or additional supervision. Crucially, it provides rigorous finite-sample coverage guarantees—ensuring the ground-truth mask lies within the predicted set with user-specified confidence. Evaluated across multiple medical imaging benchmarks, our approach empirically validates theoretical coverage while yielding margins that meaningfully reflect local uncertainty, thereby substantially enhancing both reliability and interpretability of segmentation results.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Segmentation
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty

Application Category

• Security and Privacy: Data transparency and provenance
• User Modeling, Personalization and Recommendation: Accountability, Transparency, and Ethics for personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Address uncertainty in binary image segmentation using conformal prediction.
• Construct prediction sets with morphological margins for uncertainty quantification.
• Apply method to any segmentation model without requiring predictor feedback.

Innovation

• Uses conformal prediction for uncertainty quantification
• Applies morphological dilation to segmentation masks
• Requires only predicted masks for prediction sets