title: Random Style Transfer based Domain Generalization Networks Integrating Shape and Spatial Information
authors: Lei Li Veronika A. Zimmer Wangbin Ding Fuping Wu Liqin Huang Julia A. Schnabel Xiahai Zhuang
text: 
  
   
  STDGNs with shape and spatial information
  
  
  
  L. Li et al.
  
  
  
  School of Biomedical Engineering, Shanghai Jiao Tong University, Shanghai, China School of Data Science, Fudan University, Shanghai, China 
  
  zxh@fudan.edu.cn School of Biomedical Engineering and Imaging Sciences, King’s College London, London, UKCollege of Physics and Information Engineering, Fuzhou University, Fuzhou, ChinaSchool of Management, Fudan University, Shanghai, China
  
  
  Random Style Transfer based Domain Generalization Networks Integrating Shape and Spatial Information
      Lei Li1, 2, 3 Veronika A. Zimmer 3 
  Wangbin Ding  4 
  Fuping Wu 2, 5 
  Liqin Huang 4 
  Julia A. Schnabel 3 
  Xiahai Zhuang2 ^() 
  
  
      August, 2020
  ==================================================================================================================================
  
  
  
  
  
  Deep learning (DL)-based models have demonstrated good performance in medical image segmentation.
  However, the models trained on a known dataset often fail when performed on an unseen dataset collected from different centers, vendors and disease populations.
  
  In this work, we present a random style transfer network to tackle the domain generalization problem for multi-vendor and center cardiac image segmentation.
  
  
  Style transfer is used to generate training data with a wider distribution/ heterogeneity, namely domain augmentation.
  As the target domain could be unknown, we randomly generate a modality vector for the target modality in the style transfer stage, to simulate the domain shift for unknown domains.
  
  The model can be trained in a semi-supervised manner by simultaneously optimizing a supervised segmentation and a unsupervised style translation objective.
  
  Besides, the framework incorporates the spatial information and shape prior of the target by introducing two regularization terms.
  
  We evaluated the proposed framework on 40 subjects from the M&Ms challenge2020, and obtained promising performance in the segmentation for data from unknown vendors and centers.
  
  
  
  
  
  
  § INTRODUCTION
  
  
  
  Quantification of volumetric changes during the cardiac cycle is essential in the diagnosis and therapy of cardiac diseases, such as dilated and hypertrophic cardiomyopathy.
  Cine magnetic resonance image (MRI) can capture cardiac motions and presents clear anatomical boundaries.
  For quantification, accurate segmentation of the ventricular cavities and myocardium from cine MRI is generally required <cit.>.
  
  
  Recently, deep learning (DL) based methods have obtained promising results in cardiac image segmentation.
  However, the generalization capability of the DL-based models is limited, i.e., the performance of a trained model will be drastically degraded for unseen data.
  It is mainly due to the existence of the domain shift or distribution shift, which is common among the data collected from different centers and vendors.
  
  To tackle this problem, domain adaptation and domain generalization are two feasible schemes in the literature.
  Domain adaptation algorithms normally assume that the target domain is accessible in the training stage.
  The models will be trained using labeled source data and a few labeled or unlabeled target data, and learn to align the source and target data onto a domain-invariant feature space. 
  For example, Yan et al. adopted an unpaired generative adversarial network (GAN) for vendor adaptation without using manual annotation of target data <cit.>.
  Zhu et al. proposed a boundary-weight domain adaptive network for prostate MRI segmentation <cit.>.
  
  Domain generalization, on the other hand, is more challenging due to the absence of target information.
  It aims to learn a model that can be trained using multi-domain source data and then directly generalized to an unseen domain without retraining.
  
  In clinical scenarios, it is impractical to retrain a model each time for new domain data collected from new vendors or centers.
  Therefore, improving the generalization ability of trained models would be of great practical value.
  
  
  
  
  
  Existing domain generalization methods can be summarized into three categories, 
  1) domain-invariant feature learning approaches, which focus on extracting task-specific but domain-invariant features <cit.>; 
  2) model-agnostic meta-learning algorithms, which optimize on the meta-train and meta-test domain split from the available source domain <cit.>; 
  3) data augmentation methods, which enlarge the scope of distribution learned by models for model generalization <cit.>. 
  For example, Volpi et al. utilized adversarial data augmentation to iteratively augment the dataset with examples from a fictitious target domain <cit.>.
  Zhang et al. propose a deep stacked data augmentation (BigAug) approach to simulate the domain shift for domain generalization <cit.>.
  Chen et al. employed some data normalization and augmentation for cross-scanner and cross-site cardiac MR image segmentation <cit.>.
  However, data augmentation in the source domain is not enough to mitigate the generalization gap, as the domain shift is not accidental but systematic <cit.>.
  
  
  In this work, we propose novel random style transfer based domain generalization networks (STDGNs) incorporating spatial and shape information.
  The framework focuses on improving the networks’ ability to deal with the modality-level difference instead of structure-level variations, i.e., domain augmentation.
  Specifically, we first transfer the source domain images into random styles via the style transfer network <cit.>.
  The generated new domain and known source domain images are then employed to train the segmentation network.
  We combine the two steps by connecting a style transfer and style recover network, to constrain the two domains with the same structure but different styles.
  With the domain augmentation strategy, the segmentation network is supposed to be domain-invariant and therefore can be adapted to unknown target domains.
  The framework can be trained in a semi-supervised manner, as it consists of an unsupervised style transfer task and a supervised segmentation task.
  Besides, we employ a shape reconstruction network (SRN) to include shape priors as a constrain, and utilize a spatial constrain network (SCN) to consider the large variety of 2D slices from various slice positions of 3D cardiac images <cit.>.
  
  
  
  
  
  § METHOD
   
  fig:method:framework provides an overview of the proposed framework.
  The proposed network is a cascaded modified U-Net consisting of two decoders for cardiac segmentation and style transfer, respectively.
  The generators consist of a shared encoder and two specific decoders, while the discriminator includes two parts, i.e., real/ fake discriminator and modality discriminator.
  Section <ref> presents the unified segmentation and style transfer framework based on GANs.
  To use unlabeled data, semi-supervised learning is employed, followed by a cross-task attention module to efficiently connect the segmentation and style transfer tasks (see Section <ref>).
  In Section <ref>, two losses based on the shape and spatial information are introduced as regularization terms.
  
  
  
   §.§ Deep Domain Generalization Framework
  
  Given N source domains 𝒟_s={𝒟_1, 𝒟_2, ⋯, 𝒟_N}, we aim to construct a model f(X),  
  
      f(X) → Y, X ∈𝒟_s ∪𝒟_t,
  
  where 𝒟_t are unknown target domains, and X,Y ∈ℝ^1 × H × W denote the image set and corresponding segmentation set.
  To achieve this, we propose to employ the style transfer to randomly generate new domains 𝒟_s^' with the same structure as the input source domain.
  Specifically, we assign a one-hot encoding modality vector v ∈[0,1]^N+1 for source domain image x ∈ℝ^1 × H × W, and then randomly generate a modality vector v^'∈[0,1]^N+1 for target domain image.
  The modality difference vector d_st=v^'-v, d_st∈[-1,0,1]^k is broadcasted from ℝ^k to ℝ^k × H × W and then concatenated with the input image. 
  
  As fig:method:framework shows, in the modality transfer phase, the generator G translates the source domain image x into a new target domain image x^' and predicts the segmentation ŷ, i.e., G(x, d_s t) →(x^', ŷ).
  The generated x^' is then inputted into a discriminator D=D_src, D_cls, where D_src recognizes it as real/ fake and D_cls predicts its modality.
  In the modality recovery phase, d_ts=v-v^' is concatenated with x^' as the input of G, to obtain the reconstructed original source image x^'' and its label ŷ^' for structure preservation. 
  Therefore, the adversarial loss of the GAN is adopted as follows,
  
      ℒ_adv^D=𝔼_x[D_src(x)]-𝔼_x^'[D_src(x^')]-λ_gp𝔼_x̂[(∇_x̂ D_src(x̂)_2-1)^2],
  
  
      ℒ_adv^G=𝔼_x^'[D_src(x^')],
  
  where x̂ is sampled from the real and fake images and λ_gp is the balancing parameter of gradient penalty.
  For style transfer, D_cls learns to classify x into its source modality v based on ℒ_cls^r, while G tries to transfer source modality into target modality v^' via ℒ_cls^f with
  
      ℒ_cls^r=𝔼_x, v[-log D_cls(v | x)],
  
  
      ℒ_cls^f=𝔼_x^', v^'[-log D_cls(v^'| x^')].
  
  The cycle consistency losses include ℒ_rec^lab and ℒ_rec^img. ℒ_rec^img is defined as
  
      ℒ_rec^img=𝔼_x, x^''[x-x^''_1],
  
  where x^'' is the predicted source domain image. ℒ_rec^lab is formulated in Section <ref>. 
  Hence, the final overall loss to optimize the generator and discriminator are defined as,
  
      min _Dℒ_D=-ℒ_adv^D+λ_clsℒ_cls^r,
  
  
      min _Gℒ_G=-ℒ_adv^G+λ_clsℒ_cls^f+λ_rec^imgℒ_rec^img+λ_segℒ_seg+λ_rec^labℒ_rec^lab,
  
  where all λ are balancing parameters, and ℒ_seg and ℒ_rec^lab are introduced in Section <ref>.
  
  
  
   §.§ Semi-supervised Learning with a Cross-task Attention
    
  To employ unlabeled data, we utilize a semi-supervised method via multi-task attention.
  Specifically, the proposed two-task networks include a supervised segmentation task and an unsupervised style transfer task, and the two tasks are connected via an attention mechanism.
  The style transfer decoder is trained without requirements to manually labeled data, and therefore the model can learn more discriminating features for segmentation from unlabeled images. 
  
  The assumption is that combining the features from the style transfer task can be beneficial to encode more representative and robust structural information for the segmentation task.
  However, due to the diverse feature spaces of the two tasks, a feature recalibration module is required to discard unrelated information.
  As fig:method:module (a) shows, we employ a cross-task skip connection for the last layers of the two decoders based on the squeeze and excitation (SE) module <cit.>. 
  The recalibration feature is defined as follows,
  
      F̃_seg=F_seg+F_tsl⊗σ(W_2δ(W_1 z)),
  
  where F_seg and F_tsl are the feature maps before the final output layers of the two tasks, W_1 and W_2 are fully connected layers, and z= AvgPool (F_tsl) ∈ℝ^C × 1 × 1.
  Specifically, W_1 is employed to squeeze the spatial information into a channel vector, i.e., W_1 z ∈ℝ^r × 1 × 1, while W_2 is used to recover the channel number from r to C for attention.
  
  
  
  
  
  
   §.§ Shape Reconstruction and Spatial Constraint for Segmentation
   
  To employ the shape prior, we use the SRN to reconstruct the label image to learn an intermediate representation, see fig:method:module (b).
  Specifically, the SRN model is pre-trained using available label images and corresponding predicted label images from basic U-Net <cit.>.
  It can be regarded as a constraint to regularize the predicted segmentation into a desired realistic shape. 
  As fig:method:framework shows, SRN is connected with the output of Decoder^seg to refine the shape of the prediction.
  The loss function of SRN is defined as,
  
      ℒ_SR=𝔼_y, ŷ, ŷ^'[R-(R̂ or R̂^') _2],
  
  where R denotes the reconstructed gold standard label from y, while R̂ and R̂^' indicate the reconstructed predicted label from ŷ and ŷ^', respectively.
  
  Considering the appearance of basal and apical slices can differ significantly, we utilize the SCN to predict the spatial information of each slice, i.e, its normalized distance to the central slice.
  As fig:method:module (b) shows, the SCN is connected to the bottom of U-Net via fully connected (FC) layers to predict the slice position.
  The SC loss is formulated as,
  
      ℒ_SC=𝔼_x, x^'[P-(P̂ or P̂^') _2],
  
  where P is the ground truth position, while P̂ and P̂^' are the predicted positions of x and x^', respectively.
  By incorporating the SR and SC loss, the segmentation loss and the label reconstruction loss both can be defined as,
  
      ℒ_seg=L_rec^lab=ℒ_CR + λ_SRℒ_SR + λ_SCℒ_SC,
  
  where ℒ_CR is the cross-entropy loss, and λ_SR and λ_SC are the balancing parameters.
  
  
  
  
  
  § EXPERIMENTS
  
  
  
  
   §.§ Materials
  
  
  
    §.§.§ Data Acquisition and Pre-processing.
  
  The data is from the M&Ms2020 challenge <cit.>, which provides 375 short-axis cardiac MRIs from patients with hypertrophic and dilated cardiomyopathies as well as healthy subjects. 
  All subjects were collected from three different countries (Spain, Germany and Canada) using four different MRI vendors (Siemens, General Electric, Philips and Canon).
  As tb:exp:dataset shows, the training data includes 150 labeled images from vendor A and B and 25 unlabeled images from vendor C. 
  The test data in this study contains 40 images from each vendor, i.e., A, B, C and a novel unseen vendor D. 
  
  All slices were cropped into a unified size of 144×144 centering at the heart region and were normalized using z-score.
  
  
  
    §.§.§ Gold Standard and Evaluation.
  
  The images have been manually segmented by experienced clinicians with labels of the left ventricle (LV), right ventricle (RV) blood pools, and left ventricular myocardium (Myo).
  Note that manual annotations are provided for only two cardiac time frames, i.e., end-diastole (ED) and end-systole (ES).
  These manual segmentations were considered as the gold standard.
  For segmentation evaluation, Dice volume overlap and Hausdorff distance (HD) were applied.
  
  
  
  
    §.§.§ Implementation.
  
  The framework was implemented in PyTorch, running on a computer with 1.90 GHz Intel(R) Xeon(R) E5-2620 CPU and an NVIDIA TITAN X GPU. 
  We used the SGD optimizer to update the network parameters (weight decay=0.0001, momentum=0.9). 
  The initial learning rate was set to 3e-4 and divided by 10 every 5000 iterations, and batch size was set to 20. 
  Note that the generator and discriminator are alternatively optimized every 50 epochs.
  The balancing parameters were set as follows, λ_seg=100, λ_rec^img=100, λ_SC=1, λ_SR=100, λ_gp=10, λ_cls=10, and  λ_rec^lab was initial set to 0 and then gradually increase to 100.
  
  
  
  
  
  
   §.§ Result
  
  
  
  
  
  
    §.§.§ Comparisons with Literature and Ablation Study.
  
  tb:result:Dice presents the quantitative results of different methods for cardiac segmentation in terms of Dice.
  Two algorithms were used as baselines for comparisons, i.e., fully supervised U-Net <cit.> and multi-task attention-based semi-supervised learning (MASSL) network <cit.>.
  MASSL is a two-task network consisting of two decoders, which are used for segmentation prediction and image reconstruction, respectively.
  Compared to U-Net, MASSL employed the unlabeled data for training, and therefore obtained better mean Dice though with worse Dice on RV.
  STDGNs is also semi-supervised and it fully utilized the supervised and unsupervised task relationship via a cross-task attention module.
  It obtained evidently better Dice scores in all sub-regions compared to U-Net and MASSL.
  
  
  To explore the effectiveness of the proposed shape and spatial constrain, we also adopted an ablation study.
  One can see that STDGNs-SR outperformed STDGNs in terms of Dice, and incorporating the SC module further improved the performance on LV and RV segmentation. 
  Furthermore, one can see that the performance on the RV segmentation were all worse than that on the LV and Myo for these algorithms.
  It may due to the fact that the shape of the RV is generally more irregular than that of the LV and Myo.
  
  
  
  
  
  
  
    §.§.§ Performance on the Data from Different Vendors
  
  To investigate the generalization ability of our model, we compared the performances on data from various vendors and centers.
  fig:result:boxplot (a) and (b) illustrate the Dice and HD of the three substructures on each vendor.
  It is evident that the model performed generally well on the known domains, i.e., vendor A, B and C.
  Though the label images of vendor C are unavailable, our model still obtained promising results, thanks to the semi-supervised learning.
  However, for the unknown domain vendor D, the performance drastically dropped, especially for the RV segmentation.
  To verify the generalization performance of the proposed model, we further compared the results of our model with that of U-Net on vendor D.
  As tb:result:Vendor_D shows, the proposed model performed evidently better in generalization compared to the baseline, though there is still a huge scope for improvements.
  
  
  
  
    §.§.§ Performance on the Data from Different Centers.
  
  Besides the generalization on different vendors, we also explored the effectiveness of the proposed model on different centers.
  As vendors C and D include one single center, we only present the results of two centers of vendor A and B, respectively (see fig:result:boxplot (c) and (d)).
  One can see that the results of center A6 were generally worse than that of center A1, as the data of center A6 is not included in the training dataset.
  However, compared to the performance decline on the new vendor D, the performance decrease on the new center is evidently less.
  For the centers already included in the training dataset (centers B2 and B3), the model obtained promising results in terms of both Dice and HD.
  Note that the performance on center B3 is evidently better than that of center B2, though the training data includes more center B2 data (see tb:exp:dataset).
  This may be because the center B2 data is more challenging to segment or center B3 has more similar appearance with other centers in the training data.
  
  
  
  
  
    §.§.§ Performance on the ED and ES Phase.
  
  tb:result:ED_ES presents the quantitative results of the proposed method for cardiac segmentation at the ED and ES phase.
  One can see that the performance was substantially higher in ED than in ES for all substructures in terms of Dice, but was similar in HD.
  This might be caused by the larger ventricular shape in the ED phase than the ES phase.
   
  
  
  
  
  § CONCLUSION
  
  In this work, we have proposed an end-to-end learning framework for domain generalization via random style transfer.
  The proposed algorithm has been applied to 40 cine images and obtained promising results on the unknown domain images from new vendors and centers.
  Especially, the model could generalize better on new centers than new vendors.
  The results have demonstrated the effectiveness of the proposed spatial and shape constrain schemes, which encourage the network to learn more domain-invariant features.
  Our technique can be easily extended to other segmentation tasks on data from new vendors and centers.
  A limitation of this work is that the definition of the modality vector requires the number of domains, which is normally unknown.
  Besides, the style transfer only can modify modality globally, but sometimes the target modality can be complex with high heterogeneity.
  In future work, we will replace the modality vector by other more flexible modality representations, and develop a style transfer module that can focus on the local difference among different domains.
  
  
  
    §.§.§ Acknowledgement.
  
  This work was supported by the National Natural Science Foundation of China (61971142), and L. Li was partially supported by the CSC Scholarship. 
  JA Schnabel and VA Zimmer would like to acknowledge funding from a Wellcome Trust IEH Award (WT 102431), an EPSRC program Grant (EP/P001009/1), and the Wellcome/EPSRC Center for Medical Engineering (WT 203148/Z/16/Z).
  
  
  splncs04
  
  
  