COVID-19 Detection Neural Network (COVNet)

This is a PyTorch implementation of the paper "Artificial Intelligence Distinguishes COVID-19 from Community Acquired Pneumonia on Chest CT". It supports training, validation and testing for COVNet.

Citation

If you find this code is useful for your research, please consider citing:

@article{zhang2020artificial,
  title={Artificial Intelligence Distinguishes COVID-19 from Community Acquired Pneumonia on Chest CT},
  author={Li, Lin and Qin, Lixin and Xu, Zeguo and Yin, Youbing and Wang, Xin and Kong, Bin and Bai, Junjie and Lu, Yi and Fang, Zhenghan and Song, Qi and Cao, Kunlin and others},
  journal={Radiology},
  year={2020}
}

Setup

Prerequisites

Anaconda 3.7
PyTorch 1.4
SimpleITK
batchgenerators
tensorboardX

Prepare data

Preprocess the data according to the Appendix E1 section in the paper and organize them as the following. A example of train.csv and val.csv are also provided.

data
├── caseid1
|   ├── masked_ct.nii
|   └── mask.nii.gz
├── caseid2
|   ├── masked_ct.nii
|   └── mask.nii.gz
├── caseid3
|   ├── masked_ct.nii
|   └── mask.nii.gz
├── caseid4
|   ├── masked_ct.nii
|   └── mask.nii.gz
├── train.csv
└── val.csv

COVNet

Training

Training a COVNet with default arguments. Model checkpoints and tensorboard logs are written out to a unique directory created by default within experiments/models and experiments/logs respectively after starting training.

python main.py

Validation and Testing

You can run validation and testing on the checkpointed best model by:

python test.py

utayao/COVNet