Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Encoding Structural Constraints into Segment Anything Models via Probabilistic Graphical Models

Yu Li, Da Chang, Xi Xiao

TL;DR

KG-SAM addresses the lack of anatomical topology and uncertainty in Segment Anything Model (SAM) for medical imaging by integrating a medical knowledge graph with an energy-based CRF and an uncertainty-aware fusion mechanism. The method combines SAM-derived image embeddings with explicit anatomical priors, formulating an energy function $E(\mathbf{M})$ that includes unary, pairwise, and anatomical potentials, and uses mean-field inference to produce refined segmentations and an uncertainty map. Experimental results on multi-center abdominal and prostate datasets show substantial gains, with Dice scores of $82.69\%$ on prostate and $78.05\%$ (MRI) / $79.68\%$ (CT) for abdomen; KG-De-SAM further reaches $84.52\%$, validating the modularity and effectiveness of knowledge-guided refinement. The framework is designed to be compatible with other SAM variants, enabling broader applicability to clinical segmentation tasks.

Abstract

While the Segment Anything Model (SAM) has achieved remarkable success in image segmentation, its direct application to medical imaging remains hindered by fundamental challenges, including ambiguous boundaries, insufficient modeling of anatomical relationships, and the absence of uncertainty quantification. To address these limitations, we introduce KG-SAM, a knowledge-guided framework that synergistically integrates anatomical priors with boundary refinement and uncertainty estimation. Specifically, KG-SAM incorporates (i) a medical knowledge graph to encode fine-grained anatomical relationships, (ii) an energy-based Conditional Random Field (CRF) to enforce anatomically consistent predictions, and (iii) an uncertainty-aware fusion module to enhance reliability in high-stakes clinical scenarios. Extensive experiments across multi-center medical datasets demonstrate the effectiveness of our approach: KG-SAM achieves an average Dice score of 82.69% on prostate segmentation and delivers substantial gains in abdominal segmentation, reaching 78.05% on MRI and 79.68% on CT. These results establish KG-SAM as a robust and generalizable framework for advancing medical image segmentation.

Encoding Structural Constraints into Segment Anything Models via Probabilistic Graphical Models

TL;DR

KG-SAM addresses the lack of anatomical topology and uncertainty in Segment Anything Model (SAM) for medical imaging by integrating a medical knowledge graph with an energy-based CRF and an uncertainty-aware fusion mechanism. The method combines SAM-derived image embeddings with explicit anatomical priors, formulating an energy function that includes unary, pairwise, and anatomical potentials, and uses mean-field inference to produce refined segmentations and an uncertainty map. Experimental results on multi-center abdominal and prostate datasets show substantial gains, with Dice scores of on prostate and (MRI) / (CT) for abdomen; KG-De-SAM further reaches , validating the modularity and effectiveness of knowledge-guided refinement. The framework is designed to be compatible with other SAM variants, enabling broader applicability to clinical segmentation tasks.

Abstract

While the Segment Anything Model (SAM) has achieved remarkable success in image segmentation, its direct application to medical imaging remains hindered by fundamental challenges, including ambiguous boundaries, insufficient modeling of anatomical relationships, and the absence of uncertainty quantification. To address these limitations, we introduce KG-SAM, a knowledge-guided framework that synergistically integrates anatomical priors with boundary refinement and uncertainty estimation. Specifically, KG-SAM incorporates (i) a medical knowledge graph to encode fine-grained anatomical relationships, (ii) an energy-based Conditional Random Field (CRF) to enforce anatomically consistent predictions, and (iii) an uncertainty-aware fusion module to enhance reliability in high-stakes clinical scenarios. Extensive experiments across multi-center medical datasets demonstrate the effectiveness of our approach: KG-SAM achieves an average Dice score of 82.69% on prostate segmentation and delivers substantial gains in abdominal segmentation, reaching 78.05% on MRI and 79.68% on CT. These results establish KG-SAM as a robust and generalizable framework for advancing medical image segmentation.

Paper Structure

This paper contains 14 sections, 6 equations, 1 figure, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. SAM and Uncertainty in Medical Segmentation
  4. Conditional Random Fields in Image Segmentation
  5. Methodology
  6. Overview
  7. Image Embedding Extraction from SAM
  8. Knowledge Graph Construction
  9. Energy-Based Conditional Random Field Optimization
  10. Experimental Results
  11. Dataset and Implementation Details
  12. Performance and Comparison
  13. Ablation Study
  14. Conclusion

Figures (1)

  • Figure 1: An overview of our proposed KG-SAM framework. A frozen SAM image encoder extracts deep image embeddings and an initial probability map from the input image. A knowledge graph is constructed from biomedical concepts to explicitly model anatomical priors. The core of the framework is the Energy-Based CRF module, which integrates the visual features from SAM with the anatomical constraints from the knowledge graph. Finally, the Fusion module uses the uncertainty map to adaptively weight features and produce the final, anatomically coherent segmentation. The components for CRF optimization and fusion are learnable.