COVID-19 PNEUMONIA CHEST X-RAY PATTERN SYNTHESIS BY STABLE DIFFUSION

Abstract

In this study, we fine-tuned a stable diffusion model to synthesize high resolution chest X-ray images (512×512) with bilateral lung edema caused by COVID-19 pneumonia using the class-specific prior preservation strategy. 300 positive images were selected from the MIDRC dataset as subject instances with an additional 400 negative images for class prior preservation. We synthesized images respectively using the new technique and the conventional technique for comparison. The synthetic images by the stable diffusion fine-tuned by the prior preservation technique have the Frechet inception distance (FID) of 9.2158 and kernel inception distance (KID) 0.0818 computed with the real positive images, which is superior to the synthetic images using the conventional methods such as WGAN and DDIM. The classification accuracy is 0.9975 with precision of 1.0 and recall of 0.9950 when the synthetic positive images with the real negative images were classified by a trained vision transformer (ViT). We conclude that the stable diffusion model can synthesize high-quality and high-resolution chest x-ray images using the prior preservation strategy with a small number of real images as subject instances and text prompt as guidance for the designated patterns.

1. INTRODUCTION

The severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), emerging in the late 2019, has caused massive loss of life and economic disruption globally. The World Health Organization (WHO) Coronavirus (COVID-19) Dashboard shows that the global death toll is approaching to 7 million with over 771 million confirmed cases [1]. On the other hand, the COVID-19 pandemic expedited the innovation and application of new technologies for detection, diagnosis, evaluation, and prognosis of COVID-19 disease. For example, following the early clinical observation on COVID-19 in Wuhan, China that reported that the unique bilateral or unilateral multiple mottling and ground-glass opacity patterns were identified on the chest X-ray (CXR) and CT images [2]. Based on this finding, a variety of machine learning methods, particularly using deep learning, were proposed to detect such visual patterns on the images. However, these new methods have multiple challenges regarding the accuracy, reliability, and robustness in field use case. A review on deep learning for COVID-19 image pattern detection pointed out the main challenges include the lack of image quality assurance, imbalance and diversity in the image datasets, and the generalization and reproducibility of the available models and algorithms [3].

Generative models such as generative adversarial network (GAN) and diffusion model provide an effective way to enrich the data diversity, which could help remedy some of the above challenges. Stable diffusion is a latent diffusion model (LDM) proposed in 2022 [4]. Instead of sampling from a pure Gaussian distribution, LDM takes the random noise from the encoded latent space by the pretrained embedding models. Therefore, the image generation runtime of LDM is significantly lower than the pure diffusion models. In addition, a new technology named Dreambooth is introduced to fine tune the stable diffusion model to generate high-quality, high-fidelity, and high-resolution synthetic images with desired patterns with a small dataset [5], however, it is demonstrated on non-medical images. Motivated by its success, in this study we aim to explore encoding medical visual patterns into a pretrained LDM and respectively bind them with medical knowledge.

In this paper, we first explain the importance of detecting lung edema on CXR images of COVID-19 pneumonia. Next, we briefly present the LDM and our experiments to fine tune a stable diffusion model. The synthetic images are compared with the images generated by Wasserstein GAN with gradient penalty (WGAN-GP) and denoising diffusion implicit model (DDIM) respectively. Finally, we summarize our findings in the conclusion section.

2. PULMONARY EDEMA DETECTION

