A New Bidirectional Unsupervised Domain Adaptation Segmentation Framework

Deep learning based models have become dominant for computer-aided automated analysis of medical images and achieved great success in recent years [1, 2, 6, 13, 17, 19]. Usually, a great quantity of manual annotations are required for training deep learning models, yet the annotation process is known to be expertise-demanding, labor-intensive, and time-consuming. Multimodal imaging—which plays an important role and provides valuable complementary information in disease diagnosis, prognosis, and treatment planning in clinical practice nowadays—may further increase the demand for annotations, as a model well-trained with data of one specific modality often performs poorly on another due to domain shift [20]. Unsupervised domain adaptation (UDA) [7, 8] is a quickly rising technique which aims to tackle the domain shift problem by adapting models trained in one domain (the source domain) to another (the target domain) without manual annotations in the latter [3, 16, 18, 20]. The concept of UDA is readily applicable to the lack-of-annotation problem in multimodal medical imaging [4, 5]. For example, Chen et al. [4] applied UDA for model adaptation from cardiac magnetic resonance imaging (MRI) to cardiac computed tomography (CT) based on image-level alignment. In spite of impressive results obtained for adapting MRI-trained models for CT data [4, 5], few existing methods simultaneously concern the opposite, i.e., CT to MRI. In practice, situations may arise in which the CT data have already been annotated while the MRI data are not, and the adaptation of CT-trained models for MRI becomes useful. Therefore, a comprehensive UDA method that is able to effectively accomplish bidirectional adaptations (termed BiUDA in this work) is of great clinical value.

Given the importance of BiUDA, we conducted experiments to examine the capabilities of several state-of-the-art (SOTA) UDA methods [3, 4, 5, 16, 18, 20] for bidirectional adaptations between MRI and CT using the Multi-Modality Whole Heart Segmentation (MMWHS) challenge dataset [24] and the Multi-Modality Abdominal Segmentation (MMAS) dataset [12, 14]. Taking MMWHS dataset as an example, we show the bidirectional adaptation results in Fig. 2. In general, these methods worked effectively for adapting MRI-trained models for CT data. When reversing the adaptation direction, however, most of them suffered a dramatic drop in performance and failed to produce comparable results. Similar results can also be found on the MMAS dataset. We call this phenomenon domain drop, which intuitively reflects the substantial performance gap between the bidirectional adaptations in BiUDA. Recently, Dou et al. [5] improved performance of the CT-MRI BiUDA by switching the early layers of the two encoders to fit individual modality and sharing the higher layers between two modalities. It is presumed that the improvement was due to the isolation of low-level features of each modality. Motivated by the presumption, we propose a novel framework based on the concept of disentangled representation learning (DRPL) [11, 15], to address the domain drop problem in BiUDA. Specifically, our framework decomposes an input image into a content code and a pattern code, with the former representing domain-invariant characteristics shared by different modalities (e.g., shape) and the latter representing domain-specific features isolated from each other (e.g., appearance).

Our main contributions include:

• 1)

  We propose a novel BiUDA framework based on DRPL, which effectively mitigates the domain drop problem and achieves SOTA results for both adaptation directions in experiments on the MMWHS and MMAS datasets.

• 2)

  Unlike existing works [3, 11, 15] that adopted a separate pattern encoder for each modality, we design a domain-aware pattern encoder to unify the encoding process of different domains in a single module with a domain controller. This design not only reduces parameters but also improves performance of the network. In addition, a content-pattern consistency loss is proposed to avoid the distortion of the content and pattern codes.

• 3)

  For better performance in the UDA segmentation task, we propose a label consistency loss, utilizing the target-domain-styled images and corresponding source-domain annotations for auxiliary training.

