SE(3)-Equivariant and Noise-Invariant 3D Rigid Motion Tracking in Brain MRI
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.
- J. Andre, B. Bresnahan, M. Mossa-Basha, M. Hoff, P. Smith, Y. Anzai, and W. Cohen, “Toward Quantifying the Prevalence, Severity, and Cost Associated With Patient Motion During Clinical MR Examinations,” Journal of the American College of Radiology, vol. 12, pp. 89–95, 2015.
- J. Xu, M. Zhang, E. A. Turk, L. Zhang, E. Grant, K. Ying, P. Golland, and E. Adalsteinsson, “Fetal Pose Estimation in Volumetric MRI using a 3D Convolution Neural Network,” in Medical Image Computing and Computer-Assisted Intervention, vol. 11767, 2019, pp. 403–410.
- N. White, C. Roddey, A. Shankaranarayanan, E. Han, D. Rettmann, J. Santos, J. Kuperman, and A. Dale, “PROMO: Real-time prospective motion correction in MRI using image-based tracking,” Magnetic Resonance in Medicine, vol. 63, no. 1, pp. 91–105, 2010.
- C. Malamateniou, S. J. Malik, S. J. Counsell, J. M. Allsop, A. K. McGuinness, T. Hayat, K. Broadhouse, R. G. Nunes, A. M. Ederies, J. V. Hajnal, and M. A. Rutherford, “Motion-Compensation Techniques in Neonatal and Fetal MR Imaging,” American Journal of Neuroradiology, vol. 34, no. 6, pp. 1124–1136, 2013.
- B. Bellekens, V. Spruyt, R. Berkvens, R. Penne, and M. Weyn, “A benchmark survey of rigid 3D point cloud registration algorithms,” Journal on Advances in Intelligent Systems, vol. 8, pp. 118–127, 2017.
- S. Miao, Z. J. Wang, and R. Liao, “A CNN Regression Approach for Real-Time 2D/3D Registration,” IEEE Transactions on Medical Imaging, vol. 35, no. 5, pp. 1352–1363, 2016.
- S. S. Mohseni Salehi, S. Khan, D. Erdogmus, and A. Gholipour, “Real-Time Deep Pose Estimation With Geodesic Loss for Image-to-Template Rigid Registration,” IEEE Transactions on Medical Imaging, vol. 38, no. 2, pp. 470–481, 2019.
- N. Gruver, M. A. Finzi, M. Goldblum, and A. G. Wilson, “The Lie Derivative for Measuring Learned Equivariance,” in International Conference on Learning Representations, 2023.
- T. Chen, S. Kornblith, M. Norouzi, and G. Hinton, “A Simple Framework for Contrastive Learning of Visual Representations,” in International Conference on Machine Learning, 2020, pp. 1597–1607.
- T. Cohen and M. Welling, “Group Equivariant Convolutional Networks,” in International Conference on Machine Learning, 2016, pp. 2990–2999.
- M. Winkels and T. S. Cohen, “3D G-CNNs for Pulmonary Nodule Detection,” in Medical Imaging with Deep Learning, 2018.
- E. J. Bekkers, M. W. Lafarge, M. Veta, K. A. J. Eppenhof, J. P. W. Pluim, and R. Duits, “Roto-Translation Covariant Convolutional Networks for Medical Image Analysis,” in Medical Image Computing and Computer Assisted Intervention, 2018, pp. 440–448.
- M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. S. Cohen, “3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data,” in Advances in Neural Information Processing Systems, vol. 31, 2018.
- D. Moyer, E. Abaci Turk, P. E. Grant, W. M. Wells, and P. Golland, “Equivariant Filters for Efficient Tracking in 3D Imaging,” in Medical Image Computing and Computer Assisted Intervention, 2021, pp. 73–92.
- B. Billot, D. Moyer, N. Karani, M. Hoffmann, E. A. Turk, E. Grant, and P. Golland, “Equivariant and Denoising CNNs to Decouple Intensity and Spatial Features for Motion Tracking in Fetal Brain MRI,” in Medical Imaging with Deep Learning, short paper track, 2023.
- R. P. Woods, S. T. Grafton, C. J. Holmes, S. R. Cherry, and J. C. Mazziotta, “Automated Image Registration: General Methods and Intrasubject, Intramodality Validation,” Journal of Computer Assisted Tomography, vol. 22, no. 1, p. 139, 1998.
- R. W. Cox and A. Jesmanowicz, “Real-time 3D image registration for functional MRI,” Magnetic Resonance in Medicine, vol. 42, no. 6, pp. 1014–1018, 1999.
- J. G. Pipe, “Motion correction with PROPELLER MRI: application to head motion and free-breathing cardiac imaging,” Magnetic Resonance in Medicine, vol. 42, no. 5, pp. 963–969, 1999.
- D. Hill, P. Batchelor, M. Holden, and D. Hawkes, “Medical image registration,” Physics in Medicine & Biology, vol. 46, no. 3, 2001.
- M. Jenkinson and S. Smith, “A global optimisation method for robust affine registration of brain images,” Medical Image Analysis, vol. 5, no. 2, pp. 143–156, 2001.
- J. V. Hajnal, N. Saeed, A. Oatridge, E. J. Williams, I. R. Young, and G. M. Bydder, “Detection of Subtle Brain Changes Using Subvoxel Registration and Subtraction of Serial MR Images,” Journal of Computer Assisted Tomography, vol. 19, no. 5, p. 677, 1995.
- W. M. Wells, P. Viola, H. Atsumi, S. Nakajima, and R. Kikinis, “Multi-modal volume registration by maximization of mutual information,” Medical Image Analysis, vol. 1, no. 1, pp. 35–51, 1996.
- A. Roche, G. Malandain, X. Pennec, and N. Ayache, “The correlation ratio as a new similarity measure for multimodal image registration,” in Medical Image Computing and Computer-Assisted Intervention, 1998, pp. 1115–1124.
- B. Avants, N. Tustison, G. Song, P. Cook, A. Klein, and J. Gee, “A Reproducible Evaluation of ANTs Similarity Metric Performance in Brain Image Registration,” NeuroImage, vol. 54, no. 3, pp. 33–44, 2011.
- M. Modat, D. Cash, P. Daga, G. Winston, J. Duncan, and S. Ourselin, “Global image registration using a symmetric block-matching approach,” Journal of Medical Imaging, vol. 1, no. 2, pp. 24–36, 2014.
- H. Chui and A. Rangarajan, “A new point matching algorithm for non-rigid registration,” Computer Vision and Image Understanding, vol. 89, no. 2, pp. 114–141, 2003.
- K. S. Arun, T. S. Huang, and S. D. Blostein, “Least-Squares Fitting of Two 3-D Point Sets,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 9, no. 5, pp. 698–700, Sep. 1987.
- B. Horn, “Closed-form solution of absolute orientation using unit quaternions,” Journal of Optical Society, vol. 4, pp. 629–642, 1987.
- T. Tuytelaars and K. Mikolajczyk, “Local Invariant Feature Detectors: A Survey,” Foundations in Computer Graphics and Vision, vol. 3, no. 3, pp. 177–280, 2008.
- X. Pennec and J.-P. Thirion, “A Framework for Uncertainty and Validation of 3-D Registration Methods Based on Points and Frames,” International Journal of Computer Vision, vol. 25, pp. 203–229, 1997.
- P. Montesinos, V. Gouet, and R. Deriche, “Differential invariants for color images,” in International Conference on Pattern Recognition, vol. 1, Aug. 1998, pp. 838–840.
- E. Chee and Z. Wu, “AIRNet: Self-Supervised Affine Registration for 3D Medical Images using Neural Networks,” 2018, arXiv:1810.02583 [cs].
- B. D. de Vos, F. F. Berendsen, M. A. Viergever, H. Sokooti, M. Staring, and I. Išgum, “A deep learning framework for unsupervised affine and deformable image registration,” Medical Image Analysis, vol. 52, pp. 128–143, 2019.
- R. Liao, S. Miao, P. Tournemire, S. Grbic, A. Kamen, T. Mansi, and D. Comaniciu, “An Artificial Agent for Robust Image Registration,” AAAI Conference on Artificial Intelligence, vol. 31, 2016.
- A. Q. Wang, E. M. Yu, A. V. Dalca, and M. R. Sabuncu, “A robust and interpretable deep learning framework for multi-modal registration via keypoints,” Medical Image Analysis, vol. 90, p. 102962, 2023.
- W. Freeman and E. Adelson, “The design and use of steerable filters,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 13, no. 9, pp. 891–906, 1991.
- M. Reisert and H. Burkhardt, “Efficient Tensor Voting with 3D tensorial harmonics,” in Computer Vision and Pattern Recognition Workshops, 2008, pp. 1–7.
- R. Duits and E. Franken, “Left-Invariant Diffusions on the Space of Positions and Orientations and their Application to Crossing-Preserving Smoothing of HARDI images,” International Journal of Computer Vision, vol. 92, pp. 231–264, 2011.
- M. Janssen, T. Dela Haije, F. Martin, E. Bekkers, and R. Duits, “The Hessian of Axially Symmetric Functions on SE(3) and Application in 3D Image Analysis,” in Scale Space and Variational Methods in Computer Vision, 2017, pp. 643–655.
- T. S. Cohen, M. Geiger, and M. Weiler, “Intertwiners between Induced Representations (with Applications to the Theory of Equivariant Neural Networks),” 2018, arXiv:1803.10743 [cs, stat].
- D. Worrall and G. Brostow, “CubeNet: Equivariance to 3D Rotation and Translation,” in European Conference on Computer Vision, 2018, pp. 585–602.
- M. W. Lafarge, E. J. Bekkers, J. P. W. Pluim, R. Duits, and M. Veta, “Roto-translation equivariant convolutional networks: Application to histopathology image analysis,” Medical Image Analysis, vol. 68, p. 101849, 2021.
- D. E. Worrall, S. J. Garbin, D. Turmukhambetov, and G. J. Brostow, “Harmonic Networks: Deep Translation and Rotation Equivariance,” in Computer Vision and Pattern Recognition, 2017, pp. 5028–5037.
- T. S. Cohen and M. Welling, “Steerable CNNs,” in International Conference on Learning Representations, 2017.
- G. Cesa, L. Lang, and M. Weiler, “A Program to Build E(N)-Equivariant Steerable CNNs,” in International Conference on Learning Representations, 2022.
- J. J. Bouza, C.-H. Yang, D. Vaillancourt, and B. C. Vemuri, “A Higher Order Manifold-Valued Convolutional Neural Network with Applications to Diffusion MRI Processing,” in Information Processing in Medical Imaging, 2021, pp. 304–317.
- A. Elaldi, G. Gerig, and N. Dey, “E(3) x SO(3) - Equivariant Networks for Spherical Deconvolution in Diffusion MRI,” in Medical Imaging with Deep Learning, 2023.
- X. Huang, M.-Y. Liu, S. Belongie, and J. Kautz, “Multimodal Unsupervised Image-to-image Translation,” in European Conference on Computer Vision, 2018, pp. 172–189.
- A. Chartsias, T. Joyce, G. Papanastasiou, S. Semple, M. Williams, D. E. Newby, R. Dharmakumar, and S. A. Tsaftaris, “Disentangled representation learning in cardiac image analysis,” Medical Image Analysis, vol. 58, p. 101535, 2019.
- X. Liu, P. Sanchez, S. Thermos, A. Q. O’Neil, and S. A. Tsaftaris, “Learning disentangled representations in the imaging domain,” Medical Image Analysis, vol. 80, p. 102516, 2022.
- A. E. Ilesanmi and T. O. Ilesanmi, “Methods for image denoising using convolutional neural network: a review,” Complex & Intelligent Systems, vol. 7, no. 5, pp. 179–198, 2021.
- B. Billot, D. N. Greve, O. Puonti, A. Thielscher, K. Van Leemput, B. Fischl, A. V. Dalca, and J. E. Iglesias, “SynthSeg: Segmentation of brain MRI scans of any contrast and resolution without retraining,” Medical Image Analysis, vol. 86, p. 102789, 2023.
- M. Hoffmann, B. Billot, D. N. Greve, J. E. Iglesias, B. Fischl, and A. V. Dalca, “SynthMorph: Learning Contrast-Invariant Registration Without Acquired Images,” IEEE Transactions on Medical Imaging, vol. 41, no. 3, pp. 543–558, 2022.
- M. Hoffmann, A. Hoopes, B. Fischl, and A. Dalca, “Anatomy-specific acquisition-agnostic affine registration learned from fictitious images,” in Medical Imaging 2023: Image Processing, Apr. 2023.
- W. Kabsch, “A solution for the best rotation to relate two sets of vectors,” Acta Crystallographica: Crystal Physics, Diffraction, Theoretical and General Crystallography, vol. 32, no. 5, pp. 922–923, 1976.
- J. Ashburner and K. J. Friston, “Unified segmentation,” NeuroImage, vol. 26, no. 3, pp. 839–851, 2005.
- O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in Medical Image Computing and Computer-Assisted Intervention, 2015, pp. 234–241.
- S. Ioffe and C. Szegedy, “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift,” in International Conference on Machine Learning, 2015, pp. 448–456.
- A. Kamath, J. Willmann, N. Andratschke, and M. Reyes, “Do We Really Need that Skip-Connection? Understanding Its Interplay with Task Complexity,” in Medical Image Computing and Computer Assisted Intervention, 2023, pp. 302–311.
- B. Billot, C. Magdamo, Y. Cheng, S. E. Arnold, S. Das, and J. E. Iglesias, “Robust machine learning segmentation for large-scale analysis of heterogeneous clinical brain MRI datasets,” Proceedings of the National Academy of Sciences, vol. 120, no. 9, 2023.
- D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” 2014, arXiv:1412.6980 [cs].
- A. Paszke et al., “PyTorch: an imperative style, high-performance deep learning library,” in International Conference on Neural Information Processing Systems, 2019, pp. 8026–8037.
- F. Pérez-García, R. Sparks, and S. Ourselin, “TorchIO: A Python library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning,” Computer Methods and Programs in Biomedicine, vol. 208, p. 106236, 2021.
- D. C. Van Essen et al., “The Human Connectome Project: A data acquisition perspective,” NeuroImage, vol. 62, no. 4, pp. 222–231, 2012.
- M. Hoffmann, D. C. Moyer, L. Zhang, P. Golland, B. Gagoski, P. E. Grant, and A. J. van der Kouwe, “Learning-based automatic field-of-view positioning for fetal-brain MRI,” International Society for Magnetic Resonance in Medicine, vol. 29, p. 1362, 2021.
- R. H. Byrd, P. Lu, J. Nocedal, and C. Zhu, “A Limited Memory Algorithm for Bound Constrained Optimization,” Journal on Scientific Computing, vol. 16, no. 5, pp. 1190–1208, 1995.
- P. A. Yushkevich, J. Piven, H. C. Hazlett, R. G. Smith, S. Ho, J. C. Gee, and G. Gerig, “User-guided 3D active contour segmentation of anatomical structures: significantly improved efficiency and reliability,” NeuroImage, vol. 31, no. 3, pp. 1116–1128, 2006.
- M. Finzi, G. Benton, and A. G. Wilson, “Residual Pathway Priors for Soft Equivariance Constraints,” in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 3037–3049.
- Z. Chi, Z. Cong, C. Wang, Y. Liu, E. Abaci Turk, E. Grant, M. Abulnaga, P. Golland, and N. Dey, “Dynamic Neural Fields for Learning Atlases of 4D Fetal MRI Time-series,” in Medical Imaging meets NeurIPS workshop, 2023.
Sponsor
Paper Prompts
Sign up for free to create and run prompts on this paper using GPT-5.
Top Community Prompts
Collections
Sign up for free to add this paper to one or more collections.