Target and Task specific Source-Free Domain Adaptive Image Segmentation

1Johns Hopkins University,

Abstract

Solving the domain shift problem during inference is essential in medical imaging as most deep-learning based solutions suffer from it. In practice, domain shifts are tackled by performing Unsupervised Domain Adaptation (UDA), where a model is adapted to an unlabeled target domain by leveraging the labelled source domain. In medical scenarios, the data comes with huge privacy concerns making it difficult to apply standard UDA techniques. Hence, a closer clinical setting is Source-Free UDA (SFUDA), where we have access to source trained model but not the source data during adaptation. Methods trying to solve SFUDA typically address the domain shift using pseudo-label based self-training techniques. However due to domain shift, these pseudo-labels are usually of high entropy and denoising them still does not make them perfect labels to supervise the model. Therefore, adapting the source model with noisy pseudo labels reduces its segmentation capability while addressing the domain shift. To this end, we propose a two-stage approach for source-free domain adaptive image segmentation: 1) Target-specific adaptation followed by 2) Task-specific adaptation. In the first stage, we focus on generating target-specific pseudo labels while suppressing high entropy regions by proposing an Ensemble Entropy Minimization loss. We also introduce a selective voting strategy to enhance pseudo-label generation. In the second stage, we focus on adapting the network for task-specific representation by using a teacher-student self-training approach based on augmentation-guided consistency. We evaluate our proposed method on both 2D fundus datasets and 3D MRI volumes across 7 different domain shifts where we achieve better performance than recent UDA and SF-UDA methods for medical image segmentation.

Method

Overview of the proposed framework.
In target-specific adaptation, we propose an ensemble entropy minimization loss to reduce the high entropy regions of pseudo-labels. We also introduce a selective voting strategy to enhance pseudo-labels for better supervision. Here, we use multiple augmentations to get regions that are consistently of high entropy and enhance the pseudo labels based on that information. Thus in Stage I, we adapt to the target domain by ensemble entropy minimization and supervise with enhanced pseudo-label without losing task-specific information.

In task-specific adaptation, we use a strong-weak augmentation based teacher-student self-training framework. Here, ensuring that the predictions are consistent across strong and weak augmentations helps the model learn task-specific information. We also introduce a augmentation-guided feature consistency loss that helps regularize the network for segmentation. Thus in Stage II, we perform student-teacher pseudo-label based self-training to enhance task-specific information of the entropy minimized model from Stage I.

BibTeX

 
Website Template taken from Nerfies