Unsupervised Medical Image Translation Using Cycle-MedGAN

TL;DR

Cycle-MedGAN is an unsupervised medical image translation framework that improves image quality by minimizing textural and perceptual discrepancies, enabling better PET-CT translation and MR motion correction without requiring paired datasets.

Contribution

It introduces non-adversarial cycle losses to enhance unsupervised medical image translation, addressing issues of blurriness and unrealistic details in translated images.

Findings

Outperforms existing unsupervised methods in PET-CT translation

Effective in MR motion correction

Produces sharper, more realistic images

Abstract

Image-to-image translation is a new field in computer vision with multiple potential applications in the medical domain. However, for supervised image translation frameworks, co-registered datasets, paired in a pixel-wise sense, are required. This is often difficult to acquire in realistic medical scenarios. On the other hand, unsupervised translation frameworks often result in blurred translated images with unrealistic details. In this work, we propose a new unsupervised translation framework which is titled Cycle-MedGAN. The proposed framework utilizes new non-adversarial cycle losses which direct the framework to minimize the textural and perceptual discrepancies in the translated images. Qualitative and quantitative comparisons against other unsupervised translation approaches demonstrate the performance of the proposed framework for PET-CT translation and MR motion correction.

Tables

Table 1. TABLE I: Quantitative comparison of unsupervised translation techniques

Model	(a) PET-CT translation
	SSIM	PSNR(dB)	VIF	LPIPS
UNIT	0.8485	20.14	0.2057	0.6762
Cycle-GAN	0.8963	23.35	0.3831	0.2561
Cycle-MedGAN	0.9115	24.08	0.4275	0.2233
Model	(b) MR motion correction
	SSIM	PSNR(dB)	MSE	UQI
UNIT	0.6914	18.64	0.1239	0.6953
Cycle-GAN	0.8011	22.39	0.3432	0.3282
Cycle-MedGAN	0.8118	22.96	0.3513	0.3029

Equations

$$\begin{array}{rl}\underset{G_1}{\min }\underset{D_1}{\max }{\mathcal{L}}_{\text{adv}}(G_1,D_1)=& {\mathbb{E}}_y\left[\text{log}{D}_1(y)\right]+\\ & {\mathbb{E}}_x\left[\text{log}\left(1-D_1\left(G_1(x)\right)\right)\right]\end{array}$$

$$\begin{array}{rl}{\mathcal{L}}_{\text{cyc}}(G_1,G_2)=& {\mathbb{E}}_x\left[\Vert x-G_2\left(G_1(x)\right){\Vert }_1\right]\\ & {\mathbb{E}}_y\left[\Vert y-G_1\left(G_2(y)\right){\Vert }_1\right]\end{array}$$

$$\begin{array}{r}{\mathcal{L}}_{\text{cPercep}}=\sum _{i=0}^L{\lambda }_{cp,i}\left(\Vert F_i\left(x\right)-F_i\left(\hat{\hat{x}}\right){\Vert }_1+\Vert F_i\left(y\right)-F_i\left(\hat{\hat{y}}\right){\Vert }_1\right)\end{array}$$

$$G{r}_i(x{)}_{m,n}=\frac{1}{h_i{w}_i{d}_i}\sum _{h=1}^{h_i}\sum _{w=1}^{w_i}{F}_i(x{)}_{h,w,m}{F}_i(x{)}_{h,w,n}$$

$$\begin{split}\mathcal{L}_{\small\textrm{cStyle}}=&\;\;\;\,\sum_{i=1}^{L}\lambda_{cs,i}\frac{1}{4d_{i}^{2}}\Big{(}\lVert{Gr_{i}\left(x\right)-Gr_{i}\left(\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{x}}}\right)}\rVert_{F}^{2}\\ &+\lVert{Gr_{i}\left(y\right)-Gr_{i}\left(\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{y}}}\right)}\rVert_{F}^{2}\Big{)}\end{split}$$

$$\begin{array}{rl}\underset{G_1,G_2}{\min }\underset{D_1,D_2}{\max }\mathcal{L}=& {\mathcal{L}}_{\text{adv}}(G_1,D_1)+{\mathcal{L}}_{\text{adv}}(G_2,D_2)+{\lambda }_{cP}{\mathcal{L}}_{\text{cPercep}}\\ & +{\lambda }_{cyc}{\mathcal{L}}_{\text{cyc}}(G_1,G_2)+{\lambda }_{cS}{\mathcal{L}}_{\text{cStyle}}\end{array}$$

Full text

I Introduction

