clara_pt_covid19_3d_ct_classification

Model Overview

The model described in this card is used to classify the lung region from the 3D chest CT images to COVID and non-COVID. For a detailed description, please see Harmon et al. [1].

Note: The 4.1 version of this model is only compatible with the 4.1 version of the Clara Train SDK container

Model Architecture

The model is trained using a 3D version of MONAI's DenseNet121 network [2] with Clara Train SDK v4.1. The input images are cropped to the lung region based on a lung segmentation mask before being fed to the model. These preprocessing steps involved in training the model are listed in config_train.json. The user can optionally mask the region outside the lung by setting disable_mask_intensity=False (see "Masking of Lung Region"). Note: the released pre-trained model was trained with masking disabled (disable_mask_intensity=True).

Training

The training was performed with the following:

• Script: train.sh
• GPU: At least 12GB of GPU memory.
• Actual Model Input: 192 x 192 x 64
• AMP: True
• Optimizer: ADAM
• Learning Rate: 5e-6
• Loss: "BCEWithLogitsLoss"

Dataset

This model was trained and evaluated on a global dataset with thousands of experimental cohorts collected from across the globe. See [1] for details.

In total, 2724 scans from 2617 patients were used in this study, including 1029 scans of 922 patients with RT-CPR confirmed COVID-19 and lung lesions related to COVID-19 pneumonia.

Of these, 1387 scans from 1280 patients were utilized for algorithm development, and 1337 patients were utilized for algorithm testing and evaluation.

Dataset JSON Path

The dataset.json needs to provide the path to both the image and lung segmentation mask image and target label for training. For example:

  "training": [
        {
            "image": "CT-Covid-19-August2020/volume-covid19-A-0000.nii.gz",
            "label_image": "CT-Covid-19-August2020/Mask/volume-covid19-A-0000/volume-covid19-A-0000_model.nii.gz",
            "label": [
                1
            ]
        },
        {
            "image": "LIDC-IDRI/LIDC-IDRI-0041_1.nii.gz",
            "label_image": "LIDC-IDRI/Mask/0041.nii.gz",
            "label": [
                0
            ]
        },
        …

Performance

This classification model achieved an accuracy of 91.4% using the same test set and lung segmentation masks of [1] ("Original training schema") which consisted of more than one thousand CT images collected across the globe. The model achieved an AUC value of 0.951 for predicting COVID positive cases with the provided MMAR configuration.

Training

A graph showing the training accuracy over 3k steps.

Validation

A graph showing the validation accuracy over 3k steps.

How to Use this Model

The model was validated with NVIDIA hardware and software. For hardware, the model can run on any NVIDIA GPU with memory greater than 12 GB. For software, this model is usable only as part of Transfer Learning & Annotation Tools in Clara Train SDK container. Find out more about Clara Train at the Clara Train Collections on NGC.

(Optional) Validate checksum of pretrained model ckpts

md5sum -c models/models.md5

Run docker:

For single GPU:

cd commands
./docker.sh 0

or for multi GPU:

cd commands
./docker.sh all

To train, run

cd commands; ./train.sh

To validate the checkpoint

cd commands; ./validate.sh

Full instructions for the training and validation workflow can be found in our documentation.

Input

Input: 1 channel CT image with intensity in HU and 1 lung segmentation mask.

Preprocessing:

1. "OrientationDict"
2. "ScaleIntensityRangeDict"
3. "MaskIntensityDict" (disabled by default)
4. "CropForegroundDict"
5. "ResizeDict"

Augmentation in training:

1. "RandShiftIntensityDict"
2. "RandAdjustContrastDict"
3. "RandGaussianNoiseDict"
4. "RandZoomDict"

The CT image will be cropped and resized to fit the model input based on provided lung segmentation. The lung segmentation image needs to be of the same size as the CT image. You can use clara_pt_covid19_ct_lung_segmentation to provide a segmentation mask (binary, 1 for lung, 0 for background) of each lung or provide your own lung segmentation, which is used for cropping the region to be classified.

Output

Output: 1 class probability for COVID

Masking of Lung Region

To replicate the training scheme involving in masking described in [1], one can mask (set to zero) the tissue outside the lung mask using disable_mask_intensity=False in train.sh or train_multi_gpu.sh. A model trained this way should be evaluated by setting "disable" at "MaskIntensityDict" to false in the validation and inference configurations.

Generate Saliency Map Visualization

To visualize the class activation maps (CAM) of the models using visualization techniques like GradCAM implemented in MONAI, use

cd commands; ./infer_saliency.sh

Note: You can choose from different CAM techniques by setting "cam_name" arg in "SaliencyInferer" section of config_saliency.json.

An example output of the saliency inference (*_saliency_inverted.nii.gz) overlaid on the original image is shown here:

Public Dataset Testing

Dataset

A subset of the COVID-positive cases used in [1] are available as "CT Images in COVID-19". One can furthermore use images from LIDC-IDRI as COVID-negative cases. The DICOM images from LIDC-IDRI must be converted NifTI format before training:

    medl-dataconvert -d ${SOURCE_IMAGE_ROOT} -s .dcm -e .nii.gz -o ${DESTINATION_IMAGE_ROOT}

We provide an example data list at config/example_data.json. This list is set referenced environment.json and utilized in the commands provided in this MMAR. The path in the datalist is relative to the DATA_ROOT also specified in environment.json.

Lung Mask Generation

To generate lung segmentation masks, please utilize NVIDIA's clara_pt_covid19_ct_lung_segmentation model.

Performance

On this example testing data, the pre-trained model achieved an accuracy of 87.7% and AUC of 0.987. The expected predictions and metrics using validate_ckpt.sh of our provided pre-trained model are given under docs/example_eval in metrics.csv or predictions.csv for training/validation.

Note: This example data is just given for illustrating the functions of this mmar. It is not supposed to provide a clinically useful training/testing dataset. Please refer to performance on the test set of [1] for a more clinically relevant performance measure.

Limitations

The example public dataset is given for illustrating the functions of this MMAR. It is not supposed to provide a clinically useful training or testing dataset. Please refer to performance on the test set of [1] for a more clinically relevant performance measure.

This training and inference pipeline was developed by NVIDIA. It is based on a segmentation and classification model developed by NVIDIA researchers in conjunction with the NIH. The Software is for Research Use Only. Software’s recommendation should not be solely or primarily relied upon to diagnose or treat COVID-19 by a Healthcare Professional.

This research use only software has not been cleared or approved by FDA or any regulatory agency. Clara pre-trained models are for developmental purposes only and cannot be used directly for clinical procedures.

References

[1] Harmon, Stephanie A., et al. "Artificial intelligence for the detection of COVID-19 pneumonia on chest CT using multinational datasets." Nature communications 11.1 (2020): 1-7. https://www.nature.com/articles/s41467-020-17971-2

[2] Huang, Gao, et al. "Densely connected convolutional networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017. https://arxiv.org/abs/1608.06993.

[3] https://developer.nvidia.com/clara-medical-imaging

License

End User License Agreement is included with the product. Licenses are also available along with the model application zip file. By pulling and using the Clara Train SDK container and downloading models, you accept the terms and conditions of these licenses.