Shape-Aware Joint Distribution Alignment for Cross-Domain Image Segmentation

We present an unsupervised domain adaptation method for image segmentation which aligns high-order statistics, computed for the source and target domains, encoding domain-invariant spatial relationships between segmentation classes. Our method first estimates the joint distribution of predictions for pairs of pixels whose relative position corresponds to a given spatial displacement. Domain adaptation is then achieved by aligning the joint distributions of source and target images, computed for a set of displacements. Two enhancements of this method are proposed. The first one uses an efficient multi-scale strategy that enables capturing long-range relationships in the statistics. The second one extends the joint distribution alignment loss to features in intermediate layers of the network by computing their cross-correlation. We test our method on the task of unpaired multi-modal cardiac segmentation using the Multi-Modality Whole Heart Segmentation Challenge dataset and prostate segmentation task where images from two datasets are taken as data in different domains. Our results show the advantages of our method compared to recent approaches for cross-domain image segmentation. Code is available at https://github.com/WangPing521/Domain_adaptation_shape_prior.