Medical Image Registration via Neural Fields

Shanlin Sun,Kun Han,Chenyu You,Hao Tang,Deying Kong,Junayed Naushad,Xiangyi Yan,Haoyu Ma,Pooya Khosravi,James S. Duncan,Xiaohui Xie

2024-10-11

Abstract:Image registration is an essential step in many medical image analysis tasks. Traditional methods for image registration are primarily optimization-driven, finding the optimal deformations that maximize the similarity between two images. Recent learning-based methods, trained to directly predict transformations between two images, run much faster, but suffer from performance deficiencies due to model generalization and the inefficiency in handling individual image specific deformations. Here we present a new neural net based image registration framework, called NIR (Neural Image Registration), which is based on optimization but utilizes deep neural nets to model deformations between image pairs. NIR represents the transformation between two images with a continuous function implemented via neural fields, receiving a 3D coordinate as input and outputting the corresponding deformation vector. NIR provides two ways of generating deformation field: directly output a displacement vector field for general deformable registration, or output a velocity vector field and integrate the velocity field to derive the deformation field for diffeomorphic image registration. The optimal registration is discovered by updating the parameters of the neural field via stochastic gradient descent. We describe several design choices that facilitate model optimization, including coordinate encoding, sinusoidal activation, coordinate sampling, and intensity sampling. Experiments on two 3D MR brain scan datasets demonstrate that NIR yields state-of-the-art performance in terms of both registration accuracy and regularity, while running significantly faster than traditional optimization-based methods.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper attempts to address the trade-offs between traditional methods and existing learning-based methods in medical image registration. Specifically:

1. **Traditional Optimization Methods**:
- **Advantages**: Capable of finding the optimal deformation field, ensuring maximum similarity between images.
- **Disadvantages**: Usually require strict assumptions about the types of allowable deformations to ensure the regularity of the registration, which can sometimes lead to unfavorable registration results. Additionally, these methods are often time-consuming.

2. **Learning-Based Methods**:
- **Advantages**: Very fast prediction speed once training is completed.
- **Disadvantages**: In unsupervised settings, the registration accuracy is usually not as good as optimization methods, especially when dealing with unseen image pairs, where the model's generalization ability is poor.

To address these issues, the paper proposes a new framework called NIR (Neural Image Registration), which combines the advantages of optimization methods and deep neural networks. NIR models the deformation between image pairs through neural fields, which are continuous functions that can take 3D coordinates as input and output the corresponding deformation vectors. NIR provides two methods for generating deformation fields:

1. **Directly Output Displacement Vector Field**: Used for general deformable registration.
2. **Output Velocity Vector Field and Integrate**: Used to generate the deformation field, ensuring that the deformation is diffeomorphic.

NIR updates the parameters of the neural fields through stochastic gradient descent (SGD) to find the optimal registration. To improve optimization efficiency, the paper also designs various coordinate sampling strategies, including a downsize sampler and a mini-patch sampler, and proposes a hybrid sampler to combine the advantages of these two samplers.

Experimental results show that NIR performs excellently on two 3D MRI brain scan datasets, achieving state-of-the-art levels in both registration accuracy and regularity, while being significantly faster than traditional optimization methods and requiring less memory than learning-based methods (less than 3500MB GPU memory).