PaperlayerView exact revision on arXiv

Source-linked AI summary

Real-time Deep Pose Estimation with Geodesic Loss for Image-to-Template Rigid Registration

Seyed Sadegh Mohseni Salehi, Shadab Khan, Deniz Erdogmus, Ali Gholipour

arXiv:1803.05982v4cs.CV

TL;DR

Non-convex registration and unreliable initialization limit capture range, while exhaustive or multi-scale approaches can be costly. The paper uses CNNs to estimate 3D pose from slices or volumes, with geodesic rotation loss and a correction network. The methods provide real-time, wide-capture-range registration, and contrast transfer supports testing T2-trained models on T1 newborn scans.

Abstract

from arXiv · show

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 to train regression CNNs for 3D pose estimation used in two different scenarios: slice-to-volume registration and volume-to-volume registration. 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 also 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 pose of newborn brains based on T1-weighted MRI scans. We showed that the 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 accelerated optimization-based registration algorithms can dramatically enhance the performance of automatic imaging devices and image processing methods of the future.

I. INTRODUCTION

Non-convex registration objectives and unreliable initialization limit capture range, while grid search and multi-scale methods can be too costly for time-sensitive use. Prior learning-based work addressed related tasks, but global rigid pose estimation and broad rotation ranges remained distinct challenges.

A. Background I

  • Non-convex similarity objectives constrain capture range, and initialization methods can fail on wide rotations or complex shapes; grid search and multi-scale registration are computationally expensive.This limits their usefulness in time-sensitive applications such as image-based navigation.Source: A. Background · p. 1
  • Learning-based registration was limited, although image-to-template tasks with wide plausible pose ranges were identified as amenable to learning.Human observers can quickly approximate pose through feature identification without iterative optimization.Source: A. Background · p. 1

B. Related Work

  • For global rigid registration, Liao et al. reported improved CT accuracy and robustness with reinforcement learning, but the method was relatively slow and lacked a systematic test-time stopping criterion.Source: B. Related Work · p. 2
  • Miao et al., used hierarchical CNN regression for X-ray-to-CT slice-to-volume registration, reporting improved registrations in 100 ms, about 20–45 times faster than the best intensity-based method.Experiments used rotations up to 30° and perturbations with standard deviations of 10° in each rotation parameter.Source: B. Related Work · p. 2

C. Contributions

The paper proposes CNN regression for 3D pose from volumetric or sectional images, using angle-axis rotations and geodesic loss. It targets slice-to-volume and volume-to-volume registration, with a correction network and optimization-based registration in the proposed pipeline.

C. Contributions

  • The proposed CNN regression estimates 3D pose from arbitrarily oriented volumetric or sectional images, beyond prior work focused on 2D-projected representations.Source: C. Contributions · p. 2
  • The method uses angle-axis rotation vectors and bi-invariant geodesic distance on SO(3) as its loss, and adds a correction network to estimate translation parameters.The estimated pose can initialize optimization-based registration; the paper distinguishes pose relative to a canonical space from registration of both rotation and translation.Source: C. Contributions · p. 2

II. METHODS

  • The methods section introduces rotation representation, network architectures, and a two-step training algorithm with loss functions for the regression networks.Source: §II · p. 3
  • The rotation vector compactly represents a 3D rotation with its unit vector as the axis and its magnitude as the angle, which is non-negative and at most π.For a 3D vector v, the method defines its normalized vector as the axis and its norm as the rotation angle; Rodrigues’ formula yields the rotation matrix.Source: A. 3D Rotation Representation · p. 3

B. Slice 3D Pose Estimation Network Architecture

The slice pose network uses a ResNet18 feature extractor with separate regression heads for 3D rotation and atlas-space slice location.

B. Slice 3D Pose Estimation Network Architecture

C. Volume 3D Pose Estimation Network Architecture

The volume pose network extracts image features to estimate orientation, while a correction network takes a roughly aligned moving image and reference as input to estimate rotation and translation.

C. Volume 3D Pose Estimation Network Architecture

D. Volume-to-Volume Correction Network Architecture

  • The correction network receives the reference and roughly oriented moving volume as a two-channel input and estimates both rotation and translation.Initial translations are assumed to come from center-of-gravity matching, and initial rotations from the volume pose network.Source: D. Volume-to-Volume Correction Network… · p. 4
  • The correction network shares the volume pose network’s 3D feature extractor and adds a three-parameter translation head alongside its rotation head.Source: D. Volume-to-Volume Correction Network… · p. 4

E. Training the Networks

Training combines translation MSE with rotation losses, using MSE for rotation initially and geodesic loss in the second stage. This schedule pairs MSE’s convexity advantage with geodesic loss’s accuracy.

