Multi-scale GANs for Memory-efficient Generation of High Resolution Medical Images

TL;DR

This paper introduces a multi-scale patch-based GAN method that efficiently generates high-resolution 2D and 3D medical images, overcoming computational limitations and improving image quality over traditional patch approaches.

Contribution

The novel multi-scale GAN approach enables high-resolution 2D and 3D medical image generation with constant GPU memory, outperforming existing patch-based methods in quality and artifact reduction.

Findings

Generated 3D thorax CTs of size 512x512x512

Produced high-resolution thorax X-rays of size 2048x2048

Achieved constant GPU memory demand regardless of image size

Abstract

Currently generative adversarial networks (GANs) are rarely applied to medical images of large sizes, especially 3D volumes, due to their large computational demand. We propose a novel multi-scale patch-based GAN approach to generate large high resolution 2D and 3D images. Our key idea is to first learn a low-resolution version of the image and then generate patches of successively growing resolutions conditioned on previous scales. In a domain translation use-case scenario, 3D thorax CTs of size 512x512x512 and thorax X-rays of size 2048x2048 are generated and we show that, due to the constant GPU memory demand of our method, arbitrarily large images of high resolution can be generated. Moreover, compared to common patch-based approaches, our multi-resolution scheme enables better image quality and prevents patch artifacts.

Tables

Table 1. Table 1: Quantitative results: Measurements between a generated image and its ground truth. Columns 3-5: average SSIM (higher is better), MAE and MSE (lower is better). Experiments (top and bottom): B80f to B20f image translation on a thorax CT dataset, and image generation on thorax X-Ray. Compared to ground truth (row-wise): For both experiments – our generated images, upscaled low resolution images, conventional patch-wise generation; for thorax CT – the original B80f image and a non-local means filtered B80f image. Subscript ( ⋆ ) indicates significantly worse performing methods in terms of all measures compared to ours in a paired t-test ( p < 0.0001 𝑝 0.0001 p<0.0001 ).

		SSIM	MAE	MSE
Dataset	Method	mean($\pm std$)	mean($\pm std$)	mean($\pm std$)
Thorax CT $512\times 512\times 512$	Our gen. B20f	$0.773\pm 0.025$	$0.033\pm 0.004$	$0.004\pm 0.001$
	Small gen. B20f⋆	$0.633\pm 0.024$	$0.058\pm 0.004$	$0.011\pm 0.002$
	Patch gen. B20f⋆	$0.706\pm 0.047$	$0.049\pm 0.007$	$0.008\pm 0.002$
	Original B80f⋆	$0.480\pm 0.045$	$0.065\pm 0.008$	$0.012\pm 0.003$
	Filtered B80f	$0.773\pm 0.048$	$0.031\pm 0.005$	$0.004\pm 0.001$
Thorax X-Ray $2048\times 2048$	Our gen.	$0.711\pm 0.067$	$0.104\pm 0.028$	$0.022\pm 0.014$
	Small gen.⋆	$0.673\pm 0.064$	$0.108\pm 0.027$	$0.024\pm 0.014$
	Patch gen.⋆	$0.520\pm 0.072$	$0.200\pm 0.024$	$0.069\pm 0.019$

Equations

$${\mathcal{L}}_{cGAN}(G_0,D_0)$$

$${\mathbb{E}}_{x_0,y_0}[\log D_0(x_0,y_0)]+{\mathbb{E}}_{x_0,\mathbf{z}}[\log (1-D_0(x_0,G_0(x_0,\mathbf{z})))]$$

$${\mathcal{L}}_{cGAN}(G_i,D_i)$$

$${\mathbb{E}}_{x_{p_i},\mathbf{z}}[\log (1-D_i(x_{p_i},y_{p_{i-1}},G_i(x_{p_i},y_{p_{i-1}},\mathbf{z})))],$$

Full text

1 Introduction

