Color-Quality Invariance for Robust Medical Image Segmentation

TL;DR

This paper introduces novel normalization and loss techniques to improve the robustness of medical image segmentation models against color and quality variations across different domains.

Contribution

The work presents the DCIN module and CQG loss, which together enhance domain generalization in medical image segmentation under color and quality shifts.

Findings

Up to 32.3-point increase in Dice score over baseline

Significant improvement in segmentation robustness across target domains

Effective handling of color and quality variations in unseen domains

Abstract

Single-source domain generalization (SDG) in medical image segmentation remains a significant challenge, particularly for images with varying color distributions and qualities. Previous approaches often struggle when models trained on high-quality images fail to generalize to low-quality test images due to these color and quality shifts. In this work, we propose two novel techniques to enhance generalization: dynamic color image normalization (DCIN) module and color-quality generalization (CQG) loss. The DCIN dynamically normalizes the color of test images using two reference image selection strategies. Specifically, the DCIN utilizes a global reference image selection (GRIS), which finds a universal reference image, and a local reference image selection (LRIS), which selects a semantically similar reference image per test sample. Additionally, CQG loss enforces invariance to color and quality variations by ensuring consistent segmentation predictions across transformed image pairs. Experimental results show that our proposals significantly improve segmentation performance over the baseline on two target domain datasets, despite being trained solely on a single source domain. Notably, our model achieved up to a 32.3-point increase in Dice score compared to the baseline, consistently producing robust and usable results even under substantial domain shifts. Our work contributes to the development of more robust medical image segmentation models that generalize across unseen domains. The implementation code is available at https://github.com/RaviShah1/DCIN-CQG.