SAMReg: SAM-enabled Image Registration with ROI-based Correspondence

TL;DR

The proposed SAMReg models outperform both intensity-based iterative algorithms and DDF-predicting learning-based networks across tested metrics including Dice and target registration errors on anatomical structures, and further demonstrates competitive performance compared to weakly-supervised registration approaches that rely on fully-segmented training data.

Abstract

This paper describes a new spatial correspondence representation based on paired regions-of-interest (ROIs), for medical image registration. The distinct properties of the proposed ROI-based correspondence are discussed, in the context of potential benefits in clinical applications following image registration, compared with alternative correspondence-representing approaches, such as those based on sampled displacements and spatial transformation functions. These benefits include a clear connection between learning-based image registration and segmentation, which in turn motivates two cases of image registration approaches using (pre-)trained segmentation networks. Based on the segment anything model (SAM), a vision foundation model for segmentation, we develop a new registration algorithm SAMReg, which does not require any training (or training data), gradient-based fine-tuning or prompt engineering. The proposed SAMReg models are evaluated across five real-world applications, including intra-subject registration tasks with cardiac MR and lung CT, challenging inter-subject registration scenarios with prostate MR and retinal imaging, and an additional evaluation with a non-clinical example with aerial image registration. The proposed methods outperform both intensity-based iterative algorithms and DDF-predicting learning-based networks across tested metrics including Dice and target registration errors on anatomical structures, and further demonstrates competitive performance compared to weakly-supervised registration approaches that rely on fully-segmented training data. Open source code and examples are available at: https://github.com/sqhuang0103/SAMReg.git.

Figures (9)

• Figure 1: The representation diagram of a) rigid transformation-based, b) DDF-based and c) ROI-based correspondence. The ROI-based correspondence is exhibited in a concise and elegant manner.
• Figure 2: The pipeline of the proposed SAM-Reg algorithm. The process can be divided into two primary steps: ROI embedding and ROI matching. Initially, fixed and moving images $I^x$ and $I^y$ are encoded into feature representations $F^x$ and $F^y$, which are then decoded to produce segmented masks $M_k^x$ and $M_k^y$ highlighting various ROIs $R^x_k$ and $R^y_k$. These ROIs are embedded into prototype vectors $[p^x_1,p^x_2,...,p^x_{K^x}]$ and $[p^y_1,p^y_2,...,p^y_{K^y}]$, respectively. The ROI matching step utilizes a similarity matrix $S \in \mathbb{R}^{K^x \times K^y}$ to identify and match the most similar candidate ROIs from the moving image to each fixed ROI.
• Figure 3: Qualitative comparisons of SAMReg with other leading registration methods balakrishnan2019voxelmorphhu2018weakly on 2D datasets: a) Retina and b) Aerial. Generated DDF aligns the ROI from the moving to the moved label, matching the fixed label's ROI. Colored areas in the SAMReg column highlight the corresponding ROIs, which closely resemble the fixed label, showcasing superior alignment performance.
• Figure 4: Qualitative comparisons of SAMReg with other leading registration methods balakrishnan2019voxelmorphhu2018weakly on 3D datasets: a) Prostate bosaily2015promis, b) Cardiac bernard2018deep and c) Lung LUNG. Generated DDF maps ROIs from the moving to the moved label, aligning with the fixed label's ROIs. Colored highlights in the SAMReg column indicate the matched ROIs by SAMReg, demonstrating its superior registration performance.
• Figure 5: The qualitative comparisons of a) SAM kirillov2023segment and b) fully-supervised trained segmentation models milletari2016vtragakis2023fullyroy2023mednext on inter- and intra-subject registration tasks, including the wrapping performance on label ROI (top) and pseudo ROI (bottom). SAMReg with SAM benefits from its extensive insight. Despite being fully trained, dataset-specific segmentation models show competitive performance only on familiar ROIs (label ROI) but falter on unknown ROIs (pseudo ROI).

Theorems & Definitions (1)

• Theorem 1