FlowSDF: Flow Matching for Medical Image Segmentation Using Distance Transforms

TL;DR

FlowSDF introduces a flow matching framework that models segmentation masks as signed distance functions, enabling accurate sampling, uncertainty estimation, and improved robustness in medical image segmentation.

Contribution

It presents a novel SDF-based probabilistic framework for medical image segmentation that improves sampling accuracy and uncertainty quantification.

Findings

Competitive performance on nuclei and gland segmentation datasets

Effective uncertainty estimation through variance maps

Enhanced robustness in segmentation predictions

Abstract

Medical image segmentation plays an important role in accurately identifying and isolating regions of interest within medical images. Generative approaches are particularly effective in modeling the statistical properties of segmentation masks that are closely related to the respective structures. In this work we introduce FlowSDF, an image-guided conditional flow matching framework, designed to represent the signed distance function (SDF), and, in turn, to represent an implicit distribution of segmentation masks. The advantage of leveraging the SDF is a more natural distortion when compared to that of binary masks. Through the learning of a vector field associated with the probability path of conditional SDF distributions, our framework enables accurate sampling of segmentation masks and the computation of relevant statistical measures. This probabilistic approach also facilitates the generation of uncertainty maps represented by the variance, thereby supporting enhanced robustness in prediction and further analysis. We qualitatively and quantitatively illustrate competitive performance of the proposed method on a public nuclei and gland segmentation data set, highlighting its utility in medical image segmentation applications.