A Unified Deep Learning Framework for Motion Correction in Medical Imaging

Jian Wang, Razieh Faghihpirayesh, Danny Joca, Polina Golland, Ali Gholipour

Published: 2024/9/21

Abstract

Deep learning has shown significant value in image registration, however, current techniques are either limited by the type and range of motion they can handle, or require iterative inference and/or retraining for new imaging data. To address these limitations, we introduce UniMo, a Unified Motion Correction framework that leverages deep neural networks to correct diverse motion in medical imaging. UniMo employs an alternating optimization scheme for a unified loss function to train an integrated model of 1) an equivariant neural network for global rigid motion correction and 2) an encoder-decoder network for local deformations. It features a geometric deformation augmenter that 1) enhances the robustness of global correction by addressing local deformations from non-rigid motion or geometric distortions, and 2) generates augmented data to improve training. UniMo is a hybrid model that uses both image intensities and shapes to achieve robust performance amid appearance variations, and therefore generalizes to multiple imaging modalities without retraining. We trained and tested UniMo to track motion in fetal magnetic resonance imaging, a challenging application due to 1) both large rigid and non-rigid motion, and 2) wide variations in image appearance. We then evaluated the trained model, without retraining, on MedMNIST, lung CT, and BraTS datasets. Results show that UniMo surpassed existing motion correction methods in accuracy, and notably enabled one-time training on a single modality while maintaining high stability and adaptability across unseen datasets. By offering a unified solution to motion correction, UniMo marks a significant advance in medical imaging, especially in applications with combined bulk and local motion. The code is available at: https://github.com/IntelligentImaging/UNIMO