SE(3)-Equivariant and Noise-Invariant 3D Rigid Motion Tracking in Brain MRI
arXiv:2312.13534 · doi:10.1109/TMI.2024.3411989
Abstract
Rigid motion tracking is paramount in many medical imaging applications where movements need to be detected, corrected, or accounted for. Modern strategies rely on convolutional neural networks (CNN) and pose this problem as rigid registration. Yet, CNNs do not exploit natural symmetries in this task, as they are equivariant to translations (their outputs shift with their inputs) but not to rotations. Here we propose EquiTrack, the first method that uses recent steerable SE(3)-equivariant CNNs (E-CNN) for motion tracking. While steerable E-CNNs can extract corresponding features across different poses, testing them on noisy medical images reveals that they do not have enough learning capacity to learn noise invariance. Thus, we introduce a hybrid architecture that pairs a denoiser with an E-CNN to decouple the processing of anatomically irrelevant intensity features from the extraction of equivariant spatial features. Rigid transforms are then estimated in closed-form. EquiTrack outperforms state-of-the-art learning and optimisation methods for motion tracking in adult brain MRI and fetal MRI time series. Our code is available at https://github.com/BBillot/EquiTrack.
Published at IEEE transactions on Medical Imaging
References in corpus (18)
- Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
- A Simple Framework for Contrastive Learning of Visual Representations
- A Deep Learning Framework for Unsupervised Affine and Deformable Image Registration
- SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining
- TorchIO: A Python library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning
- 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data
- Robust machine learning segmentation for large-scale analysis of heterogeneous clinical brain MRI datasets
- SynthMorph: learning contrast-invariant registration without acquired images
- Disentangled Representation Learning in Cardiac Image Analysis
- Learning Disentangled Representations in the Imaging Domain
- 3D G-CNNs for Pulmonary Nodule Detection
- AIRNet: Self-Supervised Affine Registration for 3D Medical Images using Neural Networks
- Scale-Equivariant Steerable Networks
- Intertwiners between Induced Representations (with Applications to the Theory of Equivariant Neural Networks)
- An Artificial Agent for Robust Image Registration
- The Lie Derivative for Measuring Learned Equivariance
- Dynamic Neural Fields for Learning Atlases of 4D Fetal MRI Time-series
- Deformation Robust Roto-Scale-Translation Equivariant CNNs