/PairAug

[CVPR2024] PairAug: What Can Augmented Image-Text Pairs Do for Radiology?

Primary LanguagePython

PairAug

This repo holds the Pytorch implementation of PairAug:

[CVPR2024] PairAug: What Can Augmented Image-Text Pairs Do for Radiology?

Usage

  • Create a new conda environment
conda create --name pairaug python=3.10
source activate pairaug
  • Clone this repo
git clone https://github.com/YtongXie/PairAug.git
cd PairAug

Medical image-report pair generation

0. Installation

  • Install packages for image generation
pip install -r requirements_T2I.txt

1. Data Preparation

  • Download MIMIC-CXR-JPG dataset (Need to be a credentialed user for downloading)

  • Put the image data under data/MIMIC_images_ori/ should be like:

    ├── p19_p19995997_s50123635_6fa953ea-79c237a5-4ca3be78-e3ae6427-c327e17b.png
    ├── p19_p19996061_s58482960_87923de8-5595ad44-eaa89d38-610e97e2-42cacf04.png
    ├── ...
  • Run python data/reports_jsonl.py to generate the report data lists data/metadata_train.jsonl.
  • Put the report data under data/MIMIC_reports_ori/ should be like:
    ├── p19_p19995997_s50123635_6fa953ea-79c237a5-4ca3be78-e3ae6427-c327e17b.txt
    ├── p19_p19996061_s58482960_87923de8-5595ad44-eaa89d38-610e97e2-42cacf04.txt
    ├── ...

2. InterAug generation

  • Download pretrained RoentGen weights and put it under pretrained/. Need to contact authors of paper RoentGen for download.
  • Run
  python InterAug_Step1.py --start_index 0 --end_index 2200000

It aims to generate new patient reports via ChatGPT, and save them in data/InterAug/InterAug_reports.

  • Run
  python InterAug_Step2_T2I.py --start_index 0 --end_index 2200000

It aims to generate inter-patient images via RoentGen model, and save them in data/InterAug/InterAug_images.

3. IntraAug generation

  • Run
python IntraAug_Step1.py --start_index 0 --end_index 2200000

It aims to generate intra-patient reports via ChatGPT, and save them in data/IntraAug/IntraAug_reports.

  • Run
python IntraAug_Step2_T2I.py --start_index 0 --end_index 2200000

It aims to generate intra-patient images based on generated and original reports, and save them in data/IntraAug/IntraAug_images.

Data Pruning

0. Installation

  • Install packages for data pruning
pip install -r requirements_MedClip.txt

1. Data Pruning for InterAug pairs

  • Run
python InterAug_Step3_fliter.py

It introduces a data pruning method w.r.t. the semantic alignment between generated InterAug image-report pairs to ensure the quality of the generated pairs.

2. Data Pruning for IntraAug pairs

  • Run
python IntraAug_Step3_fliter.py

It introduces a hybrid consistency score between generated IntraAug image-report pairs to ensure the quality of the generated pairs.

Acknowledgements

Thanks to diffusers for the latent diffusion model, prompt-to-prompt for prompt-to-prompt stable diffusion model, RoentGen for pretrained weights, CheXzero for medical visual-language pre-training.

Contact

Yutong Xie (yutong.xie678@gmail.com)