PromptReg: Interactive Registration by "Corresponding Prompts" for Segment Anything Model (SAM).

Abstract

Effectively establishing correspondence between two images is at the centre of image registration methods. Spatially omnipresent representations, including dense displacement fields (DDFs) and spatial (non-)rigid transformations, have been used to parameterise such correspondence. Alternatively, region-based representation uses paired regions of interest (ROIs) to represent region-level correspondence, while retaining its local and dense representation capability at pixel/voxel level if required. Thus, registration can be re-envisioned as a problem of segmenting corresponding paired ROIs in the to-be-registered images. In this work, we utilize models such as SAM, which are pre-trained on substantive datasets, to segment ROIs of the same class from two images, for a new training-free, non-iterative registration algorithm. First, a "corresponding prompt problem" is posed to find a corresponding Prompt Y on Image Y, given any vision Prompt X on Image X, such that the two respectively prompt-conditioned segmentations are a pair of corresponding ROIs from the two images. Second, we propose an "inverse prompt" solution to the corresponding prompt problem, by inverting Prompt X to the Image Y prompt space, where the Jacobian of prototypical features is used. Third, we propose a new registration algorithm that identifies multiple paired corresponding ROIs, by marginalizing the inverted Prompt X over both prompt and spatial spaces, random sampling Prompt X and spatial warping Image X. Comprehensive experiments were conducted on five applications of registering 3D prostate MR, 3D abdomen CT, 3D lung CT, 2D histopathology and, as a non-medical example, 2D aerial images. Based on metrics including Dice and target registration errors on anatomical structures, the proposed registration outperforms both intensity-based iterative algorithms and learning-based networks, even yielding competitive performance with weakly-supervised registration which requires fully-segmented training data.