A Unified Deep Learning Framework for Motion Correction in Medical Imaging

Abstract

Deep learning has shown significant value in medical image registration for motion correction, 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 for various types of motion in medical imaging. UniMo exploits 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 to correct local deformations. It features a geometric deformation augmenter that 1) enhances the robustness of global motion correction by addressing any local deformations, and 2) generates augmented data to improve the training process. UniMo is a hybrid model that uses both image intensities and shapes to achieve robust performance amid image appearance variations, and, therefore, it generalizes well to various medical imaging modalities without a need for network retraining. We trained and tested UniMo to track motion in fetal magnetic resonance imaging. Then we tested the trained model, without retraining, on various image modalities from three public datasets, including MedMNIST, lung CT, and BraTS. The results show that UniMo surpassed existing motion correction methods in terms of accuracy, and, notably, it enabled one-time training on a single modality while maintaining high stability and adaptability for inference across multiple unseen imaging datasets. By offering a unified solution, UniMo marks a significant advantage in challenging applications with a mixture of bulk motion and local deformations. https://github.com/IntelligentImaging/UNIMO