Generative adversarial networks (GANs) [3] have shown impressive results for photo-realistic image synthesis in the last couple of years [4, 6, 12]. They also offer numerous applications in medical image analysis, such as generating images for data augmentation, image reconstruction and image synthesis or domain adaptation. Despite the undeniable success and the large variety of applications, GANs still struggle to generate images of high resolution. A reason for that is the fact that generated images are easier to distinguish from real ones at higher resolutions, which hinders the training process [6]. Further reasons are computational demands and memory requirements of current network architectures. To deal with the first issue, Karras et al. [6] recently proposed a progressive learning strategy for GANs that starts with low resolution and adds finer details throughout the training. Using this strategy, the authors were able to generate high resolution 2D images of size $1024\times 1024$. However, the training process already requires $\sim$16 GB of GPU RAM, which means that larger images can only be generated with special and expensive hardware. The problem is further aggravated in 3D imaging typically demanded in clinical routine. An early attempt for the 3D application of GANs is made in [10], where furniture shapes of size $64^{3}$ are generated using a DCGAN architecture. Even though the task of generating only shapes (no texture or intensities) is fairly simple, its computational demands are borderline to most consumer-class GPUs.

The image size limitations make GANs impracticable for many medical image applications, e.g. for thoracic CTs, brain MRIs, or high resolution X-ray images. In [9] the authors claim to be forced to use only half of the image size ($128\times 128\times 54$) due to GPU memory limits of the dedicated hardware (NVIDIA DGX system). The largest 3D output size of a GAN found in literature is $128^{3}$ [11], however no memory requirements are mentioned. An usual approach to overcome such computational challenges is slice-/patch-wise generation [7, 5]. Unfortunately, those methods suffer from artifacts between patches/slices due to noncontinuous transitions. Dealing with this problem by applying patch overlaps and averaging, leads to blurry results and loss of image detail. The intuition behind patch inconsistencies is that when patches are generated independently, they do not have any global intensity information. Thus, [5] proposes to additionally observe a larger area around each patch of the input image to cope with this issue. Even though shown to be well suitable for segmentation and would probably improve patch artifacts in strictly paired image translation (e.g. CT to MR), such an approach cannot be applied to image generation from scratch (or sparsely conditioned), and its effect is limited when the image size drastically exceeds the chosen patch size.

In this work we propose a memory-efficient multi-scale GAN approach for the generation of high-resolution medical images in high quality. Our approach combines a progressive multi-scale learning strategy with a patch-wise approach, where low-resolution image content is learned first, and image patches at higher resolutions are conditioned on the previous scales to preserve global intensity information. We demonstrate the ability to generate realistic large images on thoracic X-rays of size $2048^{2}$ and 3D lung CTs of size $512^{3}$. Further, we show that w.r.t. the growing side length of an isotropic 3D image, the memory requirements for popular GANs grow cubical, while they stay constant for any image size using our approach. Although the presented method is theoretically suitable for from-scratch image generation, in this work, we apply a conditional GAN for unsupervised domain adaptation. This application features topology preserving style transfer, where in contrast to strictly paired image translation, our approach does not assume corresponding images from two domains and enables the possibility to translate arbitrary datasets to a desired domain. We show that image translation from different CT reconstruction kernels, different devices and different acquisition parameters to a particular desired domain is possible, keeping the original 3D image size. The generated images are evaluated with respect to a known ground-truth image emphasizing the quality of generated images.

2 Methods

GANs are generative models that learn to map a random noise vector $\mathbf{z}$ to an output image $y$ using a generator function $G:\mathbf{z}\rightarrow y$ [3]. An extension of regular GANs are the conditional GANs, that learn the mapping from an observed image $x$ additionally, $G:\{x,\mathbf{z}\}\rightarrow y$. To ensure that the generator produces realistically looking images that cannot be distinguished from real ones, an adversarial discriminator $D$ is enclosed in the training process, aiming to perfectly distinguish between real images and generator’s fakes.

2.1 Multi-scale Conditional GANs

The idea of multi-scale image generation to obtain high resolutions, is inspired by early works like [2], where, based on a Laplacian pyramid, multiple GANs for increasing resolutions are used, and more recent approaches [12] that achieve photo-realistic image quality by iterative resolution refinement applying successive GANs. However, those methods still assume that at a certain stage the whole full-resolution image is propagated through the network, requiring extremely high GPU capacities especially for 3D images. Here, we propose to overcome this issue by first generating the whole image in a low resolution, then iteratively proceed by generating image patches of constant size but growing resolutions. Since the GANs for each scale are trained separately, a particular maximum size is never exceeded during training, so the GPU demand stays constantly low, facilitating the generation of arbitrarily large images. Further, we enable a consistent and realistic look of the generated images by conditioning each high resolution scale on the previous lower resolution image, thus providing patches with global intensity information and preventing patch artifacts and inconsistencies (see Fig. 1 for a method overview). A further benefit of progressively learning increasing resolutions, is that each GAN has a rather simple task to learn compared to generating the whole high-resolution image from scratch. This does not only enable a simpler training process, but also ensures sharp results for each scale and thus a consistent high-resolution result.

