fRegGAN with K-space Loss Regularization for Medical Image Translation

Generative adversarial networks (GANs)[1] have gained much attention in the last few years due to their ability to generate realistic images. In the medical imaging field, they are often used for image-to-image translation problems, such as mapping magnetic resonance imaging (MRI) to computed tomography images[2, 3], T1-weighted to T2-weighted MRI[4, 5], or low- to high dose contrast-enhanced MRIs[6, 7]. Such approaches could potentially reduce healthcare costs and patient burden while maintaining or even improving the diagnostic value of a modality. While GANs have shown promising results in these tasks, they tend to suffer from a frequency bias towards low frequencies[8]. This is especially problematic for medical image-to-image translation tasks, where preserving the images’ high-frequency content (i.e., edges) is crucial. Local errors in the generated images can lead to removing essential structures, such as lesions, which can significantly impact downstream tasks. To address this problem, we propose a novel frequency-aware image-to-image translation framework based on the supervised RegGAN framework[9]. We employ a K-space loss, which can regularize the frequency content of the generated images (see Figure 1).

Obtaining well-aligned data is often a challenging task in medical image-to-image translation. Common image misalignments, like poor registration of source and target domain samples, can lead to significant degradation in the performance of supervised approaches like Pix2Pix[10], which rely heavily on pixel-wise reconstruction losses. While unsupervised training methods like CycleGAN[11] can alleviate this issue, they generally tend to have lower performance compared to supervised methods. Huang et al.[9] proposed the supervised RegGAN medical image-to-image translation approach. In the RegGAN framework, the misaligned target images are considered noisy labels, and the generator is trained with an additional registration network estimating a displacement vector field to fit the misaligned noise distribution adaptively. We are not the first to investigate the use of the frequency domain for improving the performance of GANs, but our work is the first to apply the frequency domain to the medical image-to-image translation task and to combine it with the RegGAN. In the field of image processing and generation, several recent studies [12, 13, 14] have explored the use of the frequency domain for improving the performance of various models. The closest to our idea is in the work of Cai et al. [15], where a frequency domain image translation framework was proposed for image-to-image translation. The framework exploits frequency information by adding multiple losses based on the frequency-transformed images to the optimization problem of the generator. They showed natural images that the proposed method could improve the performance of image-to-image translation and synthesis. An open question is if the proposed method can be combined with cycle consistency and applied to medical images. We summarize our contributions as follows:

1. 1.

   We further investigate the RegGAN approach by Huang et al. [9] by generally applicable modifications to the architecture and training procedure. We extend their CycleGAN approach by adding a second registration network, which we call Reg²CycleGAN.

2. 2.

   We incorporate constraints in the frequency domain (K-space) to regularize and guide the network training process. This regularization technique is motivated by MRI acquisition and reconstruction methods which leverage the distribution of image feature information in K-space to reduce noise and/or accelerate image acquisition.

3. 3.

   Our proposed method is evaluated on the publicly available BraTS dataset. The results show that our method outperforms the baseline (i.e., RegGAN and CycleGAN) methods in quantitative and qualitative metrics. We also provide detailed ablation studies to understand the effect of each modification on the final performance. Our proposed method is a step towards improving the performance of image-to-image translation and synthesis in the medical domain to bring GAN-based methods closer to clinical application.

This section begins with a brief overview of how vanilla GANs[1] are formulated and then summarizes later improvements like the LSGAN[16] and RegGAN[15]. Afterward, we describe our new K-space Loss regularization and its integration into the RegGAN and the CycleGAN. GANs are generative models composed of two networks, a generator $G$ and a discriminator $D$. The generator learns a mapping $G:X\rightarrow Y$ from source domain $X$ to target domain $Y$ such that $\widehat{y}=G(x)$ is close to $y$ w.r.t. a specific metric. The discriminator $D$ is a binary classifier trained to distinguish between real $y$ and generated $\widehat{y}$. Here, $x$ and $y$ denote sets of paired images $\{(x_{i},y_{i})\}_{i=1}^{N}$ where $x_{i}$ is an image in the source domain $X$ and $y_{i}$ is the corresponding image in the target domain $Y$. Following the notation from Kong et al.[9], the vanilla GAN optimization problem is:

