Awesome
MGCA
Multi-Granularity Cross-modal Alignment for Generalized Medical Visual Representation Learning, NeurIPS 2022.
Installation
To clone this repository:
git clone https://github.com/fuying-wang/MGCA.git
To install Python dependencies:
pip install -r requirements.txt
To install package mgca
:
pip install -e .
Dataset downloading
Datasets we used are as follows:
-
MIMIC-CXR: We downloaded the MIMIC-CXR-JPG dataset as the radiographs. Paired medical reports can be downloaded in MIMIC-CXR.
-
CheXpert: We downloaded the CheXpert dataset which consisting of 224,316 chest radiographs of 65,240 patients.
-
RSNA: We used the stage 2 of RSNA dataset in Kaggle.
-
COVIDx: We used the version 6 of COVIDx dataset in Kaggle.
-
SIIM: We downloaded the stage 1 of SIIM dataset in Kaggle.
-
Object-CXR: We downloaded the object-CXR dataset in its official website.
After downloading datasets, please check if the path in mgca/constants.py
is correct.
Data Preprocessing
We preprocessed these datasets and split the dataset into train/val/test set using the code in mgca/preprocess
.
Pre-training
Reminder: Please check Line 47 of mgca/datasets/pretrain_dataset.py
, Line 503 of mgca_module.py
and make sure the path is correct.
We pre-trained MGCA on MIMIC-CXR using this command:
cd mgca/models/mgca
CUDA_VISIBLE_DEVICES=0,1 python mgca_module.py --gpus 2 --strategy ddp
We train our framework 50 epochs on 2 pieces of RTX 3090 GPUs with batch size of 144. It takes about 1 day to pre-train this model.
Note that it is flexible to develop other pre-training models under this framework. You may create a folder in mgca/models
and complete the {MODEL_NAME}_module.py
file.
Pre-trained models can be found here.
Finetune on downstream tasks
We evlauate the performance of MGCA framework on three downstream tasks: classification, object detection and semantic segmentation. Before finetuning, we need set the path
(or ckpt_path
) argument to the path of pre-trained MGCA model.
Linear classification
We evaluate linear classification performance of our model using this command:
cd mgca/models/mgca
CUDA_VISIBLE_DEVICES=1 python mgca_finetuner.py --gpus 1 --dataset chexpert --data_pct 0.01
We can use --dataset
to set specific dataset for finetuning. Here, 3 datsets are available: chexpert, rsna and covidx.
We can use --data_pct
to set the fraction of training data for finetuning.
Object detection
We evaluate object detection performance of our model using this command:
cd mgca/models/mgca
CUDA_VISIBLE_DEVICES=0 python mgca_detector.py --devices 1 --dataset rsna --data_pct 1 --learning_rate 5e-4
Here, 2 datsets are available: rsna and object_cxr.
To run all experiments for this detection task:
sh run_det_funetune.sh
Semantic segmentation
We evaluate semantic segmentation performance of our model using this command:
cd mgca/models/mgca
CUDA_VISIBLE_DEVICES=0 python mgca_segmenter.py --gpus 1 --data_pct 1 --dataset rsna --batch_size 16 --learning_rate 5e-4
Here, 2 datsets are available: rsna and siim.
To run all experiments for this detection task:
sh run_seg_funetune.sh
TODO List
- Refactor the train/valid/test lists.
Reference
If you found our work useful in your research, please consider citing our works(s) at:
@article{wang2022multi,
title={Multi-granularity cross-modal alignment for generalized medical visual representation learning},
author={Wang, Fuying and Zhou, Yuyin and Wang, Shujun and Vardhanabhuti, Varut and Yu, Lequan},
journal={Advances in Neural Information Processing Systems},
volume={35},
pages={33536--33549},
year={2022}
}