We consider an additional condition, namely the object edges, since our use-case scenario implies topology preserving image domain translation. In the lowest resolution (LR) GAN, the whole low-resolution edge image is used as an input. In the higher-resolution (HR) GANs, two conditions are used: a patch from the image of the previous resolution, upscaled to the size of the current scale; and a patch from the edge image of the current scale. Note, that the patch size of each resolution is constant, however, as the overall image becomes larger after each HR GAN, the relative image fraction per patch decreases, while resolution grows. The objective of this learning process can be described as follows. For multiple conditional images $x_{0}\dots x_{n}$ with resolutions $0\dots n$, output images $y_{0}\dots y_{n}$, where $y_{n}$ represents the final output image, are generated using separate generators $G_{0\dots n}$ and discriminators $D_{0\dots n}$ with the objectives

[TABLE]

where $x_{p_{i}}$ and $y_{p_{i}}$ are patches of the conditional image $x_{i}$ and the generated image $y_{i}$ respectively, with $i\in[1,n]$.

2.2 Architecture and Training

Different generator’s architectures were chosen for $G_{0}$ and $G_{1\dots n}$, since the tasks of generating whole low resolution images and high resolution patches differ in a variety of requirements. The LR GAN uses a U-Net architecture, which is able to filter out many unimportant details and generalize better due to its bottleneck. Its tendency to result in more blurry images is negligible in the context of low-resolution images. For the patch generation by the HR GANs, ResNet blocks are chosen, since they are known to produce sharp results by keeping the input image resolution unchanged. The higher overfitting risk of not having a bottleneck is diminished due to the stronger conditioning (on the previous scales) and the overall large number of patches compared to the number of images used. For the discriminators $D_{0\dots n}$ a regular fully-convolutional architecture is chosen111The complete architecture is available in the supplemental materials..

Data augmentation is crucial to our method, because cascading approaches are prone to propagate errors up from lower resolutions. To deal with this issue we corrupt a percentage of the low resolution images while training each HR GAN: we apply random noise, Gaussian blurring and vary the image resolution, also to make the patch generation less dependent on perfect edge extraction, the edge images are perturbed with noise. In our experience those techniques help to overcome overadapting to either of the inputs.

3 Results and Experiments

3.1 Memory Requirements for 3D GANs

GANs are currently rarely applied to 3D images due to computational constrains, therefore in this experiment we investigate the dependence of 3D image side length and memory requirements. Three common GAN architectures are chosen as baselines: DCGAN [10], Pix2Pix [4] and progressive growing GAN (PGGAN) [6], and compared to the two architectures of our method: LR 64 for low resolution images of size $64^{3}$ and HR 32 for high resolution patches of size $32^{3}$. PyTorch [8] is used as implementation framework of choice for all networks. Since Pix2Pix and PGGAN are only implemented for 2D images, a straight-forward translation to 3D is obtained (replacing 2D convolutions by 3D convolutions, etc.). The RAM demand computation is realized using the summary approach from the keras framework [1]. The assumed lower bound of memory usage here includes one forward and backward pass for the generator and discriminator each, as well as the memory required to store the images, gradients and network’s parameters for batch size of one. The results for different image sizes are shown in Fig. 2. Naturally, all three baseline approaches have an at least cubic memory requirement growth w.r.t. the image side’s length. For those approaches calculations for size over $128^{3}$ were not even possible on the used Titan XP 12GB GPU, thus the extrapolated cubic regressed curves are shown. These results underline the infeasibility of straight-forward 3D GAN approaches for medical images, as their sizes commonly reach $512^{3}$. In contrast, our method has a constant character and is thus suitable for arbitrary image sizes with predictable memory usage.

3.2 High Resolution Medical Image Generation

Data.

Our experiments use the following datasets (see Fig. 2 for examples):

Thorax CT: 3D CT images ($\sim 512^{3}$) of 56 subjects with varying degree of chronic obstructive pulmonary disease. For each subject the data is reconstructed simultaneously with a soft (B20f) and a sharp (B80f) kernel.

