Domain Composition and Attention for Unseen-Domain Generalizable Medical Image Segmentation

TL;DR

This paper introduces DCA-Net, a novel domain generalization framework for medical image segmentation that uses domain composition and attention mechanisms to improve performance across unseen domains, validated on multi-institutional prostate MRI data.

Contribution

The paper proposes a new domain composition and attention-based network with a divergence constraint to enhance domain representation and generalization in medical image segmentation.

Findings

Outperforms state-of-the-art methods on prostate MRI segmentation

Effectively generalizes to unseen domains

Utilizes a divergence constraint to diversify basis representations

Abstract

Domain generalizable model is attracting increasing attention in medical image analysis since data is commonly acquired from different institutes with various imaging protocols and scanners. To tackle this challenging domain generalization problem, we propose a Domain Composition and Attention-based network (DCA-Net) to improve the ability of domain representation and generalization. First, we present a domain composition method that represents one certain domain by a linear combination of a set of basis representations (i.e., a representation bank). Second, a novel plug-and-play parallel domain preceptor is proposed to learn these basis representations and we introduce a divergence constraint function to encourage the basis representations to be as divergent as possible. Then, a domain attention module is proposed to learn the linear combination coefficients of the basis representations. The result of linear combination is used to calibrate the feature maps of an input image, which enables the model to generalize to different and even unseen domains. We validate our method on public prostate MRI dataset acquired from six different institutions with apparent domain shift. Experimental results show that our proposed model can generalize well on different and even unseen domains and it outperforms state-of-the-art methods on the multi-domain prostate segmentation task.