Source-Free Domain Adaptive Fundus Image Segmentation with Class-Balanced Mean Teacher

TL;DR

This paper introduces CBMT, a source-free domain adaptation method for fundus image segmentation that uses a class-balanced mean teacher approach to improve stability and handle class imbalance, achieving superior results.

Contribution

The paper proposes a novel class-balanced mean teacher model with a weak-strong augmentation scheme and loss calibration for stable, effective domain adaptation in fundus segmentation.

Findings

CBMT outperforms existing methods on multiple benchmarks.

The weak-strong augmentation scheme stabilizes pseudo-label generation.

Class-balanced loss calibration improves foreground class segmentation.

Abstract

This paper studies source-free domain adaptive fundus image segmentation which aims to adapt a pretrained fundus segmentation model to a target domain using unlabeled images. This is a challenging task because it is highly risky to adapt a model only using unlabeled data. Most existing methods tackle this task mainly by designing techniques to carefully generate pseudo labels from the model's predictions and use the pseudo labels to train the model. While often obtaining positive adaption effects, these methods suffer from two major issues. First, they tend to be fairly unstable - incorrect pseudo labels abruptly emerged may cause a catastrophic impact on the model. Second, they fail to consider the severe class imbalance of fundus images where the foreground (e.g., cup) region is usually very small. This paper aims to address these two issues by proposing the Class-Balanced Mean Teacher (CBMT) model. CBMT addresses the unstable issue by proposing a weak-strong augmented mean teacher learning scheme where only the teacher model generates pseudo labels from weakly augmented images to train a student model that takes strongly augmented images as input. The teacher is updated as the moving average of the instantly trained student, which could be noisy. This prevents the teacher model from being abruptly impacted by incorrect pseudo-labels. For the class imbalance issue, CBMT proposes a novel loss calibration approach to highlight foreground classes according to global statistics. Experiments show that CBMT well addresses these two issues and outperforms existing methods on multiple benchmarks.