E. Training the Networks

  • MSE is used first because its convexity can narrow the pose-prediction search space, although it does not accurately measure distance between rotations.Source: E. Training the Networks · p. 4
  • Rotation training switches from MSE in stage one to geodesic loss in stage two, while translation loss remains MSE.Each stage involved ten epochs.Source: E. Training the Networks · pp. 4–5
  • The geodesic loss measures the rotation angle between the estimated and ground-truth matrices and is computed using equation (10) for easier implementation.The paper describes this distance as a natural Riemannian metric on SO(3).Source: E. Training the Networks · pp. 4–5

A. Datasets

The study uses fetal and newborn MRI data, with fetal scans registered to an age-matched atlas to generate ground-truth alignment. Training and testing span slice- and volume-based pose estimation, with different orientation ranges and newborn contrasts.

A. Datasets

  • The dataset includes 93 reconstructed fetal T2-weighted MRI scans and T1- and T2-weighted scans from 40 newborns.Fetal scans were acquired at 21–37 weeks’ gestational age; the newborn data came from the developing human connectome project.Source: A. Datasets · p. 5
  • Fetal images were registered to an age-matched spatiotemporal atlas, with multi-scale rigid registration refining the alignment after mapping into fetal coordinates.Alignments have intrinsic uncertainty because anatomy differs across subjects and ages, so images were manually checked for visually correct atlas alignment.Source: A. Datasets · p. 5
  • Training used reconstructed T2-weighted scans from 36 fetuses at 28–37 weeks, with slice pose estimation trained over half-space and volume pose estimation over the full rotation space.Half-space training was chosen for slices because their 2D information makes rotations π radians apart difficult to distinguish; volume rotations were uniformly sampled over all 3D rotations.Source: 1) Training Dataset · p. 5
  • The Correction-Net was trained on a limited transformation range to refine initial predictions, while a wider-range version was evaluated as 3DReg-Net.Initial translations used center-of-gravity matching and initial rotations came from pose-estimation networks.Source: 1) Training Dataset · pp. 5–6
  • Three held-out test sets cover later-gestation fetal T2 scans and slices, earlier-gestation fetal T2 scans, and newborn T2 and T1 scans.The passage reports no overlap between test sets and the training set.Source: 2) Testing Datasets · p. 6
  • Test images received ten random rotations each, and Figure 2 plots their distance from identity in degrees, distinguishing slice samples from volume samples.The figure caption identifies half-space rotations for slice pose estimation and full-space rotations for volume pose estimation.Source: 2) Testing Datasets · p. 6, Figure 2 · p. 6

B. Intensity-Based Registration

The section defines conventional intensity-based rigid-registration comparisons for volume-to-volume and slice-to-volume tasks, alongside versions initialized by the CNN’s predicted transforms.

B. Intensity-Based Registration

  • VVR-GC and VVR-PAA use multi-scale rigid registration with gradient-descent maximization of normalized mutual information, differing in gravity-center versus moments-and-principal-axis initialization.Source: B. Intensity-Based Registration · p. 6
  • SVR-GC uses multi-scale registration initialized by center-of-gravity matching and optimizes normalized cross-correlation; SVR-Deep replaces that initialization with 3DPose-Net predictions.VVR-Deep likewise substitutes 3DPose-Net predicted transforms for VVR-GC initialization, and the learning rate was lowered for CNN-initialized VVR and SVR.Source: B. Intensity-Based Registration · p. 6

C. Results

The results show that deep pose predictions broaden registration capture range, with deep-initialized optimization strongest in volume-to-volume tests. Geodesic loss improved pose estimation over MSE, and age and contrast-transfer experiments assessed generalization.

