Unsupervised Domain Adaptation with Dual-Scheme Fusion Network for Medical Image Segmentation

1 Introduction

Deep learning has achieved state-of-the-art results in many fields. However, the model trained on one dataset cannot be directly applied to other datasets with similar contents and various distribution. Since the different distribution brings a domain shift problem and further leads to identification failure. Solving the identification failure caused by domain shift is an important issue, which can improve the generalization as well as portability and further solve the problem of lacking of labeled data and expensive cost of manual annotation. Unsupervised domain adaptation methods are an effective method to solve those problems, which enables one labeled dataset to serve multiple other unlabeled ones and greatly improve the practical value of deep learning network. The unsupervised domain adaptation methods can be categorized into two types: image adaptation and feature adaptation. [Chen et al., 2018; Bousmalis et al., 2017; Hoffman et al., 2017; Chen et al., 2019] are the representative articles of image adaptation. In [Chen et al., 2018; Bousmalis et al., 2017], segmentation network is first trained using source domain dataset, and then target data are translated into source domain for segmentation using the trained segmentation network. In [Hoffman et al., 2017; Chen et al., 2019], the source domain data are first translated into target domain, and then the segmentation network are trained by the generated target-like data with corresponding source domain annotation. The representative articles of feature adaptation are [Tsai et al., 2018; Dou et al., 2018]. Those methods directly perform feature adaptation without cross-domain translation of the image.

∗Corresponding Authors.

However, current works based on image adaptation only conduct the two-direction translation separately. This imbalanced source-to-target or target-to-source one way connection may not eliminate the domain gap, which limits the performance of the pre-trained model. To overcome this problem, we propose to combine these two complementary directions into one unified framework for acquiring a better result. Our proposed model, named Dual-Scheme Fusion Network (DSFN), adopts CycleGAN[Zhu et al., 2017]'s architecture as base structure. To make the network can simultaneously conduct the tasks of image translation and segmentation, we modified the generator to a joint translation-segmentation network. The translation and segmentation parts share a same encoder. This setting can not only reduce network's parameters, but also enable the two tasks to enhance each other. Moreover, the segmentation network also can help reserve content of the generated image. During training, the first joint cross-domain translation and segmentation network is trained by source domain images. The task of translation branch is generating target-like images and the segmentation branch aims to segment the source domain image. The first trained segmentation network can be used to segment the source-like image generated by the target domain image. Then, the generated target-like image is used to train the joint network for converting target-like image to source domain and segmen-

s

tation. The second obtained segmentation branch can be directly used to segment the target domain data. Finally, we merge the two segmentation outputs of target image to obtain a more accurate result. Moreover, we also introduce feature adaptation loss referring[Chen et al., 2019] to further narrow the feature-level gap between target and target-like data, as well as source and source-like data. The whole network is designed subtly, the backward process of our network is also the test procedure. The main contributions of our paper are as follows:

1. We propose an end-to-end unsupervised parallel dualscheme fusion network for solving the problem of unsupervised domain adaptation from two perspectives by building both source-to-target and target-to-source connections, this balanced joint information flow helps bridge the domain gap to further improve the network performance.
2. We train our network in a cycle manner and the backward process is also the testing process. The structure of our network is effective and elegant.
3. We validate our method on the BRATS[Menze et al., 2014] and MM-WHS[Zhuang and Shen, 2016] dataset. Experimental results show that our proposed approach exceed current state-of-the-art domain adaptation methods, effectively alleviating the expensive cost of manual label-

ing cross-domain medical data, which has great clinical practical value.

2 RelatedWork

Typically, domain adaptation methods generalize to other tasks by minimizing the distance of feature space between the source and the target domain. For example, Tzeng et al. [Tzeng et al., 2014] tries to learn a domain-invariant feature representation between source and target by proposing an adaptation layer and a domain confusing loss. Long et al. [Long et al., 2015] matches source and target distribution by embedding their hidden features in a reproducing kernel Hilbert space.

Inspired by the advanced of generative adversarial network [Goodfellow et al., 2014], adversarial learning becomes the main training strategy for domain adaptation methods. Some of them directly map feature representations between source and target domain, for example, Tsai et al .[Tsai et al., 2018] conducted different feature levels' domain adaptation by introducing multi-level discriminators to the segmentation network and Dou et al. [Dou et al., 2018] shared a similar idea with an innovative network structure. Some conduct pixellevel domain transfer, which has two directions as mentioned above. One is to perform target-to-sourced translation, such as [Chen et al., 2018], and the other is to perform source-totarget translation, such as [Hoffman et al., 2017] and [Zhu et al., 2020].

