Test-Time Modality Generalization for Medical Image Segmentation

TL;DR

This paper introduces a novel framework, TTMG, that enhances medical image segmentation across unseen modalities by using modality-aware style projection and instance whitening, outperforming existing methods.

Contribution

The paper proposes the TTMG framework with MASP and MSIW components to improve generalization to unseen modalities in medical image segmentation.

Findings

Achieves superior segmentation performance across eleven datasets.

Effectively generalizes to four different medical imaging modalities.

Outperforms existing domain generalization techniques.

Abstract

Generalizable medical image segmentation is essential for ensuring consistent performance across diverse unseen clinical settings. However, existing methods often overlook the capability to generalize effectively across arbitrary unseen modalities. In this paper, we introduce a novel Test-Time Modality Generalization (TTMG) framework, which comprises two core components: Modality-Aware Style Projection (MASP) and Modality-Sensitive Instance Whitening (MSIW), designed to enhance generalization in arbitrary unseen modality datasets. The MASP estimates the likelihood of a test instance belonging to each seen modality and maps it onto a distribution using modality-specific style bases, guiding its projection effectively. Furthermore, as high feature covariance hinders generalization to unseen modalities, the MSIW is applied during training to selectively suppress modality-sensitive information while retaining modality-invariant features. By integrating MASP and MSIW, the TTMG framework demonstrates robust generalization capabilities for medical image segmentation in unseen modalities a challenge that current methods have largely neglected. We evaluated TTMG alongside other domain generalization techniques across eleven datasets spanning four modalities (colonoscopy, ultrasound, dermoscopy, and radiology), consistently achieving superior segmentation performance across various modality combinations.