Cascaded Latent Diffusion for High-Resolution Chest X-ray Synthesis

This repository contains code for running and training Cheff - a cascaded chest X-ray latent diffusion pipeline. The cheff pipeline consists of three cascading phases:

1. Modeling a diffusion process in latent space
2. Translating the latent variables into image space with a decoder
3. Refinement and upscaling using a super-resolution diffusion process

Phase 1 and 2 together define an LDM. Phase 2 and 3 are trained on MaCheX, a collection of over 650,000 chest X-rays and thus, build a foundational basis for our model stack. The first phase is task-specific. For unconditional snythesis, we train on full MaCheX and for report-to-chest-X-ray we use the MIMIC subset.

How to use Cheff?

Please have a look into our tutorial notebook.

Models

We provide the weights for 5 models:

• Chest X-ray autoencoder: Click!
• Chest X-ray unconditioned semantic diffusion model: Click!
• Chest X-ray report-conditioned semantic diffusion model: Click!
• Chest X-ray super-resolution diffusion model base: Click!
• Chest X-ray super-resolution diffusion model finetuned: Click!

The tutorial notebook assumes that downloaded models are placed in trained_models.

Training

Our codebase builds heavily on the classic LDM repository. Thus, we share the same interface with a few adaptions. A conda environment file for installing necessary dependencies is environment.yml. For a pip-only install use requirements.txt. The full config files are located in configs. After adjusting the paths, the training can be started as follows:

python scripts/01_train_ldm.py -b <path/to/config.yml> -t --no-test

Training the Super-Resolution Model

The training procedure for reproducing CheffSR is located in an extra repository. You will find a script that contains the necessary configuration and routine.

Acknowledgements

This code builds heavily on the implementation of LDMs and DDPMs from CompVis: Repository here.

Citation

If you use Cheff or our repository in your research, please cite our paper Cascaded Latent Diffusion Models for High-Resolution Chest X-ray Synthesis:

@inproceedings{weber2023cascaded,
  title={Cascaded Latent Diffusion Models for High-Resolution Chest X-ray Synthesis},
  author={Weber, Tobias and Ingrisch, Michael and Bischl, Bernd and R{\"u}gamer, David},
  booktitle={Advances in Knowledge Discovery and Data Mining: 27th Pacific-Asia Conference, PAKDD 2023},
  year={2023},
  organization={Springer}
}