Due to the expensive cost of manual annotation, there is a lot of research on domain adaptation in the medical field. For example, Nie et al. [Nie et al., 2017] proposed a context-aware GAN to generate PET images from MRI images. Wolterink et al. [Wolterink et al., 2017] used CycleGAN to convert between unpaired CT and MRI images. Similarly, Chen et al. [Chen et al., 2019] conducted the MRI to CT translation and CT splenomegalia segmentation at the same time. Zhao et al. [Zhao et al., 2018] proposed a RsGAN technique to synthesize fundus image data set and generalize its segmentation network to other data sets. Chen et al. [Chen et al., 2018] adopted target-to-source translation to segment left/right lung using chest X-ray datasets. Kamnitsas et al. [Kamnitsas et al., 2017] proposed a multi-connected domain discriminator to segment unannotated MR images with traumatic brain injuries. [Dou et al., 2018] transfers a cardiac segmentation network trained with MRI to unannotated CT images. Fang and Yan [Fang and Yan, 2020] propose a new multi-scale paramid network and loss function to segment multiple partially labeled datasets.

3 Methods

The overall framework of our proposed Dual-Scheme Fusion Network for unsupervised domain adaptation is shown in Fig.1, which contains two joint translation-segmentation modules. Each joint network has an encoder-decoder structure. The decoder is divided into two branches-translation branch and segmentation branch, which share one same encoder. Each translation branch corresponds to a discriminator to distinguish the generated image from the real image. The design of our network structure is a transformation based on CycleGAN[Zhu et al., 2017]. The difference is that we replace the generator with a joint translation-segmentation module.

The background setting is that we have a source domain dataset {xs} with corresponding annotation {ys} and want to annotate the unlabeled target domain dataset {xt}. We solve this problem from two perspectives: (a) Target-to-Source: We train a segmentation network Es + Ss using source domain dataset {xs, ys}, then translate target image xt into the source domain through Et + Tt→s and segment it through the segmentation network Es + Ss; (b) Source-to-Target: we first translate the source domain image xs into the target domain through Es + Ts→t, then a segmentation network Et + St is trained using the generated xs→t with corresponding annotation ys, so the real xt can be directly segmented by Et + St. We integrate these two streams into a unified network and fuse their results to form a new prediction as final output.

The flow of data through our network is described as follows. In the training process of forward, xs is inputed into the joint network G1 to obtain the cross-domain generated image xs→t and the segmentation prediction ˆ ys. Discriminator Dt will distinguish between xt and xs→t. The generated image xs→t is then putted into the joint network G2 to obtain reconstructed image xs→t→s and the segmentation prediction ˆ ys→t .

And in the training process of backward (test process follows the same data flow of backward process), the target domain image xt is inputted into joint network G2, and we can get the xs-like cross-modality image xt→s and segmentation prediction ˆ yt. Discriminator Ds distinguish between xs and xt→s. Then, the generated xt→s is putted into the joint network G1 which outputs the reconstructed image xt→s→t and the segmentation prediction ˆ yt→s. As for feature adaptation, Discriminator Ds f distinguishes between ˆ ys and ˆ yt→s while Dtf will distinguish between ˆ ys→t and ˆ yt. Dt and Ds also have an auxiliary task to respectively distinguish between xs→t and xt→s→t, xs→t→s and xt→s. In the test process, ˆ yt and ˆ yt→s are two segmentation predictions for the target domain image xt from different perspective. We fuse the two predictions to get a higher-accuracy result fusion(ˆ yt, ˆ yt→s) as our final result for target image xt. For the selection of the fusion strategy, we find that averaging the prediction probabilities of these two results can obtain better performance.

3.1 Adversarial Loss

We train our unified network in a cycle manner, which means that we have a source to target translation and a target to source translation as well. To generate more realistic images, we adopt the adversarial loss as loss function. The adversarial learning consists of two modules: a generator and a discriminator. Our problem is an image-to-image domain translation problem. So in source to target translation, we forward a source image xs to generator Es+Ts→t and it's responsible for generating an image xs→t that looks like target image and can fool the discriminator Dt, while the discriminator Dt aims to distinguish whether the image is actually derived from target domain or synthesized by the generator. This forms a process of adversarial learning, which can be described as minEs+Ts→t maxDt Ladv(Es, Ts→t, Dt) and:

The target to source translation is likewise. The objective function is minEt+Tt→s maxDs Ladv(Et, Tt→s, Ds) and:

3.2 Reconstruction Loss

In order to make the generated image preserve structure, shape and other content information of the original image, we introduce reconstruction loss in our model to constraint the training process.

