Voxelmorph++ Going beyond the cranial vault with keypoint supervision and multi-channel instance optimisation

Mattias P. Heinrich,Lasse Hansen

2022-03-01

Abstract:The majority of current research in deep learning based image registration addresses inter-patient brain registration with moderate deformation magnitudes. The recent Learn2Reg medical registration benchmark has demonstrated that single-scale U-Net architectures, such as VoxelMorph that directly employ a spatial transformer loss, often do not generalise well beyond the cranial vault and fall short of state-of-the-art performance for abdominal or intra-patient lung registration. Here, we propose two straightforward steps that greatly reduce this gap in accuracy. First, we employ keypoint self-supervision with a novel network head that predicts a discretised heatmap and robustly reduces large deformations for better robustness. Second, we replace multiple learned fine-tuning steps by a single instance optimisation with hand-crafted features and the Adam optimiser. Different to other related work, including FlowNet or PDD-Net, our approach does not require a fully discretised architecture with correlation layer. Our ablation study demonstrates the importance of keypoints in both self-supervised and unsupervised (using only a MIND metric) settings. On a multi-centric inspiration-exhale lung CT dataset, including very challenging COPD scans, our method outperforms VoxelMorph by improving nonlinear alignment by 77% compared to 19% - reaching target registration errors of 2 mm that outperform all but one learning methods published to date. Extending the method to semantic features sets new stat-of-the-art performance on inter-subject abdominal CT registration.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The main goal of this paper is to address the issue of insufficient accuracy in existing deep learning-based image registration methods when dealing with large deformation registration. Specifically, most current research focuses on inter-patient brain image registration, and these methods perform poorly when handling regions outside the skull (such as the abdomen or lungs). To solve this problem, the authors propose two main improvements:

1. **Keypoint Self-Supervision**: A new network head is introduced to predict discretized heatmaps, which can better capture large deformations and improve robustness.
2. **Multi-Channel Instance Optimization**: Single-instance optimization using handcrafted features and the Adam optimizer replaces the multiple learning fine-tuning steps.

With these improvements, the proposed method achieves significant performance enhancements across multiple datasets. Notably, in handling challenging chronic obstructive pulmonary disease (COPD) cases, the nonlinear alignment error is reduced by 77%, achieving a target registration error of 2 millimeters, surpassing all published learning methods except one. Additionally, when extended to abdominal CT registration, it also achieves state-of-the-art performance.