$$\min_{G}\max_{D}\mathcal{L}_{GAN}(D,G)=\mathbb{E}_{y}\log D(y)+\mathbb{E}_{x}\log\left(1-D(G(x))\right)$$

(1)

Since optimizing $G$ and $D$ simultaneously is not possible, the optimization problem is solved by alternating between $D$ and $G$[1]. The vanilla GAN training procedure has stability issues, which motivated the proposition of several modifications. The most relevant for our experiments is the Least Squares GAN (LSGAN)[16], which replaces the binary cross entropy loss with the mean squared error loss, resulting in the optimization task:

	$\displaystyle\min_{D}\mathcal{L}_{D,LSGAN}(D,G)$	$\displaystyle=\mathbb{E}_{y}\left(D(y)-1\right)^{2}+\mathbb{E}_{x}\left(D(G(x))\right)^{2}$		(2)
	$\displaystyle\min_{G}\mathcal{L}_{G,LSGAN}(D,G)$	$\displaystyle=\mathbb{E}_{x}\left(D(G(x))-1\right)^{2}$

In the work of Kong et al.[9], the robustness of Pix2Pix[10], which uses an L1-loss for supervision, was analyzed in the presence of misaligned source and target domain image data. Based on their experiments on the publicly available BraTS benchmark, it is concluded that the L1-loss reconstruction term only works with well-aligned images, which is often not the case in medical imaging. Therefore, the RegGAN method proposes to replace the L1-loss with the following correction loss:

$$\min_{G,R}\mathcal{L}_{corr}(G,R)=\mathbb{E}_{x,y}\left|y-G(x)\circ R(G(x),y)\right|$$

(3)

Here $R(G(x),y)$ denotes a deformation vector field (DVF) which is estimated based on $G(x),y$ using VoxelMorph[17] and $\circ$ is the resample operator. Adding the smoothness constraint $\mathcal{L}_{smooth}(R)=\mathbb{E}_{x,y}\left|\nabla R(G(x),y)\right|^{2}$ for the DVF, the final RegGAN optimization problem is then formulated as follows:

	$\displaystyle\min_{D}\mathcal{L}_{D,reg}(D,G)$	$\displaystyle=\mathcal{L}_{D,LSGAN}(D,G)=\mathbb{E}_{y}\left(D(y)-1\right)^{2}+\mathbb{E}_{x}\left(D(G(x))\right)^{2}$		(4)
	$\displaystyle\min_{G,R}\mathcal{L}_{G,reg}(D,G,R)$	$\displaystyle=\mathcal{L}_{G,LSGAN}(D,G)+\lambda_{1}\mathcal{L}_{corr}(G,R)+\lambda_{2}\mathcal{L}_{smooth}(R)$

For unsupervised image-to-image translation, Zhu et al.[11] proposed to use two generators $F,G$, and two discriminators $D_{X},D_{Y}$, and to add a cycle consistency loss $\mathcal{L}_{cyc}(G,F)=\mathbb{E}_{x}\left|F(G(x))-y\right|+\mathbb{E}_{y}\left|G(F(y))-x\right|$. In additional experiments, they showed that an identity loss $\mathcal{L}_{id}(G,F)=\mathbb{E}_{y}\left|G(y)-y\right|+\mathbb{E}_{x}\left|F(x)-x\right|$ preserves content information. All four models are trained jointly, and the optimization problem for the generators is formulated as follows:

$$\min_{G,F}\mathcal{L}_{G,CycGAN}(D_{X},D_{Y},G,F)=\mathcal{L}_{G,LSGAN}(D_{Y},G)+\mathcal{L}_{G,LSGAN}(D_{X},F)+\lambda_{3}\mathcal{L}_{cyc}(G,F)+\lambda_{4}\mathcal{L}_{id}(G,F)$$

(5)

