FlowSDF: Flow Matching for Medical Image Segmentation Using Distance Transforms

TL;DR

FlowSDF addresses the challenge of robust medical image segmentation by learning a conditional flow that transports a simple prior to a distribution of smooth, implicit segmentation maps expressed as a signed distance function. The method conditions on the input image to generate multiple plausible segmentations and corresponding uncertainty maps, enabling robust predictions and uncertainty-aware decision making. By combining image-guided flow matching with an SDF-based representation, FlowSDF yields smoother boundaries, improved sample diversity, and practical uncertainty estimates, while reducing artefacts common in binary-mask approaches. Empirical results on MoNuSeg and gland segmentation demonstrate competitive performance and potential benefits for clinical decision support through probabilistic segmentation outputs.

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.

Figures (9)

• Figure 1: Overview of FlowSDF, which models a conditional flow $\Psi_t(\widetilde{m}|x)$ by relating a learned conditional vector field $v_{\theta}$ to the conditional probability path $p_t(\widetilde{m}|\widetilde{m}_1,x)$. During inference, a sample $\widetilde{m}_0\sim p_0$ is drawn from a known prior distribution (left) and the is numerically solved to obtain segmentation masks $\widetilde{m}_1$ that resemble those from the conditional distribution $q(\widetilde{m}|x)$ (right).
• Figure 2: Given an image $x$ (left), its binary segmentation mask (center) can be transformed into a truncated mask (right). The zoomed area shows some segmentation objects in detail denoted by $\mathcal{S}$ and their boundaries $\partial \mathcal{S}$ embedded in the domain $\Omega$.
• Figure 3: Comparison of the effect of the transitioning process on the resulting thresholded segmentation masks when using an mask $\widetilde{m}$ and a binary mask $m$ for a given image $x$. Note that the representation allows for a more natural distortion process in its thresholded masks, which evolves along the object boundaries instead of directly introducing "hole"-like structures at random pixel positions as it is the case in the thresholded masks when directly using the binary segmentation mask.
• Figure 4: Schematic of the corruption/generative process (top row for different $t \in [0,1]$) on an segmentation mask $\widetilde{m}$ for a given image $x$. Note that the governing mechanism is bidirectional, hence to obtain samples during inference a numerical integrator is used. The corresponding binary segmentation masks $m_t$ are shown in the second row. While they are not directly used neither during training nor sampling, they can be obtained easily by thresholding and they demonstrate the induced corruption mechanism through the usage of the .
• Figure 5: Exemplary sampled segmentation masks for both data sets. The predicted masks $\widetilde{m}_1$ are directly obtained from the sampling, whereas the thresholded masks $m$ are shown to additionally enable a visual comparison with the depicted ground truth $m_{\mathrm{gt}}$. Furthermore, we provide the condition image with the overlaid thresholded mask $x\otimes m$.