Partial Unbalanced Feature Transport for Cross-Modality Cardiac Image Segmentation.

Abstract

Deep learning based approaches have achieved great success on the automatic cardiac image segmentation task. However, the achieved segmentation performance remains limited due to the significant difference across image domains, which is referred to as domain shift. Unsupervised domain adaptation (UDA), as a promising method to mitigate this effect, trains a model to reduce the domain discrepancy between the source (with labels) and the target (without labels) domains in a common latent feature space. In this work, we propose a novel framework, named Partial Unbalanced Feature Transport (PUFT), for cross-modality cardiac image segmentation. Our model facilities UDA leveraging two Continuous Normalizing Flow-based Variational Auto-Encoders (CNF-VAE) and a Partial Unbalanced Optimal Transport (PUOT) strategy. Instead of directly using VAE for UDA in previous works where the latent features from both domains are approximated by a parameterized variational form, we introduce continuous normalizing flows (CNF) into the extended VAE to estimate the probabilistic posterior and alleviate the inference bias. To remove the remaining domain shift, PUOT exploits the label information in the source domain to constrain the OT plan and extracts structural information of both domains, which are often neglected in classical OT for UDA. We evaluate our proposed model on two cardiac datasets and an abdominal dataset. The experimental results demonstrate that PUFT achieves superior performance compared with state-of-the-art segmentation methods for most structural segmentation.