Pulmonary edema of pneumonia is a key indicator of acute respiratory distress syndrome (ARDS) [6], which is used to estimate the severity of SARS-CoV-2 infection and the disease prognosis. A cohort study found that the time of mechanical ventilation for patients was quantitatively related to the degree of lung edema score [7]. The radiographic assessment of the lung edema (RALE) scoring system is a quantitative tool to measure the alveolar opacities on chest radiographs that reflect the degree of pulmonary edema [8]. A multicenter study found the progressive increase of RALE score was associated to higher mortality and longer use of ventilator [9]. The mRALE is a modified RALE scoring system whose score range is simplified from 0 to 24, where mRALE = 0 means no pulmonary edema and mRALE = 24 indicates severe bilateral lung edema. It offers a handy tool to assess the severity of COVID-19 pneumonia [10]. If we can use a generative model to synthesize CXR images with text guidance to control the presence of pulmonary edema, it could provide an effective solution to enrich the diversity and robustness of the current dataset for automated diagnosis and pattern detection for COVID-19 or other new or rare disease.

3. DIFFUSION MODEL FOR IMAGE SYNTHESIS

Diffusion models outperform existing image generative models such as VAE (variational autoencoder), GAN, and flow models [11]. The adversarial optimization nature of GAN makes it vulnerable to mode collapse and low diversity in training [12]. In contrast, the diffusion models use the Markovian process to model the image generation process by adding random noise sampled from Gaussian distributions during the forward diffusion process, then it uses a neural network to predict and remove the noise to generate new images during the reverse diffusion process.

3.1. Denoising diffusion models

The general architecture of a denoising diffusion probabilistic model (DDPM) [13] is illustrated as below.

The diffusion process is divided into $T$ steps, where the Gaussian noise $ϵ~\left(\mathrm{0,1}\right)$ is incrementally added to the real image ($x_0$) weighted by a random $\beta \in \left(\mathrm{0,1}\right)$ constrained by ${\alpha }^2+{\beta }^2=1^2.\alpha$ is considered a scaler to weight the original pixel value in the image at any step $t=T,T-1,\dots ,1$. With use of the reparameterization trick, the pixel values at any time step $t$ is presented as:

$$\begin{gathered} x_t=\sqrt{1-a_t{a}_{t-1}{a}_{t-2\dots }{a}_1}*ϵ+\sqrt{a_t{a}_{t-1}{a}_{t-2\dots }{a}_1}*x_0 \\ =\sqrt{1-{\bar{\alpha }}_t}*ϵ+\sqrt{{\bar{\alpha }}_t}*x_0,\text{let}\bar{\alpha }={\prod }_0^T{a}_t \end{gathered}$$ (1)

Then the forward diffusion process at step $t\in (1,T)$ can be modeled by a function $q\left(x_t\mid x_0\right)$ defined by:

$$q\left(x_t\mid x_0\right)=\sqrt{{\bar{a}}_t}{x}_0+\sqrt{1-{\bar{a}}_t}ϵ,ϵ~N\left(\mathrm{0,1}\right)$$ (2)

In Eq. 2, the term ${\bar{a}}_t$ can be computed by a linear scheduler with the value between 1 × 10⁻⁴ and 2 × 10⁻² [13]. We can use a closed-form solution of $q\left(x_t\mid x_0\right)$ to add noise $ϵ$ to the original image $x_0$ to attain the noised image $x_t$ at step t.

In the reverse diffusion process, we use a neural network to model the noised image at any step $t$ defined by:

$$p_{\theta }\left(x_{t-1}\mid x_t\right)≔N\left(x_{t-1};{\mu }_{\theta }\left(x_t,t\right),{ϵ}_{\theta }\left(x_t,t\right)\right)$$ (3)

where $\theta$ represents the learnable parameters of the neural network. To improve the training stability, we can replace ${ϵ}_{\theta }\left(x_t\right)$ with ${\beta }_t$ [13], representing the noise variance at step t defined by ${ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\right)$. To optimize the model parameter $\theta$, we use the Kullback–Leibler (KL) divergence loss as objective for model optimization, where $p\left(x_{t-1}\mid x_t\right)$ is the predict distribution and $q\left(x_t\mid x_{t-1}\right)$, is the ground true distribution:

$$L_{\mathit{\text{DDPM}}}={\mathbb{E}}_{t,x_0,ϵ}\left[{‖ϵ-{ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\right)‖}^2\right]$$ (4)

