Multi-Feature Vision Transformer via Self-Supervised Learning for COVID-19 Diagnosis
Introduction
This is the code to repoduce the study of [Multi-Feature Vision Transformer via Self-Supervised Representation Learning for Improvement of COVID-19 Diagnosis](arxiv update soon). Please cite our study if you are using this dataset or referring to our method.
Network Structure # Multi-Feature-Semi-Supervised-Learning_COVID-19 (Pytorch)
Result
- Test-1
Method | Labeled Sample (%) | Precision | Recall | F1-Scores | Top-1(%) |
---|---|---|---|---|---|
MF-ViT CA | 10 | 0.91 | 0.91 | 0.91 | 91.10 |
MF-ViT CA | 30 | 0.93 | 0.93 | 0.93 | 93.27 |
MF-ViT CA | 100 | 0.95 | 0.95 | 0.95 | 95.03 |
Usage
-
Dataset and Trained model weights:
- Download them from Kaggle. CXR folder are all origianl CXR images and Enh folder are all corresponding enhanced images.
-
Preparation:
- Create a folder named MoCo-COVID to save all downloaded files from this repo and files from Kaggle in one folder.
-
Train MoCo-COVID:
- python MoCo-COVID/moco_pretraining/moco/main_covid_mocov3based_single_img_type_5draws_mocov3structure_mocov2loss_vitsmall.py -a vit_small -b 16 --optimizer=adamw --lr=1.5e-4 --weight-decay=.1 --epochs=30 --warmup-epochs=4 --stop-grad-conv1 --moco-m-cos --moco-t=.2 --multiprocessing-distributed --world-size 1 --rank 0 --aug-setting chexpert --rotate 10 --exp-name --train_data data --cos (train_data: data=original CXR image; Train_Mix=enhanced CXR image)
-
Finetune MoCo-COVID-LP:
- python MoCo-COVID/moco_pretraining/moco/main_vit_covid_test_val_single_img_type_5draws_rev_v2loss_v3structure_vitsmall.py -a vit_small --lr 3 --batch-size 16 --epochs 90 --exp-name --pretrained --maintain-ratio --rotate --aug-setting chexpert --train_data data --optimizer sgd --cos (--pretraind = pretrained weights)
-
Finetune MoCo-COVID-FT:
- python MoCo-COVID/moco_pretraining/moco/main_vit_covid_test_val_single_img_type_5draws_rev_v2loss_v3structure_vitsmall.py -a vit_small --lr 3 --batch-size 16 --epochs 90 --exp-name --pretrained --maintain-ratio --rotate --aug-setting chexpert --train_data data --optimizer sgd --cos --semi-supervised (--pretraind = pretrained weights)
-
Finetuen MF-ViT CA
- python MoCo-COVID/main_vit_covid_test_val_single_img_type_5draws_rev_v2loss_v3structure_crossvit_2vits_2additionaloutputs_trainval_sum.py -a vit_small --batch-size 32 --exp-name --pretrained --pretrained_enh --maintain-ratio --rotate --aug-setting chexpert --lr 1.5e-4 --cos --ep 25