Unsupervised Domain Adaptation via Similarity-based Prototypes for Cross-Modality Segmentation

TL;DR

This work tackles unsupervised domain adaptation for cross-modality medical image segmentation by introducing a unified, end-to-end framework that explicitly aligns features to class-wise prototypes. It leverages a cycle-consistent architecture with two generators to translate between domains and computes class-wise prototypes from embedded representations, storing them in feature dictionaries to enable cross-image, contrastive learning. The method introduces a cosine-based class-wise similarity loss and a dictionary-driven prototype contrastive loss, improving intra-class consistency and inter-class separability while reducing domain gaps. Empirical results on the MRI-to-CT task in the MMWH dataset show state-of-the-art Dice scores and competitive boundary metrics, validating the effectiveness of prototype-based alignment for cross-modality segmentation in clinical imaging contexts.

Abstract

Deep learning models have achieved great success on various vision challenges, but a well-trained model would face drastic performance degradation when applied to unseen data. Since the model is sensitive to domain shift, unsupervised domain adaptation attempts to reduce the domain gap and avoid costly annotation of unseen domains. This paper proposes a novel framework for cross-modality segmentation via similarity-based prototypes. In specific, we learn class-wise prototypes within an embedding space, then introduce a similarity constraint to make these prototypes representative for each semantic class while separable from different classes. Moreover, we use dictionaries to store prototypes extracted from different images, which prevents the class-missing problem and enables the contrastive learning of prototypes, and further improves performance. Extensive experiments show that our method achieves better results than other state-of-the-art methods.

Figures (5)

• Figure 1: Illustration of our adaptation procedure. On the one hand, our method performs class-wise adaptation to align semantic features to their prototypes, on the other hand, we align class-wise prototypes across domains using the contrastive loss.
• Figure 2: Our framework has a cycle structure, and mainly consists of $G_S$ and $G_T$. These modules have the same structure and output the translated image $\hat{x}$, segmentation result $\hat{y}$ and embedded projection $\hat{z}$. Prototypes $c$ are obtained by performing a class-wise average operation on $\hat{z}$ under supervision from $\hat{y}$. During training, only $c_s$ and $c_{s \rightarrow t}$ are stored into feature dictionaries $B_s$ and $B_t$. Besides the widely used circle consistency loss ${L_{cycle}}$, segmentation loss $L_{seg}$, adversarial loss $L_{adv}^{img}$ and $L_{adv}^{seg}$, we additionally use the proposed loss $L_{sim}$ and $L_{cl}$ to perform explicit feature alignment.
• Figure 3: Details of $G_S$ and prototypes $c_s$. $G_S$ consists of a feature encoder $E_s$, an image generator $T_t$, and symmetric heads $F_s$ and $P_s$. These heads have same structures but different output channels. $F_s$ outputs the segmentation results ${\hat{y}_s}$ ,while $P_s$ outputs the embedded representation $\hat{z}_s$. The prototypes $c_s$ are obtained by performing class-wise average on $\hat{z}_s$ under supervision of ${\hat{y}_s}$. And we store these prototypes in Dictionary $B_s$.
• Figure 4: Visual comparison of representative methods. The structures of MYO, LAC, LVC, AA are denoted by yellow, green, red, blue colors, respectively.
• Figure 5: t-SNE visualization of foreground features in Fig.\ref{['fig:visual']}. Left: The results without the proposed losses, Right: The results with the proposed losses.