UPL-TTA: Uncertainty-Aware Pseudo Label Guided Fully Test Time Adaptation for Fetal Brain Segmentation

Abstract

Test Time Adaptation (TTA) is promising to improve a deep learning model’s robustness when encountering images from an unseen domain. Existing TTA methods are with low performance due to the insufficient supervision signal from unannotated target domain images, or limited by specific requirements on the pre-training strategy and network structure in the source domain. We aim to separate the pre-training in the source domain and adaptation in the target domain, in order to achieve high-performance and more generalizable TTA without assumptions on the pre-training strategy. To solve this problem, we propose UPL-TTA, an Uncertainty-aware Pseudo Label guided fully Test Time Adaptation method. Specifically, we introduce Test Time Growing (TTG) to duplicate the prediction head of the source model with perturbations at image and feature levels in the target domain. The different predictions obtained in these duplicated prediction heads are used to obtain pseudo labels for the unlabeled target domain images as well as their uncertainty maps, which can identify reliable pseudo labels. Pixels with unreliable pseudo labels are regularized by imposing entropy minimization on the mean prediction of the multiple heads. UPL-TTA was validated bidirectionally on a cross-modality fetal brain segmentation dataset. Compared with no adaptation, it significantly improved the average Dice in the two different target domains by 3.95% and 6.12%, respectively, and outperformed several state-of-the-art TTA methods.