This repository contains the Pytorch implementation of the paper "Attention-aware non-rigid image registration for accelerated MR imaging"
by Aya Ghoul, Jiazhen Pan, Andreas Lingg, Jens Kübler, Patrick Krumm, Kerstin Hammernik, Daniel Rueckert, Sergios Gatidis, Thomas Küstner
published on IEEE Transactions on Medical Imaging 2024
GMARAFT is a deep learning-based method for the alignment/registration of both fully sampled and accelerated magnetic resonance (MR) images. We combined RAFT, GMA and a denoising network to derive reliable motion estimation across different sampling trajectories (Cartesian and radial) and acceleration factors of up to 16x for cardiac motion and 30x for respiratory motion. The model was tested on the downstream task of motion-compensated reconstruction.
Figure 1: Illustration of GMARAFT for non-rigid image registration with exemplary Cartesian VISTA undersampling.
For a detailed description of the required conda environment, please refer to the provided requirements file. Additionally, the implementation relies on the Merlin (Machine Enhanced Reconstruction Learning and Interpretation Networks) library. Refer to this link for build instructions. Trained networks from the above manuscript can be obtained from the corresponding author upon reasonable request.
The model's robustness was tested against three undersampling masks
To train the model, creating the data list in the loader needs to be first customized. Then, the config .json file can be edited to adjust training hyperparameters and paths. The run-pairwise.py script will train an image-to-image registration network (described in the paper). Model weights will be saved to a path specified in the config .json file.
To assess the model's quality by computing the photometric error between a fixed and warped test image, simply customize the eval.py script with the appropriate path to the model's weights and image data.
Figure 2: Motion estimation for relaxation (top) and contraction (bottom) motion in a representative test subject using Elastix, VoxelMoprh, Vit-v-net, TransMorph, XMorpher and our proposed method in a fully sampled case. VoxelMorph, Vit-V-Net, TransMorph and XMorpher are trained using fully sampled images. Color-encoded motion estimations (top row) are shown together with the moving images overlayed with quiver plots (bottom row).
Figure 3: Motion-compensated reconstructions of the end-systolic and end-diastolic phase are shown in a fully sampled case for a representative subject. Images were reconstructed using iterative Sense with Elastix, VoxelMorph, Vit-V-Net, TransMorph, XMorpher and the proposed model motion estimates in comparison to the fully sampled ground truth (GT) image. VoxelMorph, Vit-V-Net, TransMorph and XMorpher are trained using fully sampled images.
Figure 4: Cardiac motion estimation is shown for the fully sampled and retrospectively undersampled acquisitions with VISTA (Cartesian) and radial (non-Cartesian) undersampling for R=16 acceleration.
Figure 5: Respiratory motion estimation for the fully sampled and Cartesian vdPD accelerated acquisitions with accelerations R=16 and R=30. Deformation fields are overlaid on the moving image. Images of motion-compensated reconstructions are depicted next to the used color-encoded deformation fields.
- Registration performance
Figure 6: Violin plots of the residual photometric error values in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.
Figure 7: Violin plots of the percentages of non-positive values in the determinant of the Jacobian matrix values in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.
- Downstream task: Motion-compensated reconstruction image quality metrics
Figure 8: Violin plots of the SSIM in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.
Figure 9: Violin plots of the PSNR in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.
Figure 10: Violin plots of the NRMSE in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.
Figure 11: Violin plots of the HFEN in fully-sampled, R=8 and R=16 test data for VISTA and radial undersampling.