Data-Agnostic Augmentations for Unknown Variations: Out-of-Distribution Generalisation in MRI Segmentation

TL;DR

The paper tackles the challenge of out-of-distribution generalization in MRI segmentation by evaluating data-agnostic augmentations MixUp and Auxiliary Fourier Augmentation (AFA) within nnU-Net across cardiac cine MRI and prostate bpMRI datasets. It systematically introduces a broad, severity-aware image variation model and compares base augmentations, MixUp, and AFA (and combinations) under synthetic corruptions and real-world distribution shifts. Results show that MixUp and AFA, particularly when combined with base augmentations, yield robust improvements in transformed data while preserving performance on clean data, supported by kVGM-based analysis of latent representations. These findings suggest simple, easy-to-integrate augmentation strategies can enhance reliability of medical segmentation in diverse clinical settings, with practical implications for deploying robust MRI-deployable segmentation models.

Abstract

Medical image segmentation models are often trained on curated datasets, leading to performance degradation when deployed in real-world clinical settings due to mismatches between training and test distributions. While data augmentation techniques are widely used to address these challenges, traditional visually consistent augmentation strategies lack the robustness needed for diverse real-world scenarios. In this work, we systematically evaluate alternative augmentation strategies, focusing on MixUp and Auxiliary Fourier Augmentation. These methods mitigate the effects of multiple variations without explicitly targeting specific sources of distribution shifts. We demonstrate how these techniques significantly improve out-of-distribution generalization and robustness to imaging variations across a wide range of transformations in cardiac cine MRI and prostate MRI segmentation. We quantitatively find that these augmentation methods enhance learned feature representations by promoting separability and compactness. Additionally, we highlight how their integration into nnU-Net training pipelines provides an easy-to-implement, effective solution for enhancing the reliability of medical segmentation models in real-world applications.

Figures (15)

• Figure 1: Corruptions (severity level 3) in cardiac cine MRI images and the T2w channel of prostate bpMRI images as a result of our image variation model.
• Figure 2: Trend of DSC per severity for test sets corrupted with bias field, ghosting, k-space subsampling, Rician noise and spike noise. We notice that MixUp and AFA are effective in mitigating the challenges posed by complex image corruptions.
• Figure 3: PCA projection of learned features, for final features from nnU-Net trained with different augmentation techniques for samples from the transformed test sets (top: ACDC, bottom: P158) with the corresponding kVGM metric.
• Figure 4: Visualization of the 14 data variations, alongside the original image (top-left) for a test sample in the ACDC dataset. All transforms visualised at severity 3.
• Figure 5: Visualization of the 14 data variations, alongside the original image (top-left) for a test sample in the P158 dataset. All transforms visualised at severity 3.