Conformal confidence sets for biomedical image segmentation

TL;DR

This paper introduces a conformal inference method for biomedical image segmentation that provides spatial uncertainty guarantees, ensuring the true mask is contained within the confidence set with high probability.

Contribution

It adapts conformal inference to image segmentation, developing a calibration procedure that guarantees coverage and improves localization accuracy.

Findings

Confidence sets contain true masks with desired probability.

Using transformed scores improves the tightness of bounds.

Method validated on a polyp tumor dataset.

Abstract

We develop confidence sets which provide spatial uncertainty guarantees for the output of a black-box machine learning model designed for image segmentation. To do so we adapt conformal inference to the imaging setting, obtaining thresholds on a calibration dataset based on the distribution of the maximum of the transformed logit scores within and outside of the ground truth masks. We prove that these confidence sets, when applied to new predictions of the model, are guaranteed to contain the true unknown segmented mask with desired probability. We show that learning appropriate score transformations on a learning dataset before performing calibration is crucial for optimizing performance. We illustrate and validate our approach on a polpys tumor dataset. To do so we obtain the logit scores from a deep neural network trained for polpys segmentation and show that using distance transformed scores to obtain outer confidence sets and the original scores for inner confidence sets enables tight bounds on tumor location whilst controlling the false coverage rate.