In recent years, the machine learning community has achieved tremendous leaps in performance. This owes to increasingly available computational resources and open-source access to large datasets. From another perspective, radiological scans are vital tools in modern medicine. They enable diagnostics, disease tracking and patient treatment monitoring. This has led to the utilization of recent advances in computer vision, especially Deep Convolutional Neural Networks (DCNNs), in the field of medical image analysis. For example, DCNNs have been adapted for lesion classification in Magnetic Resonance Images (MRI) [1], 3D image segmentation [2] and anomaly detection [3] among other applications [4, 5].

A branch of deep learning is generative models which are utilized for dataset generation and augmentation. Amongst them, Generative Adversarial Networks (GANs) [6] are the prominent choice, with a large body of research focusing on theoretical and architectural analysis [7, 8]. In 2016, the pix2pix framework, a supervised GAN-based framework, has introduced the task of image-to-image translation from a source domain image, e.g. a day-time image, to a corresponding target domain image, e.g. a night-time image, provided that both domains have the same underlying structure [9]. This task has been adapted to the field of medical image analysis by using pix2pix for applications such as low-dose Computed Tomograph (CT) denoising [10], Positron Emission-computed Tomography (PET) to MR translation [11], splenomegaly segmentation [12] and MR to CT translation [13]. Other specialized architectures have been introduced for tasks such as compressed sensing MR reconstruction [14] and retinal image super-resolution [15]. Recently, we proposed MedGAN as a new framework for image translation [16]. It extends pix2pix with a cascaded U-net generator architecture [17] and additional non-adversarial losses, such as perceptual [18] and style transfer loss functions [19]. Since then it has been applied to tasks such as PET denoising [16], MR motion artifacts correction [20] and medical image in-painting [21].

However, these methods are supervised. In other words, training such models requires paired datasets where the images from the source domain are paired in a pixel-wise sense with the corresponding images in the target domain. Nevertheless, the acquisition of such paired datasets in real-life situations is often challenging. This is due to difficulties in obtaining co-registered cross-modality data from different scanners and acquisition sequences or multi-modal data for some organs such as the heart due to technical challenges or the required extensive planning and acquisition. Consequently, unsupervised image translation techniques, which are trainable with no paired examples but with image samples from both domains, are especially important for medical image translation.

Several methods for unsupervised image-to-image translation have been developed. UNIT is an unsupervised translation framework which maps the input-target images into a common latent space using a pair Variational Auto Encoder-GANs (VAE-GANs) before reconstruction in the desired image domain [22]. It has been utilized in the medical domain for the translation of T1-weighted and T2-weighted MR scans [23]. Cycle-GAN is another unsupervised translation approach which is based on the combination of adversarial losses and the pixel-wise cycle-consistency loss [24]. It has been adapted for medical translation tasks such as CT to MR bidirectional-translation [25] and CT denoising [26]. Other unsupervised frameworks exist with an overview available in [27].

In this work, we introduce a new framework for unsupervised medical image translation titled Cycle-MedGAN. This work expands the Cycle-GAN framework by introducing two new non-adversarial loss functions analogous to the perceptual and style transfer losses utilized in MedGAN. However, unlike MedGAN, the calculation of such losses does not require any explicit pairing of the input datasets during training. The training procedure is unsupervised using unpaired data, while validation is conducted on paired datasets.To validate the performance of the proposed framework, qualitative and quantitative comparisons against unsupervised frameworks, such as Cycle-GAN and UNIT, are conducted. The comparisons are carried out on the two medical tasks, MR motion artifact correction and PET to CT translation.

II Method

The proposed Cycle-MedGAN framework is based on the traditional Cycle-GAN framework with the inclusion of new non-adversarial losses. The baseline Cycle-GAN model is illustrated in Fig. 1.

II-A The Cycle-GAN Framework

Cycle-GAN is an unsupervised framework which allows bidirectional translation between the source domain $X$, e.g. PET images, and the target domain $Y$, e.g. CT images. It consists of two mapping functions $G_{1}:X\rightarrow{Y}$ and $G_{2}:Y\rightarrow{X}$, where $G_{1}$ and $G_{2}$ are two generator networks parametrized using DCNNs. Each of the generator networks is trained adversarially using a corresponding discriminator network, $D_{1}$ and $D_{2}$. For illustration, the first generator network $G_{1}$ receives as input a source domain image, $x\in X$, and outputs a synthetic translation, $\hat{y}=G_{1}(x)$. $D_{1}$ receives as input both the synthetic output $\hat{y}$ and an unpaired image randomly sampled from the desired target domain, $y\in Y$. The two networks, $G_{1}$ and $D_{1}$, are pitted against each other in competition, where $D_{1}$ acts as a binary classifier attempting to distinguish between the translated samples and the target domain samples. On the other hand, $G_{1}$ attempts to improve the quality of the translated output, thus deceiving the discriminator. This training procedure is formulated as a min-max optimization task over the adversarial loss function $\mathcal{L}_{\small\textrm{adv}}(G_{1},D_{1})$:

