Unsupervised learning of spatially varying regularization for diffeomorphic image registration

keywords: Deformable Image Registration, Spatially Varying Regularization

This is a PyTorch implementation of my paper:

Chen, Junyu, et al. "Unsupervised learning of spatially varying regularization for diffeomorphic image registration." arXiv preprint arXiv:2412.17982 (2024).

Introduction

Spatially varying regularization accommodates the deformation variations that may be necessary for different anatomical regions during deformable image registration. Historically, optimization-based registration models have harnessed spatially varying regularization to address anatomical subtleties. However, most modern deep learning-based models tend to gravitate towards spatially invariant regularization, wherein a homogenous regularization strength is applied across the entire image, potentially disregarding localized variations. In this paper, we propose a hierarchical probabilistic model that integrates a prior distribution on the deformation regularization strength, enabling the end-to-end learning of a spatially varying deformation regularizer directly from the data.

Registration tasks

• Atlas-to-subject registration on IXI dataset (brain MRI) [code]
• Atlas-to-subject registration on AutoPET dataset (whole-body CT) [To be added]
• Intra-subject registration on ACDC and M&Ms dataset (cardiac MRI) [To be added]

Examples of the learned spatially varying regularizer

Hypernetwork for continuous regularization control

We further incorporated the concept from HyperMorph, enabling the learning of a set of regularization hyperparameters for continuous control of spatially varying regularization at the test time.

Citation:

If you find this code is useful in your research, please consider to cite:

@article{chen2024unsupervised,
title={Unsupervised learning of spatially varying regularization for diffeomorphic image registration},
author={Chen, Junyu and Wei, Shuwen and Liu, Yihao and Bian, Zhangxing and He, Yufan and Carass, Aaron and Bai, Harrison and Du, Yong},
journal={arXiv preprint arXiv:2412.17982},
year={2024}
}

About Me