Kong et al.[9] also proposed combining cycle consistency with their registration loss

$$\min_{G,F,R}\mathcal{L}_{G,RegCycGAN}(D_{X},D_{Y},G,F,R)=\mathcal{L}_{G,reg}(D_{Y},G,R)+\mathcal{L}_{G,LSGAN}(D_{X},F)+\lambda_{3}\mathcal{L}_{cyc}(G,F)+\lambda_{4}\mathcal{L}_{id}(G,F)$$

(6)

which turns the unsupervised training task into a supervised one. This has the beneficial effect of reducing the number of solutions for the generators and improving the overall quality.

We perform two sets of experiments. First, we set a baseline and solely investigate our changes to the Reg²CycleGAN. Secondly, we explore the combination of the RegGAN and Reg²CycleGAN with our K-space loss (see equation (8)). All our experiments are trained on the public BraTS 2021 dataset[18, 19, 20]. BraTS has 1251 MRI scans of patients with brain tumors. The dataset contains four different modalities, but we only use the T1w and T2w modalities to train our models. All MRI scans are available as NIfTI files and were acquired with other clinical protocols and various scanners from multiple data-contributing institutions. In addition, all volumes are pre-processed, i.e., co-registered to the same anatomical template, interpolated to the same resolution ($1\text{ mm}^{3}$), and skull-stripped. Hence, the dataset is well aligned, and each volume has a size of $240\times 240\times 155$ voxels. To train and evaluate our models, we randomly split the dataset on patient-level into a training, validation, and testing set with 1000, 51, and 200 volumes, respectively. We further pre-process the dataset by normalizing each volume to the range $[0,1]$ based on the $0.5$ and $99.5$ percentile of the volume. Afterward, we sliced the 3d volumes in axial direction into 2d images and removed all slices without information. This results in 139221, 6891, and 27956 2d slices for training, validation, and testing, respectively. Based on the survey papers[21, 22, 23], we selected to most used metrics for evaluating our models: peak signal-to-noise ratio (PSNR), structural similarity index (SSIM), and multi-scale structural similarity index (MS-SSIM). The SSIM is a measure of the similarity between two images. It is based on the mean and variance of the local luminance and contrast of the images. The MS-SSIM is a multi-scale version of the SSIM. It is defined as the geometric mean of the SSIM at different scales. The SSIM and MS-SSIM are in the range $[0,1]$ and a higher value indicates a better synthetization. Both metrics are implemented using the TorchMetrics 0.10.3 library[24].

We have proposed a novel image-to-image translation approach in the medical imaging field called fRegGAN. Our method extends the RegGAN approach by incorporating a Fourier-space based loss to regularize and guide the network training process. Our extensive experiments on the BraTS 2021 dataset show that our proposed method outperforms the baseline methods in terms of both quantitative metrics and qualitative visual inspection. The frequency regularization loss helps the optimization process find solutions that preserve more details and are closer to the ground truth. However, in our evaluation, we found that comparison based on quantitative metrics with reference images has some limitations. Skull stripping (to anonymize patients) resulted in some artificial artifacts between input and output volumes (i.e., T1w and T2w). Our results prove that incorporating Fourier-space information can improve image-to-image translation’s performance and visual quality in the medical domain. Our findings suggest that fRegGAN could have broader implications for the medical imaging field. For example, frequency regularization, an essential aspect of our proposed method, might be beneficially applied to image synthesis tasks using diffusion models. The improvements observed in image-to-image translation tasks could potentially translate to enhancements in the quality and accuracy of synthesized images, particularly in complex tasks such as those involving multiple modalities. Moreover, the idea of k-space to k-space synthesis, although not directly addressed in our study, could serve as a promising future research direction. Such an approach might further aid clinical decision-making, diagnostics, and treatment planning by providing more accurate and realistic synthetic medical images.

I.M.B., M.L., F.K., and A.H. conceived the method and experiments, I.M.B. and A.H. conducted the experiments, I.M.B., M.L., M.D., A.H., analyzed the results. All authors reviewed the manuscript.