The DDPM succeeded in synthesizing high-quality and high-resolution images, but it must follow the Markovian process with many reversed steps (800~1000) to synthesize images. In addition, DDPM uses the identical large dimensionality for the latent variables as the input images. Thus, DDPM consumes expensive runtime and memory for both training and image generation which constrained its application to medical imaging.

3.2. Stable diffusion

The stable diffusion method is based on the latent diffusion model (LDM) architecture with multiple pretrained embedding models to reduce the dimensionality of the feature maps, resulting in significantly improved computational resource use during model training and inference [4].

There are three main components in the LDM architecture: a VAE, a text encoder $\left({\tau }_{\theta }\right)$, and a UNet with multi-head cross-attentions layers. The VAE encoder converts the input images into a lower dimensional latent representation as input to the forward diffusion process, and the decoder is responsible for converting the latent representations back to real images during the reverse diffusion process. The text encoder takes the text prompts as guidance for image synthesis and converts them into embedding vectors with a language model ${\tau }_{\theta }$. The encoded information is mapped to the UNet via the multi-head cross-attention layers. In addition, other spatially aligned inputs such as semantic map, representation, and images can be concatenated into the embedding as extra conditionings. Finally, the inner diffusion model (UNet) takes the encoded feature representations and acts as a conditional image generator by augmenting denoising process with cross-attentions.

The LDM optimization is similar as the general diffusion models. The main difference is that it converts the pixel inputs into a latent input $z_t$ through the encoder of the VAE to replace the generated image $x_t$ with a constraint ${\tau }_{\theta }\left(y\right)$ to the UNet for the diffusion and denoising process. The loss objective is correspondingly revised as:

$$L_{LDM}={\mathbb{E}}_{t,z_0,ϵ,y}\left[{‖ϵ-{ϵ}_{\theta }\left(z_t,t,{\tau }_{\theta }\left(y\right)\right)‖}^2\right]$$ (5)

Benefited by the low dimensionality of the latent space, the denoising processing of the LDM is significantly faster than the pure diffusion models.

4. CONCEPT BINDING BY PRIOR PRESERVATION

4.1. Fine-tuning with class prior preservation

Another main challenge for using the stable diffusion model for synthesizing the CXR images with pulmonary edema patterns is the insufficient number of training examples. To address this issue, we used the prior preservation technique to fine tune the stable diffusion 2.1 model pretrained by large image datasets including ImageNet-21K and MS-COCO [4]. The prior preservation technique was proposed as DreamBooth by Google Research in 2022 [5]. This fine-tuning process, illustrated in Figure 2, uses a class-specific prior preservation loss combined with the reconstruction loss as the loss objective to train the LDM to bind the concept of a new subject with the corresponding visual patterns via the text encoder and the VAE.

Figure 2.

Fine-tuning with prior preservation

Accordingly, the loss objective is revised as

$$L={\mathbb{E}}_{x,c,ϵ,{ϵ}^{\prime },t}[{\omega }_t{\Vert {\hat{x}}_{\theta }\left({\alpha }_{t}x+{\sigma }_tϵ,c\right)-x\Vert }_2^2+\lambda {\omega }_{t^{\prime }}{\Vert {\hat{x}}_{\theta }\left({\alpha }_{t^{\prime }}{x}_{pr}+{\sigma }_{t^{\prime }}{ϵ}^{\prime },c_{pr}\right)-x_{pr}\Vert }_2^2]$$ (6)

where the first term is to constrain the generator to synthesize images with the desired lung edema pattern by the subject identifier “bilateral lung edema mRALE 24” encoded by the text-prompt encoder. The second term of the loss serves as the prior preservation term to supervise the generator from overfitting with the non-severe CXR images (mRALE <18) binding with encoded prior class identifier “chest x-ray”.

4.2. Experiments