Datasets. The proposed BiUDA framework is evaluated using the MMWHS challenge dataset [24] and the MMAS dataset. The MMWHS dataset includes 20 MRI (47–127 slices per scan) and 20 CT (142–251 slices per scan) cardiac scans. Four cardiac anatomic structures, including ascending aorta (AA), left atrium blood cavity (LABC), left ventricle blood cavity (LVBC), and left ventricle myocardium (LVM), are annotated. The MMAS dataset includes 20 MRI scans (21–33 slices per scan) from the CHAOS Challenge [12] and 30 CT scans (35–117 slices per scan) from [14]. Multiple organs are manually annotated, including liver, right kidney (R-Kid), left kidney (L-Kid), and spleen. For a fair comparison, every input slice is resized to 256$\times$256 pixels and augmented in the same way as SIFA [4], including random crop, flip and rotation. A 5-fold cross-validation strategy is employed to test our framework.

Evaluation metrics. The Dice coefficient (Dice) and F1 score are used as the basic evaluation metrics. The performance upper-bounds are established by separately training and testing two segmentation networks (one for each modality) using data of the same modality, and denoted by ‘M2M’ (MRI to MRI) and ‘C2C’ (CT to CT), respectively. Followed by the same adaptation direction reported in  [4], we define the adaptation from MRI (as suorce domain) to CT (as target domain) as the forward adaptation, and vice versa. Since the levels of difficulty are markedly different for MRI- and CT-based cardiac segmentation due to distinct modal characteristics (see the rows for M2M and C2C in Table 1), it would be difficult to directly compare Dice or F1 scores obtained via the forward and backward adaptations. Instead, we resort to the performance drop, which is computed by subtracting the UDA performance from the corresponding target domain upper-bound, e.g., subtracting MRI-to-CT UDA performance from C2C. Lastly, to intuitively reflect the quality of BiUDA using a single metric, we further calculate the average performance drop by averaging the bidirectional performance drops.

Implementation. The content encoder $E_{c}$ is based on PSP-101 [22], accompanied by a fully convolutional network as the segmenter $S$. The domain-aware pattern encoder $E_{p}$ comprises several downsampling units followed by a global average pooling layer and a $1\times 1$ convolution layer. Each unit contains a convolution layer with $stride=2$, followed by a batch normalization layer and a ReLU layer. The discriminator $D$ has a similar structure to $E_{p}$, with its units consisting of a convolution layer with $stride=2$, an instance normalization layer and a leaky-ReLU layer. The generator $G$ is composed of a set of residual blocks with adaptive instance normalization (AdaIN) [10] layers and several upsampling and convolution layers. During the inference stage, only $E_{c}$ and $S$ are used to obtain the segmentation results. The whole framework is implemented with PyTorch on an NVIDIA Tesla P40 GPU. We use a mini-batch size of 8 for training, and train the framework for 30,000 iterations. We use the SGD optimizer with an initial learning rate of $2.5\times 10^{-4}$ for $E_{c}$, and the Adam optimizer with an initial learning rate of $1.0\times 10^{-3}$ for $E_{p}$ and $G$. The alternating training scheme [9] is adopted to train the discriminator $D$ using the Adam optimizer with an initial learning rate of $1.0\times 10^{-4}$. The polynomial decay policy is adopted to adjust all learning rates. The hyper-parameters $\lambda_{1},\lambda_{2},\lambda_{3}$ and $\lambda_{4}$ in Eq. (7) are empirically set to 0.01, 1.0, 0.5, and 0.01, respectively, although we find in our experiments that the results are not very sensitive with respect to the exact values of these parameters. The upper-bound networks for M2M and C2C are implemented using the standard PSP101 backbone and trained with the same settings as $E_{c}$.

