This repository contains code for reproducing the following paper:
We also include models from the MoCo pretraining process for groups interested in fine-tuning them on their own data. Prior to using this code or pretrained models please consult the Disclaimer.
First, follow the
official instructions for
installing PyTorch. Then, navigate to the root CovidPrognosis
directory and
run
pip install -e .
After that you should be able to run the examples in cp_examples
.
For pretraining, you'll need to download the
MIMIC-CXR-JPG or
CheXpert datasets.
Once you've downloaded the data, add the path to configs/data.yaml
and it
should be used as a default.
The cp_examples
directory contains three subdirectories corresponding to the
training stages in the paper:
moco_pretrain
: Momentum-contrast (MoCo) pretraining (e.g., with MIMIC, CheXpert, or both)sip_finetune
: Fine-tuning of MoCo models for single-image prediction tasks (i.e., single-image adverse event prediction or oxygen requirements prediction)mip_finetune
: Fine-tuning of MoCo models for multi-image prediction tasks
Our code is built on top of the PyTorch Lightning framework.
The examples scripts for MoCo pretraining and SIP fine-tuning are set up for public X-ray data sets - due to consideration of patient privacy we do not release COVID data and use the public data as examples.
We provide pretrained models that use the MIMIC-CXR-JPG and CheXpert datasets. These datasets are subject to additional terms and conditions as indicated on their respective websites. For a list of models, see here. For an example of how to download and train with the models, please look at the SIP Fine-tuning example.
The following publication describes MIMIC-CXR:
And this publication describes CheXpert:
This code and accompanying pretrained models are provided with no guarantees regarding their reliability, accuracy or suitability for any particular application and should be used for research purposes only. The models and code are not to be used for public health decisions or responses, or for any clinical application or as a substitute for medical advice or guidance.
If you use this code or models in your scientific work, please cite the following paper:
@misc{sriram2021covid19,
title={COVID-19 Deterioration Prediction via Self-Supervised Representation Learning and Multi-Image Prediction},
author={Anuroop Sriram and Matthew Muckley and Koustuv Sinha and Farah Shamout and Joelle Pineau and Krzysztof J. Geras and Lea Azour and Yindalon Aphinyanaphongs and Nafissa Yakubova and William Moore},
year={2021},
eprint={2101.04909},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
The code is MIT licensed, as found in the LICENSE file.