Taking the forward process as example, reconstruction means that xs→t translated by xs can be reconstructed back to source domain by generator Et+Tt→s and will be consistent with the original xs, ie Tt→s (Et (xs→t)) ∼ xs. Its objective function is:

Similarly, in the backward process, it is Ts→t (Es (xt→s)) ∼ xt:

3.3 Segmentation Loss

As mentioned above, our network contains two segmentation networks to solve the unsupervised domain adaptation problem. As shown in Fig. 1, in the forward process, Es + Ss is our first segmentation network by training the source domain image xs, Et + St is the second segmentation network obtained by training the generated image xs→t. The loss function of two segmentation network is:

As for backward process, segmentation network is employed in the test procedure to get the final segmentation result of the target domain image xt. First, feeding target domain data xt into the joint network G2 through the segmentation branch Et + St in the joint network for getting the first segmentation prediction ˆ yt. At the same time, xt→s is generated by the translation branch Et+Tt→s. Then xt→s will be fed into the joint network G1, so the segmentation branch Es+Ss of G1 can output the second segmentation prediction for xt, denoted as ˆ yt→s. Finally, we fuse ˆ yt and ˆ yt→s to get the final segmentation result fusion{ˆ yt, ˆ yt→s}.

3.4 Feature-adaptation Loss

In order to further improve the segmentation results, we introduce feature adaptation loss referring to [Chen et al., 2019]. As mentioned above, the segmentation network Es + Ss is trained on xs and used to induce xt→s. But there is still domain gap between them that makes the segmentation result of xt→s inaccurate. To reduce their differences in feature level, we introduce a discriminator Ds f and an auxiliary task for Dt, when Ds f cannot distinguish between segmentation predictions from xs or xt→s, and Dt can't differentiate between xs→t and xt→s→t, the features of xs and xt→s are consistent. The corresponding loss function is:

and

Similarly, to reduce the domain gap between xt and xs→t, we also introduce a discriminator Dtf and an auxiliary task for Ds. And the loss function is:

and

In summary, the overall objective function is:

where we set λrec = 10, λf = 0.1, λaux = 0.1 and λseg = 1 to balance different losses.

3.5 Network Structure and Training Details

Our translation-segmentation joint network follows the structure settings from [Chen et al., 2019]. Overall, our joint network contains an encoder connecting to two decoder branches. The encoder contains three residual blocks[He et al., 2016] and down-sampling layers in sequence, then 8 residual blocks are followed and two dilated residual blocks continue to enlarge the receptive field, and two convolutional layers are connected at the end. For the translation branch, decoder contains a convolutional layer, 4 residual blocks and 3 deconvolutional layers to restore the resolution, then a convolutional layer and tanh activation function is followed. As for the segmentation branch, it contains a convolutional layer and an upsampling layer.

For the discriminator, we follow the setting of patchGAN[Isola et al., 2017] and use five convolutional layers to convolve in turn, with the first three stride being 2 and the last two being 1. Their kernel size is 4×4 and the channel number is 64, 128, 256, 512 and 1, respectively.

During the training process, we substitute the least-squares loss for the log likelihood objective in the adversarial loss referring to CycleGAN[Zhu et al., 2017], which can make the training process more stable[Mao et al., 2017]. We use the Adam solver[Kingma and Ba, 2014] with a batch size of 8 and initialize the learning rate to 0.0002 for translation task and 0.001 for segmentation task. The various parts of the network are trained in the following order: Es+Ss → Ts→t → Dt → Ds f →Et+St →Tt→s →Ds →Dtf .

4 Experiments

We validate our proposed method on two datasets - MultiModality Brain Tumor Segmentation Challenge 2018[Menze et al., 2014] and the Multi-Modality Whole Heart Segmentation Challenge 2017[Zhuang and Shen, 2016]. The first dataset contains four modalities of MRI imaging: T1, T1CE, T2, and FLAIR. There are a total of 75 patient data. Because experts always annotate whole tumor on T2 modality, we take T2 as the source domain with whole tumor annotation and aims to segment three target domain - T1, T1CE, FLAIR. Our experiments are conducted in an unpaired manner.

The second dataset is composed of 20 unpaired MR and CT imaging data, and their ground truth contains the annotation of four heart structures - the ascending aorta (AA), the left atrium blood cavity (LAC), the left ventricle blood cavity (LVC), and the myocardium of the left ventricle (MYO). We use MR images and their corresponding labels as the source domain dataset to segment the target domain CT images.

For both datasets, we randomly select 80% patient data for each modality as the training set and 20% as the testing set. All data subtract their mean and are divided by their standard deviation for normalization, then we switch them to [-1,1] before feeding to the network. We augment the data by rotation, cropping, etc., and each slice is resized to 256×256. The target labels of these two datasets are only used for evaluation.

