Domain Generalization for Robust MS Lesion Segmentation

Recently, deep-learning methods have achieved human-level performance on multiple sclerosis (MS) lesion segmentation. However, most established methods are not robust enough for practical use in the real world. They cannot generalize well to images obtained from different clinical sites, or if training and testing datasets contain different MRI modalities. To address these robustness issues, and to bring the deep neural networks closer to clinical use, we propose the addition of data augmentation and modality dropout during training for achieving unsupervised domain generalization. We hypothesize that employing data augmentations can close the gap between different datasets and render the trained models more generalizable. We further hypothesize that the random dropout technique can help the model learn to predict results given any combination of MRI modalities. We conducted an extensive set of comparisons using three publicly available datasets and demonstrate that our method performs better than the baseline without any augmentation, and approaches the performance of fully supervised methods. To provide a fair comparison with other MS lesion segmentation methods, we evaluate our methods on the test set of the Longitudinal MS Lesion Segmentation Challenge using the models trained on the other two datasets. The overall score of our approach is substantially higher than the current transfer-learning-based methods and is comparable to the state-of-the-art supervised methods.