Adapted generative latent diffusion models for accurate pathological analysis in chest X-ray images

Abstract

Respiratory diseases have a significant global impact, and assessing these conditions is crucial for improving patient outcomes. Chest X-ray is widely used for diagnosis, but expert evaluation can be challenging. Automatic computer-aided diagnosis methods can provide support for clinicians in these tasks. Deep learning has emerged as a set of algorithms with exceptional potential in such tasks. However, these algorithms require a vast amount of data, often scarce in medical imaging domains. In this work, a new data augmentation methodology based on adapted generative latent diffusion models is proposed to improve the performance of an automatic pathological screening in two high-impact scenarios: tuberculosis and lung nodules. The methodology is evaluated using three publicly available datasets, representative of real-world settings. An ablation study obtained the highest-performing image generation model configuration regarding the number of training steps. The results demonstrate that the novel set of generated images can improve the performance of the screening of these two highly relevant pathologies, obtaining an accuracy of 97.09%, 92.14% in each dataset of tuberculosis screening, respectively, and 82.19% in lung nodules. The proposal notably improves on previous image generation methods for data augmentation, highlighting the importance of the contribution in these critical public health challenges.

Keywords: Chest X-ray; Deep learning; Lung nodules; Stable diffusion; Tuberculosis.

Fig. 1

Representative examples of the used datasets of this work. The first column represents a normal case, while the second column represents a pathological case (tuberculosis or lung nodules, depending on the dataset). First row: examples of the Montgomery dataset (tuberculosis). Second row: examples of the Shenzhen dataset (tuberculosis). Third row: examples of the JSRT dataset (lung nodules)

Fig. 2

Detail of the methodology proposed in this work with 2 different steps (image generation and Pathological screening)

Fig. 3

Structure of the contrastive unpaired translation model, showing the patches that are compared to compute the loss

Fig. 4

Architecture of the adapted Stable Diffusion model following the DreamBooth training framework (keeping the U-Net unfrozen but the text encoder and the VAE frozen). It must be noted that the U-Net input vector is different in training and testing

Fig. 5

Detailed description of the classification deep network architecture used for the screening task. The intermediate output sizes are shown at the top of each layer

Fig. 6

Evolution of the F1-Score for the tuberculosis screening using the Montgomery dataset when adding the images generated by the Stable Diffusion model to the baseline, when trained during 2500, 5000, 7500 and 10,000 steps, respectively

Fig. 7

Evolution of the automatic tuberculosis screening performance in terms of the F1-Score when adding the generated images to the Shenzhen dataset with the different chosen configurations of Stable Diffusion (with 2500, 5000, 7500 and 10,000 training steps)

Fig. 8

Evolution of the F1-Score for the lung nodule screening (JSRT dataset) when adding the novel set of generated images regarding the different configurations of the Stable Diffusion model (with 2500, 5000, 7500, 10,000 and 12,500 training steps)

Fig. 9

Representative examples of pathological images generated from normal cases by the Stable Diffusion (SD) model. Each column presents the examples of the Montgomery County dataset, Shenzhen and JSRT, respectively. First row: original normal image. Second row: pathological image generated by SD trained 5000 steps. Third row: pathological image generated by SD trained 7500 steps. Fourth row: pathological image generated by SD trained 10,000 steps

Fig. 10

Representative examples of normal images generated from pathological cases by the Stable Diffusion (SD) model. Each column shows the examples of the Montgomery County, Shenzhen and JSRT datasets, respectively. First row: original pathological image. Second row: normal image generated by SD trained 5000 steps. Third row: normal image generated by SD trained 7500 steps. Fourth row: normal image generated by SD trained 10,000 steps

Fig. 11

Examples of low quality images generated by Stable Diffusion. (a) Image generated by a Stable Diffusion model that was trained a small amount of training steps. (b) Image generated by a Stable Diffusion model that was trained during a too high number of training steps. (c) Generated image with lack of anatomical coherence, which displays 4 clavicles