FlowSDF: Flow Matching for Medical Image Segmentation Using Distance Transforms

2024-05-28

arXiv.org

AI Summary

To address imprecise mask representation and weak uncertainty estimation in medical image segmentation, this paper proposes the first image-guided conditional flow matching framework based on signed distance functions (SDFs). The method models segmentation masks as implicit probability distributions, leveraging the continuous geometric prior of SDFs to enhance deformation modeling naturalness and prediction robustness. By directly regressing vector fields from noisy SDFs to target SDFs, it natively generates pixel-wise variance-based uncertainty maps, enabling high-fidelity sampling and statistical analysis. Evaluated on public nuclear and gland segmentation benchmarks, our approach significantly outperforms existing state-of-the-art methods. Comprehensive qualitative and quantitative experiments demonstrate superior segmentation accuracy, robustness to anatomical variability and noise, and effective uncertainty quantification—particularly critical for clinical decision support and model calibration.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Segmentation
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Medical image segmentation accuracy
• Representing signed distance function
• Generating uncertainty maps for robustness

Innovation

• Flow matching framework
• Signed distance function
• Uncertainty maps generation