Weakly Supervised Contrastive Learning for Chest X-Ray Report Generation

This is the pytorch implementation for our paper:

Weakly Supervised Contrastive Learning for Chest X-Ray Report Generation

at Findings of EMNLP-2021.

Requirements

• torch>=1.6.0
• torchvision>=0.8.0

Datasets

We use two datasets (MIMIC-ABN and MIMIC-CXR) in the paper.

For MIMIC-ABN, you can download the dataset from release/mimic_abn and then put the files in data/mimic_abn.

For MIMIC-CXR, you can download the dataset from release/mimic_cxr and then put the files in data/mimic_cxr.

Note: you need to sign user agreements then download x-ray images from the official website.

Run on MIMIC-ABN

Run bash run_mimic_abn.sh to train a model on the MIMIC-ABN data.

Run on MIMIC-CXR

Run bash run_mimic_cxr.sh to train a model on the MIMIC-CXR data.

Evaluation

We use public code sources to evaluate our models.

For CE metrics, please follow this repo to use ChexPert and install relative packages. Then refer to label_on_fly.py and run_label.sh in our chexpert-labeler folder to label your reports (which will generate a csv file with ChexPert labels). In the end, use calculate_metric.py to compute clinical accuracy.

For NLG metrics, please refer to pycocoevalcap.

Download Models

You can download the models we trained for each dataset from release/pretrained_models.

Citation

If you find this repository useful, please cite our paper:

@article{yan2021weakly,
  title={Weakly Supervised Contrastive Learning for Chest X-Ray Report Generation},
  author={Yan, An and He, Zexue and Lu, Xing and Du, Jiang and Chang, Eric and Gentili, Amilcare and McAuley, Julian and Hsu, Chun-Nan},
  journal={arXiv preprint arXiv:2109.12242},
  year={2021}
}

Acknowledgments

This project is built on top of R2Gen. Thank the authors for their contributions to the community!