Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Medical Image Registration Meets Vision Foundation Model: Prototype Learning and Contour Awareness

Hao Xu, Tengfei Xue, Jianan Fan, Dongnan Liu, Yuqian Chen, Fan Zhang, Carl-Fredrik Westin, Ron Kikinis, Lauren J. O'Donnell, Weidong Cai

TL;DR

Deformable medical image registration often relies solely on intensity similarity, risking inaccurate alignment in complex anatomies. The paper introduces a SAM-assisted registration framework that integrates SAM-generated segmentation masks, prototype contrastive learning and alignment, and contour-aware loss within a cross-fusion transformer to enforce both global semantic consistency and local boundary accuracy. It contributes explicit anatomical information injection, region-wise prototype learning, and contour-aware optimization, validated on Abdomen CT and ACDC MRI with state-of-the-art results. The approach enhances robustness and precision in challenging scenarios, enabling more reliable clinical image analysis.

Abstract

Medical image registration is a fundamental task in medical image analysis, aiming to establish spatial correspondences between paired images. However, existing unsupervised deformable registration methods rely solely on intensity-based similarity metrics, lacking explicit anatomical knowledge, which limits their accuracy and robustness. Vision foundation models, such as the Segment Anything Model (SAM), can generate high-quality segmentation masks that provide explicit anatomical structure knowledge, addressing the limitations of traditional methods that depend only on intensity similarity. Based on this, we propose a novel SAM-assisted registration framework incorporating prototype learning and contour awareness. The framework includes: (1) Explicit anatomical information injection, where SAM-generated segmentation masks are used as auxiliary inputs throughout training and testing to ensure the consistency of anatomical information; (2) Prototype learning, which leverages segmentation masks to extract prototype features and aligns prototypes to optimize semantic correspondences between images; and (3) Contour-aware loss, a contour-aware loss is designed that leverages the edges of segmentation masks to improve the model's performance in fine-grained deformation fields. Extensive experiments demonstrate that the proposed framework significantly outperforms existing methods across multiple datasets, particularly in challenging scenarios with complex anatomical structures and ambiguous boundaries. Our code is available at https://github.com/HaoXu0507/IPMI25-SAM-Assisted-Registration.

Medical Image Registration Meets Vision Foundation Model: Prototype Learning and Contour Awareness

TL;DR

Deformable medical image registration often relies solely on intensity similarity, risking inaccurate alignment in complex anatomies. The paper introduces a SAM-assisted registration framework that integrates SAM-generated segmentation masks, prototype contrastive learning and alignment, and contour-aware loss within a cross-fusion transformer to enforce both global semantic consistency and local boundary accuracy. It contributes explicit anatomical information injection, region-wise prototype learning, and contour-aware optimization, validated on Abdomen CT and ACDC MRI with state-of-the-art results. The approach enhances robustness and precision in challenging scenarios, enabling more reliable clinical image analysis.

Abstract

Medical image registration is a fundamental task in medical image analysis, aiming to establish spatial correspondences between paired images. However, existing unsupervised deformable registration methods rely solely on intensity-based similarity metrics, lacking explicit anatomical knowledge, which limits their accuracy and robustness. Vision foundation models, such as the Segment Anything Model (SAM), can generate high-quality segmentation masks that provide explicit anatomical structure knowledge, addressing the limitations of traditional methods that depend only on intensity similarity. Based on this, we propose a novel SAM-assisted registration framework incorporating prototype learning and contour awareness. The framework includes: (1) Explicit anatomical information injection, where SAM-generated segmentation masks are used as auxiliary inputs throughout training and testing to ensure the consistency of anatomical information; (2) Prototype learning, which leverages segmentation masks to extract prototype features and aligns prototypes to optimize semantic correspondences between images; and (3) Contour-aware loss, a contour-aware loss is designed that leverages the edges of segmentation masks to improve the model's performance in fine-grained deformation fields. Extensive experiments demonstrate that the proposed framework significantly outperforms existing methods across multiple datasets, particularly in challenging scenarios with complex anatomical structures and ambiguous boundaries. Our code is available at https://github.com/HaoXu0507/IPMI25-SAM-Assisted-Registration.

Paper Structure

This paper contains 27 sections, 12 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Segment Anything Model
  4. Deformable Medical Image Registration
  5. Prototype Learning
  6. Method
  7. SAM Segmentation Mask Generation
  8. Encoder and Cross Fusion Decoder
  9. Prototype Contrastive Learning and Alignment
  10. Prototype Feature Extraction
  11. Prototype Contrast and Alignment
  12. Contour-Aware Loss
  13. Overall Loss
  14. Similarity Loss
  15. Smooth Loss
  16. ...and 12 more sections

Figures (3)

  • Figure 1: The overview of the proposed SAM-assisted medical image registration framework includes the use of the Segment Anything Model (SAM) to generate segmentation masks, optimization of the feature space through prototype contrast and alignment, enhancement of contour alignment via contour-aware loss, and a cross fusion transformer for image and mask feature fusion registration.
  • Figure 2: The overall architecture of our Encoder and Cross Fusion Decoder.
  • Figure 3: Visualization of registration results for our proposed method and compared methods on ACDC Cardiac MRI dataset and Abdomen CT dataset. The comparison methods are all weakly supervised registration methods utilizing SAM masks.