[Paper Review] Real-time Deep Registration With Geodesic Loss.

Seyed Sadegh Mohseni Salehi, Shadab Khan+2 more|arXiv (Cornell University)|Mar 15, 2018

Fetal and Pediatric Neurological Disorders25 references18 citations

TL;DR

This paper proposes a real-time deep learning method for 3D medical image registration using regression CNNs that predict 3D pose (rotation and translation) via angle-axis representation, with geodesic loss significantly improving accuracy and capture range. The approach achieves sub-100ms inference on fetal and newborn brain MRI, generalizes across image contrasts via conditional GAN-based domain adaptation, and enhances optimization-based registration pipelines.

ABSTRACT

With an aim to increase the capture range and accelerate the performance of state-of-the-art inter-subject and subject-to-template 3D registration, we propose deep learning-based methods that are trained to find the 3D position of arbitrarily oriented subjects or anatomy based on slices or volumes of medical images. For this, we propose regression CNNs that learn to predict the angle-axis representation of 3D rotations and translations using image features. We use and compare mean square error and geodesic loss for training regression CNNs in two different scenarios: 3D pose estimation from slices and 3D to 3D registration. As an exemplary application, we applied the proposed methods to register arbitrarily oriented reconstructed images of fetuses scanned in-utero at a wide gestational age range to a standard atlas space. Our results show that in such registration applications that are amendable to learning, the proposed deep learning methods with geodesic loss minimization can achieve accurate results with a wide capture range in real-time (<100ms). We tested the generalization capability of the trained CNNs on an expanded age range and on images of newborn subjects with similar and different MR image contrasts. We trained our models on T2-weighted fetal brain MRI scans and used them to predict the 3D position of newborn brains based on T1-weighted MRI scans. We showed that trained models generalized well for the new domain when we performed image contrast transfer through a conditional generative adversarial network. This indicates that the domain of application of the trained deep regression CNNs can be further expanded to image modalities and contrasts other than those used in training. A combination of our proposed methods with optimization-based registration algorithms can dramatically enhance the performance of automatic imaging devices and image processing methods of the future.

Motivation & Objective

• Address the limited capture range and slow speed of state-of-the-art inter-subject and subject-to-template 3D registration methods in medical imaging.
• Enable real-time, accurate 3D pose estimation from 2D slices or 3D volumes for arbitrarily oriented anatomical structures.
• Improve generalization across diverse image contrasts and age ranges by leveraging domain adaptation via conditional GANs.
• Demonstrate that geodesic loss outperforms mean square error in learning rotation and translation parameters for 3D registration.
• Integrate learned initialization with optimization-based registration to boost performance of future automatic imaging systems.

Proposed method

• Train convolutional neural networks (CNNs) to regress the angle-axis representation of 3D rotations and translations from input medical images.
• Compare training with mean square error (MSE) and geodesic loss to optimize for rotational and translational accuracy.
• Apply the models to 3D pose estimation from 2D slices and full 3D volume-to-volume registration tasks.
• Use conditional generative adversarial networks (cGANs) to transfer image contrast from T2-weighted fetal MRI to T1-weighted newborn MRI for domain generalization.
• Integrate the deep regression predictions as initial estimates in optimization-based registration pipelines to improve convergence and accuracy.
• Evaluate performance on in-utero fetal MRI across a wide gestational age range and on newborn T1-weighted scans with different contrasts.

Experimental results

Research questions

• RQ1 Can geodesic loss improve the accuracy and capture range of deep learning-based 3D pose estimation in medical image registration compared to mean square error?
• RQ2 To what extent can a CNN trained on T2-weighted fetal brain MRI generalize to T1-weighted newborn brain MRI with different image contrasts?
• RQ3 How effective is the combination of deep regression with optimization-based registration in enhancing overall registration performance?
• RQ4 Can conditional GANs enable successful image contrast transfer to extend the domain applicability of trained regression models?
• RQ5 What is the inference speed of the proposed method, and can it achieve real-time performance (<100ms) for clinical deployment?

Key findings

• The proposed method achieves real-time performance with inference times under 100ms, enabling practical deployment in clinical workflows.
• Geodesic loss significantly improves registration accuracy and capture range compared to mean square error, especially for rotational components.
• Trained CNNs generalize well to newborn brain MRI scans with different contrasts when combined with conditional GAN-based image translation.
• The model maintains high accuracy when tested on an expanded age range beyond the training distribution, indicating strong robustness.
• The integration of learned initializations with optimization-based methods leads to faster convergence and improved final registration accuracy.
• The approach demonstrates transferability across imaging modalities, suggesting broad applicability beyond the specific training data.