[TABLE]

and $\mathcal{L}_{\small\textrm{adv}}(G_{2},D_{2})$ is the analogous loss function for the second pair of networks, $G_{2}$ and $D_{2}$, formed by replacing the input images as $y$ and the translated outputs as $\hat{x}$.

Training the framework merely with the adversarial losses is not sufficient since it may lead to mode collapse, where a set of different input images are mapped into a single image in the target domain [24]. Therefore, an additional constraint regularizing the mapping functions is essential. This is achieved by Cycle-GAN which enforces the two mapping functions, $G_{1}$ and $G_{2}$, to be cycle-consistent with each other. In other words the two generator networks should invert each other such that $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{x}}}=G_{2}(G_{1}(x))\approx x$ and $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{y}}}=G_{1}(G_{2}(y))\approx y$. This behaviour can be incentivized by using the pixel-wise cycle-consistency loss for both generators:

[TABLE]

II-B Non-Adversarial Cycle Losses

Cycle-GAN relies on the cycle-consistency loss to avoid mismatches which could occur due to unsupervised training using unpaired images. However, it has been discussed in the literature that pixel-wise losses fail to capture the perceptual aspect of human judgement on image quality [28]. Thus, when used in translation tasks, they often lead to results which lack sharpness and fine-detailed structures [18, 19]. To circumvent this issue, feature-based loss functions were introduced as additional constraints to enhance the quality of translated output quality. For instance, the MedGAN framework utilized a combination of perceptual and style transfer losses [16]. However, for unsupervised image translation, the utilization of such loss functions is not viable. An unpaired translation model cannot be trained by penalizing the feature-based deviation of the translated image from the unknown ground truth image.

In this work, we propose an adaptation of the above feature-based loss functions for the task of unsupervised image translation. The penalized deviation is instead between the input images, $x$ or $y$, and the cycle-reconstructed images, $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{x}}}$ or $\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{y}}}$. This process is illustrated in Fig. 2. The first proposed loss function is the cycle-perceptual loss, $\mathcal{L}_{\textrm{cPercep}}$. Analogous to the perceptual loss introduced in [16, 18], this loss aims at minimizing the perceptual discrepancies and enhancing the global consistency of the output images. This is achieved by extracting intermediate feature maps, using a pre-trained feature extractor network, for both the input and the cycle-reconstructed images. The cycle-perceptual loss then calculated as the mean absolute error (MAE) between the extracted feature maps for both generators:

[TABLE]

where $F_{i}$ and $F_{j}$ indicate the extracted feature maps from the $i^{\textrm{th}}$ layer of the feature extractor network. $L$ is the total number of layers, $\lambda_{cp,i}$ is the weight given to each layer.

The second proposed loss function is the cycle-style loss, which is typically utilized for style transfer applications [19]. This loss aims at matching the texture, style and fine details of the input images onto the cycle-reconstructed images. As a result, this motivates the generator architectures to produce more detailed translated outputs. The cycle-style loss is computed by first calculating the feature map correlations over the depth dimension. For $G_{1}$, this is represented by the Gram matrices, $Gr_{i}(x)$ and $Gr_{i}(\hat{\vphantom{\rule{1.0pt}{5.71527pt}}\smash{\hat{x}}})$, whose elements are defined as:

[TABLE]

where $h_{i}$, $w_{i}$ and $d_{i}$ are the spatial height, width and depth of the extracted feature map of the $i^{\textrm{th}}$ layer of the feature extractor network.

The style loss is then calculated as the weighted average of the squared Frobenius norm of the Gram matrices:

[TABLE]

with $\lambda_{cs,i}$ is the weight given to the Gram matrices of the $i^{\textrm{th}}$ layer.

For the Cycle-MedGAN framework, a combination of the adversarial, cycle-consistency, perceptual consistency and style consistency losses is utilized. The final min-max optimization task is given by:

[TABLE]

where $\lambda_{cP}$, $\lambda_{cyc}$ and $\lambda_{cS}$ are the weights given for the cycle-perceptual, cycle-consistency and cycle-perceptual losses respectively.

III Datasets and Experiments

