Multi-site COVID-Net CT Classification

This is the PyTorch implemention of our paper Contrastive Cross-site Learning with Redesigned Net for COVID-19 CT Classification by Zhao Wang, Quande Liu, Qi Dou

Abatract

This paper proposes a novel joint learning framework to perform accurate COVID-19 identification by effectively learning with heterogeneous datasets with distribution descrepancy. We build a powerful backbone by redesigning the recently proposed COVID-Net in aspects of network architecture and learning strategy to improve the prediction accuracy and learning efficiency. On top of it, we further explicitly tackle the cross-site domain shift by conducting separate feature normalization in latent space. Moreover, we propose a con-trastive training objective to enhance the domain invariance of semantic embeddings for boosting the classification performance on each dataset. We develop and evaluate our method with two public large-scale COVID-19 diagnosis datasets from real CT images. Extensive experiments show that our approach consistently improves the performances on both datasets, as well as outperforms existing state-of-the-art multi-site learning methods.

Usage

Setup

We suggest using Anaconda to setup environment on Linux, if you have installed anaconda, you can skip this step.

wget https://repo.anaconda.com/archive/Anaconda3-2020.11-Linux-x86_64.sh && zsh Anaconda3-2020.11-Linux-x86_64.sh

Then, we can install packages using provided environment.yaml.

git clone https://github.com/med-air/Contrastive-COVIDNet
cd Contrastive-COVIDNet
conda env create -f environment.yaml
conda activate pytorch0.4.1

Dataset

We employ two publicly available COVID-19 CT datasets:

• SARS-CoV-2 dataset
• COVID-CT dataset

Download our pre-processed datasets from Google Drive and put into data/ directory.

Pretrained Model

You can directly download our pretrained model from Google Drive and put into saved/ directory for testing.

Training

cd code
python main.py --bna True --bnd True --cosine True --cont True

Test

cd code
python test.py

Citation

If you find this code and dataset useful, please cite in your research papers.

@article{wangcontrastive,
   author={Wang, Zhao and Liu, Quande and Dou, Qi},
   title={Contrastive Cross-site Learning with Redesigned Net for COVID-19 CT Classification},
   journal={IEEE Journal of Biomedical and Health Informatics},
   DOI={10.1109/jbhi.2020.3023246},
   year={2020},
   volume={24},
   number={10},
   pages={2806-2813}
}

Questions

For further questions, pls feel free to contact Zhao Wang

References

[1] E. Soares, P. Angelov, S. Biaso, M. Higa Froes, and D. Kanda Abe, “Sars-cov-2 ct-scan dataset: A large dataset of real patients ct scans for sars-cov-2 identification,” medRxiv, 2020.

[2] J. Zhao, X. He, X. Yang, Y. Zhang, S. Zhang, and P. Xie, “Covid-ct-dataset: A ct scan dataset about covid-19,” 2020.

[3] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga et al., “Pytorch: An imperative style, high-performance deep learning library,” in Advances in neural information processing systems, 2019, pp. 8026–8037.