Our experiments are conducted from two aspects: comparison with other state-of-the-art unsupervised domain adaptation methods and ablation experiments. For evaluation, we select two metrics commonly used in medical image segmentation tasks: Dice coefficient[%] and Hausdorff Distance. Dice coefficient measures the overlapping part of our prediction results and the groundtruth. Hausdorff distance is a distance defined between two sets in the metric space.

4.1 Comparison Experiments

The methods we compare with are: CycleGAN[Zhu et al., 2017], CyCADA [Hoffman et al., 2017], AdaptSegNet[Tsai et al., 2018], SIFA[Chen et al., 2019]. In addition, we also compare the results without domain adaptation (directly input the target image into the trained source domain segmentation network).

Table 1 and Table 2 respectively show the comparison result between our method and others on the brain tumor segmentation dataset and cardiac segmentation dataset. It can be seen that our method has a significant improvement compared to no domain adaptation result by improving Dice Score from 22.6% to 66.1% on tumor segmentation task and from 17.2% to 75.8% on cardiac structures segmentation task, and reducing the Hausdorff Distance from 46.8mm to 15.6mm and from 44.8mm to 11.4mm respectively, which shows that our method is very effective in solving domain degradation problem. More importantly, our proposed method achieves better results than the current state-of-the-art methods on these two datasets. Especially when compared with SIFA[Chen et al., 2019], which only consider the source-to-target direction, we can increase the Dice by 6.8% on the BRATS dataset and by 2.8% on the MM-WHS datasetet after adopting both source-to-target and target-to-source direction and fusing their results, demonstrating the effectiveness of our proposed method.

Fig.2 and Fig.3 present the visual results compared with other methods. We can note that our outputs are not only closer to the ground truth, but also the wrong semantic prediction results are far less than others, indicating that our fusion strategy can indeed make the results more robust.

4.2 Ablation Experiments

We perform ablation experiments on feature adaptation and fusion to verify the influence of each component for our proposed Dual-Scheme Fusion Network.

The ablation experiments are set as follows, we conduct experiments respectively with or without feature adaptation, in each case, we compare the results of three schemes: sourceto-target scheme, target-to-source scheme, and their fusion solution. Table 3 and Table 4 present the ablation study results, Result1 represents the target-to-source method, Result2 represents the source-to-target method and Fusion represents our Dual-Scheme Fusion method.

It can be seen that, in both datasets, the results of adding feature adaptation are improved in all three cases compared with no feature adaptation. Among them, the improvement on the MM-WHS dataset is particularly obvious, and the Dice score can increase over 10% compared to no feature adaptation, demonstrating the effectiveness of feature adaptation.

And regardless of whether feature adaptation is added or not, compared with experiments conducted from source-totarget or target-to-source perspective alone, the experimental results are more accurate after using our dual-scheme fusion strategy on both datasets. And Dice score can improve about 3% on BRATS dataset and 4% on MM-WHS dataset without feature adaptation. If feature adaptation is added, our method can further improve the Dice score by about 3% on BRATS and 2% on MM-WHS, which reflects the effectiveness of fusion strategy.

Fig.4(a) and Fig.4(b) present the visual results on BRATS dataset and MM-WHS dataset respectively. The first row shows the results without feature adaptation while the second row shows them with feature adaptation. It can be seen that overall the results of adding feature adaptation are more accurate than not adding. The first column is the input images, prediction1 is the results from target-to-source scheme, prediction2 is from source-to-target scheme and fusion result means adopting our fusion strategy. We can see that the edges of fusion images are smoother than the two separate prediction results and are closer to the ground truth. Moreover, from Fig.4(b) we can see that the misprediction outside the cardiac structures from the two prediction results also disappear after fusion strategy, which means through our fusion strategy, the results from target-to-source and source-to-target perspectives can complement each other and a smoother and more robust result can be obtained.

5 Conclusion

In this paper, we address the problem of domain adaptation from two complementary perspectives and unify these two approaches into an end-to-end network in an elegant and efficient way. Through the cycle structure, not only can the unpaird data training problem be solved, but also a clever combination of these two perspectives is obtained. Finally, we choose an effective fusion method to acquire the final result. For further improvement, we also add feature adaptation loss. The experimental results show that our proposed method greatly surpasses the baseline and is superior to other state-of-the-art methods.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grants 62041105, 6182211, the Science and Technology Major Project of Hubei Province (Next-Generation AI Technologies) under Grant 2019AEA170, Natural Science Foundation of Hubei Province under Grants 2018CFA050. The numerical calculations in this paper have been done on the supercomputing system in the Supercomputing Center of Wuhan University.