A Plug-and-Play Image Registration Network

TL;DR

This paper introduces PIRATE, a novel deep learning-based deformable image registration method that explicitly incorporates data fidelity and a CNN prior, with PIRATE+ further enhancing performance through end-to-end training with deep equilibrium models.

Contribution

The paper proposes PIRATE, a plug-and-play framework for image registration that integrates explicit data fidelity and CNN priors, and introduces PIRATE+ with end-to-end training for improved accuracy.

Findings

Achieves state-of-the-art performance on OASIS and CANDI datasets.

Effectively combines data fidelity with learned CNN priors.

Enhances registration accuracy through deep equilibrium model fine-tuning.

Abstract

Deformable image registration (DIR) is an active research topic in biomedical imaging. There is a growing interest in developing DIR methods based on deep learning (DL). A traditional DL approach to DIR is based on training a convolutional neural network (CNN) to estimate the registration field between two input images. While conceptually simple, this approach comes with a limitation that it exclusively relies on a pre-trained CNN without explicitly enforcing fidelity between the registered image and the reference. We present plug-and-play image registration network (PIRATE) as a new DIR method that addresses this issue by integrating an explicit data-fidelity penalty and a CNN prior. PIRATE pre-trains a CNN denoiser on the registration field and "plugs" it into an iterative method as a regularizer. We additionally present PIRATE+ that fine-tunes the CNN prior in PIRATE using deep equilibrium models (DEQ). PIRATE+ interprets the fixed-point iteration of PIRATE as a network with effectively infinite layers and then trains the resulting network end-to-end, enabling it to learn more task-specific information and boosting its performance. Our numerical results on OASIS and CANDI datasets show that our methods achieve state-of-the-art performance on DIR.