CXR Segmentation by AdaIN-based Domain Adaptation and Knowledge Distillation

High-accuracy image segmentation often serves as the first step in various medical image analysis tasks [6, 20]. Recently, deep learning (DL) approaches have become the state-of-the-art (SOTA) techniques for medical image segmentation tasks thanks to their superior performance compared to the classical methods [6].

The performance of DL-based segmentation algorithm usually depends on large amount of labels, but segmentation masks are scarce due to expensive and time-consuming annotation procedures. Another difficulty in DL-based segmentation is the so-called domain shift, i.e., a segmentation network trained with data in a specific domain often undergoes drastic performance degradation when applied to unseen test domains.

For example, in the field of chest X-ray radiograph (CXR) analysis [2], segmentation networks trained with normal CXR data often produce under-segmentation when applied to abnormal CXRs with severe infectious diseases such as viral or bacterial pneumonia [18, 27]. The missed regions from under-segmentation mostly contain crucial features, such as pulmonary consolidations or ground-glass opacity, for classifying the infectious diseases. Thus, highly-accurate lung segmentation results without under-segmentation are required to guarantee that DL-based classification algorithms fully learn crucial lung features, while alleviating irrelevant factors outside the lung to prevent shortcut learning [7].

To solve the label scarcity and domain shift problems, there have been extensive researches to train segmentation networks in a semi-supervised manner using limited training labels, or in a self-supervised or unsupervised manner even without labeled dataset [1, 28, 29, 23, 32, 19, 17, 3]. However, these approaches appear different from each other, and there exist no consensus in regard to how these different approaches can be synergistically combined. Inspired by success of StarGANv2 for transferring style between various domains [4], here we propose a style transfer-based knowledge distillation framework that can synergistically combine supervised segmentation, domain adaptation and self-supervised knowledge distillation tasks to improve unsupervised segmentation performance. Specifically, our framework is designed to deal with difficult but often encountered situations where segmentation masks are only available for normal data, but the trained method should be applied to both normal and abnormal images.

The key idea is that a single generator trained along with the adaptive instance normalization (AdaIN) [11] can perform supervised segmentation as well as domain adaptation between normal and abnormal domains by simply changing the AdaIN codes, as shown in Fig. 1(a) and (b), respectively. The network is also trained in a self-supervised manner using another AdaIN code in order to force direct segmentation results (illustrated as a red arrow in Fig. 1(c)) to be matched to indirect segmentation results through domain adaptation and subsequent segmentation (illustrated as a blue arrow in Fig. 1(c)). Since a single generator is used for all these tasks, the network can synergistically learn common features from different dataset through knowledge distillation.

To validate the concept of the proposed framework, we train our network using labeled normal CXR dataset and unlabeled pneumonia CXR dataset, and test the model performance on unseen dataset composed of COVID-19 pneumonia CXRs [27, 13]. We further evaluate the network performance on domain-shifted CXR datasets for both normal and abnormal cases, and compared the results with other SOTA techniques. Experimental results confirm that our method has great promise in providing robust segmentation results for both in-domain and out-of-domain dataset. We release our source code ¹¹1https://github.com/yjoh12/CXR-Segmentation-by-AdaIN-based-Domain-Adaptation-and-Knowledge-Distillation.git for being utilized in various CXR analysis tasks presented in Supplementary Section S1.

To make this paper self-contained, here we review several existing works that are necessary to understand our method.

We present a novel framework, which can perform segmentation, domain adaptation and self-supervised learning tasks within a single generator in a cost-effective manner (Supplementary Section S9). The proposed network can fully leverage knowledge learned from each task by utilizing shared network parameters, thus the model performance can be synergistically improved via knowledge distillation between multiple tasks, so that achieves SOTA performance on the unsupervised abnormal CXR segmentation task. The experimental results demonstrate that the proposed unified framework can solve domain shift issues with great generalizability, even on dataset with way-shifted distribution. Last but not the least, the proposed model does not need any pre-processing techniques but shows superior domain adaptation performance, which presents a promising direction to solve the generalization problem of DL-based segmentation methods.

Current deep learning (DL)-based CXR analysis tasks not just diagnose disease but provide explainable results such like saliency map or quantified severity level to assist clinicians [16, 10, 21, 9]. Accordingly, we further investigated applications of the proposed method using COVID-19 pneumonia CXR dataset [13]. For generating saliency map, we referenced a public source code ²²2https://github.com/priyavrat-misra/xrays-and-gradcam.git. In addition, for severity quantification, we followed array-based methods [21, 30].

Fig. S1 shows that Proposed and Proposed+$\ell_{self}$ models provide the most suitable segmentation mask for saliency map or severity array generation. The results demonstrate that our methods can be utilized for various automatic CXR analysis tasks.

We provide details of the proposed framework, as shown in Table S1, S2, S3 and S4. Note that each input and output dimension for the domain adaption task is single channel ($C_{da}$), whereas the segmentation task requires two channels ($C_{seg}$) to separate foreground and background channels.

AdaIN has been proposed as an extension of the instance normalization [11]. AdaIN layer receives a content input $x$ and a AdaIN code $a$, and simply aligns the channel-wise mean and variance of $x$ to match those of desirable style by:

