MoreStyle: Relax Low-frequency Constraint of Fourier-based Image Reconstruction in Generalizable Medical Image Segmentation

TL;DR

This work tackles single-source domain generalization in medical image segmentation by introducing MoreStyle, a plug‑and‑play augmentation framework that relaxes low‑frequency Fourier constraints to diversify image styles. It combines an adversarial style augmentation pipeline with an auxiliary reconstruction decoder, a Fourier Spectrum Diversity Loss, and an uncertainty-weighted Intersection-Union loss to maximize gain from style shifts while stabilizing training. The key contributions are the ASA mechanism to generate diverse styled images, the $\mathcal{L}_{FSD}$ loss for broad Fourier-domain variation, and the $\mathcal{L}_{UIU}$ loss to focus learning on style-induced hard pixels, yielding consistent improvements on retinal OC/OD and prostate segmentation, and benefiting integration with MedSAM or CCSDG. The approach demonstrates faster convergence and practical potential for deploying robust segmentation models across varied clinical imaging settings, though it has limitations with geometric variability which points to avenues for future work with geometric augmentation.

Abstract

The task of single-source domain generalization (SDG) in medical image segmentation is crucial due to frequent domain shifts in clinical image datasets. To address the challenge of poor generalization across different domains, we introduce a Plug-and-Play module for data augmentation called MoreStyle. MoreStyle diversifies image styles by relaxing low-frequency constraints in Fourier space, guiding the image reconstruction network. With the help of adversarial learning, MoreStyle further expands the style range and pinpoints the most intricate style combinations within latent features. To handle significant style variations, we introduce an uncertainty-weighted loss. This loss emphasizes hard-to-classify pixels resulting only from style shifts while mitigating true hard-to-classify pixels in both MoreStyle-generated and original images. Extensive experiments on two widely used benchmarks demonstrate that the proposed MoreStyle effectively helps to achieve good domain generalization ability, and has the potential to further boost the performance of some state-of-the-art SDG methods. Source code is available at https://github.com/zhaohaoyu376/morestyle.

Figures (3)

• Figure 1: We adopt the first residual block of pre-trained ResNet-18 he2016deep to extract image features. (a) Normalized feature distribution on the source images (the upper) and images generated by proposed MoreStyle (the lower). (b) Visualization of feature distribution on the source images and generated images given by methods including MaxStyle chen2022maxstyle, FDA yang2020fda, FACT xu2021fourier, and MoreStyle on RIGA+ dataset almazroa2018retinalhu2022domaindecenciere2014feedback with tSNE. (c) Visualization of source image (the upper) and style-augmented images generated by proposed MoreStyle (the lower).
• Figure 2: Pipeline of the proposed MoreStyle. In our proposed MoreStyle framework, we first generate style-augmented images $\hat{x}_i$ using Adversarial Style Augmentation (ASA) and a reconstruction decoder $\mathcal{D_\theta}$, with a noise encoder $\mathcal{N_\epsilon}$ generating style mixing $\gamma_i$ and noise perturbation $\beta_i$ through adversarial training to affect $\mathcal{D_\theta}$. This image reconstruction is guided by the Fourier Spectrum Diversity (FSD) loss $\mathcal{L}_{FSD}$. Lastly, both style-augmented $\hat{x}_i$ and source images $x_i$ are input into the segmentation network, supervised by a customized Uncertainty-weighted Intersection-Union (UIU) loss $\mathcal{L}_{UIU}$.
• Figure 3: Visualization of results by MoreStyle and several SOTA methods.