Quantitative and qualitative analyses. To validate the efficacy of our framework in addressing the domain drop problem, extensive experiments are conducted on the two datasets. The competing algorithms include several SOTA UDA methods  [3, 4, 5, 16, 18, 20] in both the computer vision and medical image fields. We reimplement all compared methods with the origin released codes and apply with the default configurations. Table 1 presents the average performance drops in Dice and F1 score for the four cardiac structures as well as the mean values across structures on the MMWHS dataset. Table 2 presents the average performance drops for four abdominal structures and the mean values across them on the MMAS dataset. As can be seen, our framework outperforms all competing methods by large margins for all structures on these datasets. In addition, Fig. 2 visualizes the forward and backward performance drops and corresponding gaps between the bidirectional drops on the MMWHS dataset. For most competing UDA methods, there exist considerable gaps between the performances of the forward and backward adaptations. In contrast, our method significantly narrows this gap, and achieves the lowest performance drops in both adaptation directions. To further show the effectiveness of our method, we visualize the segmentation results by ours and competing methods in Fig. 3 and Fig. 4 for qualitative comparisons. As we can see, the segmentations by our method are much closer to the ground truth, especially for the backward adaptation. To summarise, the comparative experiments indicate that our framework can effectively address the domain drop problem, which has been overlooked by other UDA methods. Accordingly, our proposed framework presents new SOTA results for BiUDA segmentation on both the MMWHS and MMAS datasets.

Ablation study. We conduct ablation studies with incremental evaluations to validate the effectiveness of our novel modules, including the content-pattern consistency loss, the label consistency loss, and the domain-aware pattern encoder. The results on the MMWHS dataset are shown in Table 3. As we can see, after adding the content-pattern consistency loss, the average performance drop in Dice is reduced to 25.73%. In addition, the introduction of the label consistency loss further reduces the average drop to 12.98%, benefiting from pseudo training pairs in the target domain. Lastly, with the domain controller, the unified domain-aware pattern encoder learns pattern information from different modalities simultaneously, further reducing the average drop to 8.49%.

In this section, we summarize existing SOTA UDA methods and discuss the differences of our proposed framework with them. The core idea of AdaptSegNet [20], BDL [16], and CLAN [18] is to align the source domain with the target domain in feature space. These approaches work effectively on data of similar patterns and contents (e.g., natural images), since it is easy to align features from domains closer to each other. However, for our multi-modal medical image data, the difference between modalities presents a sharp change. This sharp change makes the source domain far away from the target domain, and makes these feature-aligning methods fail to deliver a decent performance. In contrast, SIFA [4] was tailored for medical scenarios and proposed additional alignments in the image and annotation spaces. Nonetheless, it only works effectively in the easier UDA direction (i.e., MRI to CT on the MMWHS dataset) but yields restricted performance in the reverse, more difficult direction (i.e., CT to MRI on the MMWHS dataset), as shown in Fig. 2.

Making the model insensitive to different patterns for better UDA is a more elegant way. Specifically, we can explicitly make the model aware of the content and pattern of a given image, and then apply the pattern from the target domain to the content from the source domain for training the UDA model. For example, ACE [21] utilizes a VGG net to extract the pattern code from the target-domain image and then integrates it into the source-domain image for UDA training. In contrast, DISE [3] leverages DRPL to extract the content and pattern codes from input of two domains, and adopts the decompose-recompose training strategy for these codes to realize UDA. This method has been verified effective on the BiUDA problem thus has a great potential in medical imaging UDA application. Different from DISE with separated pattern encoders, we propose a unified encoder that helps the model better understand the pattern difference. As shown in Table 3, with the unified pattern encoder, the framework achieves about 4% reduction in Dice drop. In addition, under the supervision of the CPC and LC losses, the pattern codes extracted by our method are more effective than those extracted from the VGG net. For this reason, the proposed method outperforms DISE and ACE, as shown in Table 1 and Table 2.

This work presented a novel, DRPL-based BiUDA framework to address the domain drop problem. The domain-aware pattern encoder was proposed for obtaining representative pattern codes from both the source and target domains and meanwhile simplifying the network complexity. To minimize potential data distortion in the process of domain adaptation, the content-pattern consistency loss was devised. In addition, the label consistency loss was proposed to achieve higher UDA segmentation performance. The comparative experiments indicates that our framework achieves new SOTA results for the challenging BiUDA segmentation tasks on both MMWHS and MMAS datasets.