where $f$ and $g$ compute affine parameters from AdaIN code $a$, $\mu$ and $\sigma$ represent mean and variance, respectively. In this way, AdaIN simply scales the normalized content input with $\sigma(a)$, and shifts with $\mu(a)$.

For ablation study, we analyzed contribution of different losses for the supervised segmentation task. We cumulatively added each loss to the baseline model, and compared segmentation performance on abnormal CXR.

Fig. S2 compares performance of each configuration. In configuration (b) with additional segmentation loss, domain adaptation performance was superior to the baseline. However, we observed that the segmentation results have several concave regions (blue boxes), which failed to resemble the general shape of lung structure. In configuration (c) with additional style loss, we observed that the segmentation results resemble normal lung better than (b), thanks to the style loss, which can extract common features of normal and abnormal CXR via the shared layer of the style encoder. In configuration (d) with additional cycle-consistency loss, we observed rather degraded lung segmentation performance as depicted as blue boxes, which have more concave regions compared to that of (c). The cycle-consistency loss, which tries to revert the generated lung mask back to the original image, may disturb the network to segment extremely consolidated lung regions.

Based on the ablation study results, we set the configuration (c) as our supervised segmentation loss.

Domain adaptation loss basically follows StarGANv2 [4], given by

where $\lambda_{cycle}$,$\lambda_{style}$ and $\lambda_{div}$ are hyper-parameters and $\ell_{adv}$ is the adversarial loss defined by

where $\mathcal{S}$ and $\mathcal{T}$ are source and target domains, which are chosen randomly from $\mathcal{X}$ and $\mathcal{Y}$ so that all domain combinations can be considered. Furthermore, the learnable part of the AdaIN code $a_{da}^{\mathcal{T}}$ is generated either from the encoder AdaIN coder generator $F_{d}$ or the style encoder $S$ given a reference target $t\in\mathcal{T}$.

The cycle-consistency loss $\ell_{cycle}$ is defined as follows:

Similar to the cycle-consistency loss $\ell_{cycle}$ for the images, we introduce the style loss $\ell_{style}$ to enforce the cycle-consistency in the AdaIN code domain. More specifically, once a fake image is generated using a domain-specific AdaIN code, the style encoder with the fake image as an input should reproduce the original AdaIN code. This can be achieved by minimizing the following style loss:

Finally, to make the generated fake images diverse, the difference between two fake images that are generated by two different AdaIN codes should be maximized. This can be achieved by maximizing the following loss:

where an additional ${a^{\prime}}_{da}^{\mathcal{T}}$ is generated either from the encoder AdaIN coder generator $F_{d}$ or the style encoder $S$ given an additional reference image.

The proposed method was implemented with PyTorch library [22]. We applied Adam optimizer [15] to train the models and set the batch size 1. The model was trained using a NVIDIA GeForce GTX 1080 Ti GPU. Hyper parameters were chosen to be $\lambda_{cycle}=2$, $\lambda_{style}=1$, $\lambda_{div}=1$, $\lambda_{seg}=5$, $\lambda_{inter}=10$ and $\lambda_{intra}=1$. Learning rate was optimized to $0.0001$. Once training iteration reaches certain fixed iteration points throughout the total iterations, the learning rate was reduced by factor of 10.

The network was trained for 20K iterations to simultaneously train domain adaptation and supervised segmentation tasks. We adopted early stopping strategy based on validation performance of abnormal chest X-ray radiograph (CXR) segmentation results. In terms of the training sequence, the self-supervised training started after training the domain adaptation and supervised segmentation tasks until they guaranteed certain performances. For self-supervised learning, the network was continued to be trained in self-supervised manner for additional 5K iterations.

At the inference phase, as for post-processing steps, two largest contours were automatically selected based on contour areas, and any holes within each contour were filled. The post-processing technique was identically applied to all the comparative model outputs for fair comparison.

For comparative study, baseline models for domain adaptation and supervised segmentation tasks, i.e., CycleGAN [33], MUNIT [12], StarGANv2 [4] and U-Net [24], were trained with identical conditions to that of the proposed model. For comparing performance of the unified domain adaptation and segmentation network, we inferenced pre-trained networks optimized for abnormal CXR segmentation, i.e., XLSor [28] and lungVAE [25], by utilizing their official source codes. ³³3https://github.com/raghavian/lungVAE⁴⁴4https://github.com/rsummers11/CADLab/tree/master/Lung_Segmentation_XLSor

We analyzed typical error cases, which failed to be segmented, and the error cases were grouped into three categories: (a) Over-segmentation on background pixels, (b) distorted lung shape, and (c) distorged lung boundary, as shown in Fig. S3.

The proposed unified framework costs less training computation resources, compared to training individual domain adaptation and segmentation networks. Table S5 shows total network parameters utilized for either training or inference of comparative networks.

Once the model is trained, only the generator with pre-built AdaIN codes are used at the inference phase, thus the model costs only the single generator.

As shown in Table S5, Proposed and Proposed+$\ell_{self}$ models need the least number of network parameters, with most promising segmentation performance. Specifically, compared to Proposed model, Proposed+$\ell_{self}$ model only needs a single inference without preceding domain adaptation task, with comparable segmentation performance to that of Proposed model.