COVID-19 Lung Lesion Segmentation Using a Sparsely Supervised Mask R-CNN on Chest X-rays Automatically Computed from Volumetric CTs

 Vignav Ramesh,  Blaine Rister,  Daniel L. Rubin

arXiv / Full Paper (PDF) / Papers With Code / Mask R-CNN Code

Pre-trained models

Training dataset	Train/test split	Data augmentation (y/n)	URL
X-rays Only	60/40	y	Download
Mixed	60/40	y	Download
X-rays Only	80/20	y	Download
Mixed	80/20	y	Download
X-rays Only	80/20	n	Download
Mixed	80/20	n	Download

Environment setup

Our models were trained on a single GPU (Tesla P4 GPU provided by Google Colab, 16 GB memory). The code is implemented using TensorFlow v1, but is compatible with TensorFlow v2 and can be ported to the most recent version of TensorFlow if desired. To install all required dependencies, run the following:

pip install -r requirements.txt

Afterwards, set up the Mask R-CNN model:

git clone --quiet https://github.com/rvignav/Mask_RCNN.git
cd ~/Mask_RCNN
pip install -q PyDrive
pip install -r requirements.txt
python setup.py install
cp ~/Mask_RCNN/samples/balloon/balloon.py ./lesion.py
sed -i -- 's/balloon/lesion/g' lesion.py
sed -i -- 's/Balloon/Lesion/g' lesion.py

Data

Dataset	URL
Training Dataset 1 (X-rays Only)	Download
Training Dataset 2 (Mixed)	Download
Test Dataset	Download

Training

The following commands can be used to train the Mask R-CNN model:

# Train a new model starting from pre-trained ImageNet weights
python lesion.py train --dataset='/path/to/data/' --weights=imagenet

# Train a new model starting from pre-trained COCO weights
python lesion.py train --dataset='/path/to/data/' --weights=coco

# Continue training a model that you had trained earlier
python lesion.py train --dataset='/path/to/data/' --weights=/path/to/weights/

# Continue training the last model you trained. This will find
# the last trained weights in the model directory.
python lesion.py train --dataset='/path/to/data/' --weights=last

To train with data augmentation, run:

python lesion.py train --dataset='/path/to/data/' --weights=imagenet/coco/last --aug='y'

Alternatively, use the following Colabs to train the models:

X-rays only:

Mixed:

The CT to X-ray re-projection algorithm can be executed in isolation as follows:

python ct2xray.py <path/to/CT/volume> <path/to/mask/volume>

Cite

@misc{ramesh2021covid19,
      title={COVID-19 Lung Lesion Segmentation Using a Sparsely Supervised Mask R-CNN on Chest X-rays Automatically Computed from Volumetric CTs}, 
      author={Vignav Ramesh and Blaine Rister and Daniel L. Rubin},
      year={2021},
      eprint={2105.08147},
      archivePrefix={arXiv},
      primaryClass={eess.IV}
}