Class-Aware Feature Alignment for Domain Adaptative Mitochondria Segmentation

Abstract

Unsupervised domain adaptation (UDA) has gained great popularity in mitochondria segmentation, aiming to improve the adaptability of models from the labeled source domain to the unlabeled target domain via domain alignment. However, existing UDA methods only focus on aligning domains on the prediction level, while ignoring the feature space containing more adequate information than the predictions. In this paper, we propose a class-aware domain adaptation method for mitochondria segmentation on the feature level, which relies on the prototype representation to achieve more fine-grained alignment. In particular, we first extract the feature centroids of classes from the source domain as prototypes. Leveraging the extracted prototypes as a bridge, we constrain that features belonging to the same class but from different domains are pulled closer to each other, achieving the class-aware alignment. Meanwhile, we derive a segmentation prediction directly from feature space based on the distance between target features and source prototypes. By incorporating a pseudo label to supervise the learning of this prediction, the feature distribution gap across domains is further reduced. Furthermore, to take full advantage of the potential of target domain, we propose an intra-domain consistency constraint to maintain consistent predictions of samples perturbed differently from the target image. Extensive experiments on different datasets demonstrate the superiority of our proposed method over existing UDA methods. Code is available at https://github.com/Danyin813/CAFA.