Counterfactual Disease Removal and Generation in Chest X-Rays Using Diffusion Models

Ahmet Berke Gökmen, Ender Konukoglu
ETH Zurich Computer Vision Lab
Work completed during an internship at Computer Vision Lab.
Paper Code
Teaser

Our Method is Able to Generate Clinically Accurate Counterfactual Images While Preserving Patient Identity.


Abstract

Recent advancements in medical imaging have leveraged deep learning, particularly latent manipulation in autoencoders and StyleGAN-based models, to edit chest X-rays for disease modification. However, these methods often rely on fixed latent directions, limiting their flexibility and precision. Although diffusion models have recently been explored for medical imaging, they frequently alter critical patient-specific details, compromising identity preservation. We introduce a novel diffusion-based approach that enables disease generation and removal in chest X-rays while preserving patient identity, eliminating the need for specific disease masks. By fine-tuning attention layers, our model retains defining anatomical features, ensuring realistic yet flexible edits. A pretrained classifier guides the diffusion process, aligning modifications with clinical relevance. We evaluate our approach using Fréchet Inception Distance (FID) for image quality and Contrastive Language-Image Pretraining (CLIP) accuracy.

Method

Fine Tuning Diffusion Models

Method

The process begins with DDIM Inversion, which reconstructs the input image in the latent space. A Bone Segmentor generates a bone mask to isolate relevant structures for manipulation. At each diffusion step, Bone Masking combines the original and manipulated latent states, ensuring that the anatomical integrity of the bones is preserved. The Attention Fine-Tuning module iteratively refines the outputs with a classifier-guided approach, minimizing reconstruction errors using an L1 loss. This ensures accurate and identity-preserving manipulation, as highlighted by the feedback loop of loss backpropagation.




Qualitative Results of Disease Removal

First rows show the original image while second rows show the counterfactual one. Click on titles to expand images.


Pleural Effusion Removal

Cardiomegaly Removal

Atelectasis Removal

Qualitative Results of Disease Synthesis

First rows show the original image while second rows show the counterfactual one. Click on titles to expand images.


Pleural Effusion Synthesis

Cardiomegaly Synthesis

Atelectasis Synthesis

BibTeX

@misc{,
      title={Counterfactual Disease Removal and Generation in Chest X-Rays Using Diffusion Models}, 
      author={Ahmet Berke Gokmen and Ender Konukoglu},
      year={2024},
      eprint={},
      archivePrefix={},
      primaryClass={cs.CV},
      url={}, 
}