C. Results

  • Deep initialization performed best for slice-to-volume pose estimation; geodesic loss improved results over MSE, while optimization without deep initialization failed in most cases.The geodesic-loss improvement was significant in the 80–90° and 90–100° rotation bins.Source: 1) Slice 3D Pose Estimation · p. 6, Figure 3 · p. 7, Table I · p. 8
  • Below 5 mm, the estimated slice-location error was for most slices, while boundary slices with limited or similar features were harder to localize.Slice-location estimation across different anatomies and ages is challenging and can be augmented with slice-level motion tracking when motion is fast and continuous.Source: 1) Slice 3D Pose Estimation · pp. 6–7, Figure 4 · p. 7
  • Tests on original fetal MRI slices used the estimated pose and location to retrieve corresponding slices from age-matched atlases, with representative outputs shown for five fetuses.The pipeline extracted the fetal brain from each slice before passing it to 3DPose-Net.Source: 1) Slice 3D Pose Estimation · p. 7, Figure 5 · p. 8, Figure 6 · p. 8
  • VVR-GC worked well for 0–60° rotations but failed for almost all rotations above 80°; VVR-PAA improved high-angle results but remained unreliable and lost performance at low angles.The Correction-Net performed best among the trained CNN models, with only a small fraction of failures at 130–180°; VVR-Deep had the most accurate results and fewest failures.Source: 2) Volume 3D Pose Estimation · p. 7
  • VVR-Deep performed best overall, while the Correction-Net produced comparable results with far fewer failures and substantially faster registration.VVR-Deep initializes optimization with 3DPose-Net predictions.Source: 2) Volume 3D Pose Estimation · p. 7, Table II · p. 8
  • Pose estimates generalized across tested ages and scan settings, with slightly higher errors at younger ages; T1-weighted newborn scans generalized poorly until contrast transfer to T2-weighted images reduced pose error.The age tests included younger fetuses and ex-utero newborns; the modality tests used 7 newborns and 70 samples.Source: 3) Generalization property of the… · p. 9, Figure 10 · p. 10, Figure 11 · p. 10
  • All deep-learning algorithms ran in real time on a GPU, and the conditional GAN contrast-transfer stage averaged about 50 ms.The reported GPU timing was measured on an NVIDIA GeForce GTX 1080; the authors also describe the methods as potentially useful for real-time image-based surgical navigation and motion-robust imaging.Source: 4) Testing times · p. 9, Table III · p. 9

IV. DISCUSSION

The paper presents deep CNN pose estimation for medical-image registration, using geodesic loss and a two-step correction strategy. It also tests generalization across fetal ages and newborn MRI contrasts, including T1-to-T2 contrast transfer.

IV. DISCUSSION

  • Geodesic loss significantly improves performance over MSE on rotation vectors, especially in 3D when images contain sufficient pose information.The networks use feature-extraction layers and regression heads, with a πtanh output nonlinearity for angle-axis rotation prediction.Source: §IV · p. 10
  • A two-step pose-estimation and correction-network approach supports accurate real-time inter-subject rigid registration across rotation and translation ranges.The initial pose is estimated in atlas space, then supplied with a reference image as two input channels to the correction network; center-of-gravity matching can provide initial translations in real time.Source: §IV · p. 10
  • Training on older fetuses and testing on younger ones evaluates age generalization despite rapid early-gestational changes in brain shape and size.The study uses a predefined scale inferred from gestational age and a fetal brain MRI atlas.Source: §IV · p. 10
  • Newborn T2 scans were handled well, while T1 performance dropped; translating T1 scans into T2-like images with a conditional GAN restored strong performance.The translation network was trained on paired T1- and T2-weighted newborn images and applied as real-time preprocessing before pose estimation.Source: §IV · p. 10
  • The authors propose that paired-image conditional GANs can extend these pose-estimation models to other modalities and contrasts.This broader generalization claim is conditioned on having paired images for training.Source: §IV · p. 10

V. CONCLUSION

The paper develops deep pose-estimation networks for slice-to-volume and volume-to-volume image-to-template registration. The methods provide real-time registration over plausible 3D rotations and can initialize optimization-based registration.

V. CONCLUSION

  • The study develops and evaluates deep pose-estimation networks for both slice-to-volume and volume-to-volume registration.The conclusion frames these as learning-based image-to-template registration scenarios.Source: §V · p. 10
  • The proposed CNN methods provide real-time registration with a wide capture range across plausible 3D rotations and good initialization for optimization-based methods.The paper contrasts this with highly evolved multi-scale optimization methods that can converge to wrong local minima because of non-convex cost functions.Source: §V · p. 11
  • The authors propose combining these techniques with accelerated optimization-based methods to enhance imaging devices and image-processing methods.Source: §V · p. 11

APPENDIX A

Appendix A describes conditional GAN preprocessing that translates newborn T1-weighted images into T2-like images. Its training combines adversarial learning with an L1 distance objective and uses paired T1/T2 scans.

APPENDIX A

  • The conditional GAN is trained to map T1-weighted images to T2-like images using paired image inputs.The generator produces synthetic T2-like images, while the discriminator distinguishes real from generated T2 contrast in T1/T2 pairs.Source: Appendix A · p. 11
  • The generator objective combines the conditional GAN loss with a weighted L1 loss.The L1 distance between generated and real T2 scans encourages sharper images.Source: Appendix A · p. 11
  • The GAN was trained on 33 paired newborn T1/T2 images, yielding 3,300 paired slices, and tested on seven newborn brain images.Training subjects were not used in the test set.Source: Appendix A · p. 11

Related papers

Loading 1803.05982v4…