Low-dose thorax CT: Low-dose 4D lung CT images (120kVp, 40mAs) of 12 patients acquired during free breathing, resulting in 166 3D volumes in total. Image intensities are rescaled with a lung window setting and the images are cropped with a bounding box around the lungs resulting in about $320^{3}$ voxels.

Thorax X-ray: An open access chest X-ray dataset from the Indiana University containing frontal and lateral images usually around the size of $2000^{2}$. In our experiments about 1500 frontal images of size $2048^{2}$ are used.

Experimental Setup. Our experiments are based on unpaired image domain translation, so the trained GANs are conditioned on the object’s edges. A few scenarios are considered: (1) B80f to B20f CTs: Since images reconstructed with a sharp B80f kernel are very noisy, a reconstruction with a soft B20f kernel is advantageous for automatic quantitative computations (e.g. emphysema index). For this purpose, the GANs are trained in a 5-fold cross-validation manner on the B20f images from the thorax dataset. In test phase, the edges from the B80f images of the test patients are extracted and translated into the B20f domain using the trained GANs. (2) Low-dose to high-dose: In this experiment, we show that translations between different image domains of different devices and protocols are possible with our method, since it does not require paired data. Here, the GANs are trained on the high-dose B20f images from the thorax dataset, and in test phase, the edges extracted from the low-dose data are used to enable a translation between the two datasets. (3) Large 2D image generation: To show that our approach is not constricted to 3D images, in this experiment huge X-ray images are generated. The GANs are trained on $90\%$ of the thorax X-ray datasets, whereas the rest of the images is used for testing.

Results. For evaluation of experiment (1) the actual B20f reconstructions of the test B80f images serve as ground truth and for experiment (3) the real X-ray images. Here structural similarity index (SSIM), mean absolute error (MAE) and mean squared error (MSE) between the ground truth and the generated images are used as image quality criterions (Tab. 1). Our method is compared to other standard methods: generate a smaller image with an LR GAN ($64^{3}$ is about our computational limit) and upscale it; develop a straight-forward patch-based approach for the high resolution images (HR GAN without conditioning on previous scales) and apply stitching. The comparison networks were trained in the exact same manner as for our method. In all experiments isotropic patch sizes of 64 for the LR GAN, and 32 for the HR GANs are used. To avoid padding artifacts on patch borders, only the network’s receptive field parts of the generated patches are considered. However, for our method no patch overlaps are needed, whereas for the conventional patch-based method, patch overlaps of 5 pixels are applied. The significantly higher SSIM values and smaller pixel-wise distances for our method show its ability to outperform conventional GAN-based approaches. Visually, the generated images have a consistent appearance, barely any artifacts and high resolution, enabling the visibility of tiny structures within the lungs. In contrast, when using conventional patch-wise approaches, patch artifacts and inconsistencies are clearly present (Fig. 2).

We further apply a state-of-the-art denoising method, the non-local means filter, to match the B80f and B20f images. In terms of metrics, this approach is comparable to ours, however its computation time is $\sim 100$ times longer and it does not apply to any other domain transfer tasks. Also, the filtered images lack important details and have an inconsistent appearance visually. 222See supplementary for larger/ further images

For the low-dose dataset no ground truth is available, so we evaluate experiment (2) qualitatively and compare the feature space distributions of both domains using the Fréchet inception distance (FID) between: the original low-dose and the original high-dose images (150); the translated and the original high-dose images (131); the translated and the original low-dose images (157). The visual correspondence of the target domain and the translated images is also underlined by the smaller FID value. The large FID between the original low-dose and the translated images indicates that different features are extracted for the same image of various appearances, and thus emphasizes the need for domain translation.

4 Discussion and Conclusion

In this work, we propose a multi-scale GAN-based approach for the generation of arbitrarily large 3D medical images of high resolution and realistic homogeneous appearance. We show that in contrast to existing GANs, that produce the whole image at once, the presented patch-based method, only requires constant GPU memory w.r.t. the image size. Also compared to trivial patch-wise methods, our sophisticated multi-resolution scheme provides higher quality images in terms of image consistency and resolution. Here we show a use-case scenario for medical image domain translation based on conditional GANs, however our method is also suitable for image generation from scratch and it enables various applications and possibilities for medical images.

Appendix 0.A Supplementary

0.A.1 Network Architectures

0.A.2 3D thorax CT example image slices (axial, coronal,sagittal)

0.A.3 3D brain MRI example image slices (axial): T1 to T2 translation

0.A.4 2D thorax X-ray example images