We used the CXR image dataset from the 2023 MIDRC mRALE Mastermind Challenge [10], which contains 2,599 frontal CXR images annotated with mRALE score by medical experts. Due to hardware limitations, we randomly selected 200 images with mRALE=24 as severe lung edema subject instances and 400 images with mRALE<18 as prior class instances.

The experiments were implemented in Amazon SageMaker on a ml.p3.8xlarge instance with four Nvidia Tesla V100 GPUs and 64 GB GPU RAM. The LDM was optimized by the Adam optimizer with initial learning rate of 2 × 10⁻⁶ with mini-batch size of 1 for 20 epochs.

4.3. Assessment of synthetic image quality

The fine-tuned stable diffusion model can generate CXR images with severe bilateral edema patterns invoked by the controlled text prompt “bilateral lung edema mRALE 24 [pattern identifier] chest x-ray [class identifier]”, where the subject identifier “bilateral lung edema mRALE 24” is to control the desired patterns synthesis and the prior class identifier “chest x-ray” is to improve image diversity and to reduce overfitting. The synthesized image samples are shown in Figure 3, where the three CXR image labeled as “clear” on the upper row are real non-severe CXR images, and the other three labeled as “edema” on the lower row are synthetic CXR images with severe lung edema.

Figure 3.

Comparison of real (clear) and synthetic (edema) CXR images

We respectively trained a Wasserstein GAN with gradient penalty (WGAN-GP) and a denoising diffusion implicit model (DDIM) to synthesize the CXR images with lung edema pattern for comparison.

We computed the Frechet inception distance (FID) and Kernel inception distance (KID), considered more reliable and suitable for small dataset [14], between the synthetic images and the real severe lung edema CXR images. They provide the quantitative measurement for the fidelity and diversity of synthetic images and lower scores imply better image quality. The results are presented in Table I. The result indicates the synthetic image quality by the stable diffusion model fine-tuned with prior preservation outperforms all the comparators.

Table I.

Comparison of image quality (Lower score is better)

Model	FID	KID
Stable Diffusion with prior preservation	9.2158	0.0818
Stable Diffusion without prior preservation	20.1367	0.1135
DDIM	75.2278	0.2428
WGAN-GP	77.3235	0.2254

4.3. Classification accuracy

We also fine-tuned a pretrained vision transformer (ViT) with the whole MIDRC CXR dataset to measure the accuracy of the synthetic images synthesized by the proposed stable diffusion model carrying the desired visual patterns.

We respectively used 200 synthetic images as positive and 200 non-severe real images as negative for the classification test. The stable diffusion fine-tuned by prior preservation had high accuracy (accuracy of 0.9975, precision of 1.0, and recall of 0.995). Upon removing the class prior preservation loss, a large portion of synthetic images (16.5%) were classified as negative, thereby confirming successful binding with the knowledge representing by text.

5. CONCLUSION

We present a new method to fine-tune the state-of-the-art stable diffusion model to synthesize CXR images with pulmonary edema patterns using a medical knowledge driven text prompt. We show that the class prior preservation technique can effectively control overfitting in training and enrich image diversity. It also helps synthetize high-quality and high-resolution medical images with rare pathological patterns since normal medical images are more accessible compared to the images with disease.

Compared to DreamBooth fine-tuning method proposed by Google Research, our new method used strictly defined medical terms which were mapped to the targeted visual patterns on the CXR images. In our experiments, if the prior class identifier “chest x-ray” was removed, the synthetic images will look least authentic which is confirmed by the comparison in Table I. Similarly, if the subject identifier “bilateral lung edema mRALE 24” is omitted from the input text prompt, the synthetic images will be obvious artificial. These observations confirmed that well-defined key word control is an effective method to exert control not only to the semantic meaning of the input text prompt, but also to the visual pattern generated by the fine-tuned stable diffusion model. This new finding inspires a new approach to use stable diffusion model to synthesize complex visual patterns of medical images.

Figure 1.

General architecture of diffusion model