The Cycle-MedGAN framework was evaluated on two different tasks, PET to CT translation and the correction of motion artifacts in MR. For PET-CT translation, a dataset of the head region from 46 anonymized volunteers was acquired using a joint PET-CT scanner (Siemens Biograph mCT). For training, 1935 two-dimensional slices from 38 patients were utilized while the remaining 420 slices from 8 separate patients were used for validation. For the second application, T1-weighted MR data for the head region was acquired for 17 anonymized volunteers using a 3T MR scanner (Siemens Biograph mMT) with a fast spin echo sequence. Two different scans were acquired for each volunteer, one with voluntary rigid motion (head tilting) of the head and another under resting conditions [29]. Another 980 slices from 14 patients were used for training and 105 slices from the remaining 3 patients were used for validation. For both applications, the resolution of extracted data slices was re-sampled from their original resolutions to an isotropic voxel size of $1\textrm{mm}^{3}$ and two-dimensional images of pixel dimensions $256\times 256$ were extracted.

Analogous to Cycle-GAN, random shuffling was applied between the different subjects in the collected datasets to ensure no explicit pairing between the source and target domains occur during the training procedure. The discriminator from a BiGAN network was utilized as the feature extractor network in all experimentations [30]. The BiGAN was pre-trained on a separate dataset of whole-body CT data for image reconstruction. This was conducted to extract plausible feature maps which would improve the quality of translated images.

To evaluate the performance of the proposed framework, qualitative and quantitative comparisons with other unsupervised translation techniques were carried out. More specifically, the UNIT framework [22] and the Cycle-GAN [24] were considered as performance baselines for the comparison. To ensure faithful representation of the baseline methods, verified open-source implementations were utilized along with the recommended hyper-parameters in the original publications [31, 32]. Besides the pre-trained feature extractor, the Cycle-MedGAN framework has an identical architecture as the Cycle-GAN framework which is described in details in [24]. The quantitative comparisons utilized the following metrics: Peak Signal to Noise Ratio (PSNR), Structural Similarity Index (SSIM) [33], Learned Perceptual Image Patch Similarity (LPIPS) [34] and Visual Information Fidelity (VIF) [35]. All models were trained using the ADAM optimizer and a batch size of 64. Training was for 50 epochs, lasting approximately 24 hours, using an NVIDIA Titan X GPU.

IV Results and Discussion

The results of the proposed Cycle-MedGAN framework are presented in Table I and Fig. 3 in comparison with UNIT and Cycle-GAN. Qualitatively, the worse performance is exhibited by the UNIT framework. For both medical datasets in the comparative study, UNIT results in inhomogeneous global deformations as well as substantial blurs and distortion in the translated images. This is also reflected quantitatively, with UNIT resulting in the worst scores in Table I across the chosen metrics. In contrast, the resultant images produced by Cycle-GAN framework have a global structure which closely matches that of the target ground-truth images. However, finer details are not accurately translated by Cycle-GAN, such as the bone structures in the resultant CT images, and the motion blurring due to rigid motion in MR is not completely removed. The proposed Cycle-MedGAN framework builds upon the traditional Cycle-GAN architecture by introducing addition non-adversarial cycle losses to regulate the generator architectures. This results in an enhanced visual quality in the translated images (e.g. sharp edge delineation in MR images). The resultant CT images have noticeably sharper and more consistent bone structures compared to the other frameworks. Additionally, motion blurring due to rigid motion in MR is minimized. This enhancement in performance is analogously reflected quantitatively in Table I with Cycle-MedGAN surpassing the comparison baselines on all chosen metrics.

Despite the added level of quality by the Cycle-MedGAN framework, the proposed technique is not free from drawbacks. First, the resultant images by the Cycle-MedGAN framework potentially overlook important diagnostic information in the translation process. Thus, the framework is not intended for diagnostic applications but rather for post-processing tasks. An example of such tasks is using the synthetic CT images for PET attenuation correction or the calculation of organ volumes from motion corrupted MR images. Additionally, phase information contains vital information for motion correction in MR. Therefore, in future studies, we plan on expanding the framework with multi-channel three-dimensional inputs for complex-valued data. Furthermore, we plan to expand the comparison study to include performance on clinical applications and the effect of each individual non-adversarial loss in comparison to different loss combinations.

V Conclusion

Cycle-MedGAN is a new framework for unsupervised image translation. It builds upon the widely utilized Cycle-GAN framework with the additional utilization of new non-adversarial cycle loss functions, namely the cycle-perceptual loss and the cycle-style loss. The new loss functions use intermediate feature maps, extracted from a pre-trained feature extractor network, to direct the generator architectures to minimize perceptual and textural discrepancies in the results. Quantitative and qualitative comparisons with other unsupervised translation techniques indicate that the proposed framework enhances the translated outputs for the tasks of PET to CT translation and MR motion correction.

In the future, we plan to enhance the framework from the architectural aspect by incorporating three-dimensional complex-valued data. Additionally, the diagnostic performance of the framework will be investigated by experienced radiologists conducting subjective evaluations of the results.