RP3D-Diag

Nature Communications | The official codes implementation for paper "Large-scale Long-tailed Disease Diagnosis on Radiology Images"

ArXiv Version

In this paper, we build up an academically accessible, large-scale diagnostic dataset, present a knowledge enhanced model architecture that enables to process arbitrary number of input scans from various imaging modalities, and initialize a new benchmark for multi-modal multi-anatomy long-tailed diagnosis. Our method shows superior results on it. Additionally, our final model serves as a pre-trained model, and can be finetuned to benefit diagnosis on various external datasets.

Dataset

Overview of RP3D-DiagDS. There are 40,936 cases (195,010 scans) across 7 human anatomy regions and 9 diverse modalities covering 5569 disorders mapped into 931 ICD-10-CM codes.

Model

The overview of our model RadDiag. Three parts demonstrate our proposed visual encoders and fusion module, together with the knowledge enhancement strategy respectively. a, The three types of vision encoder, i.e., ResNet-based, ViT-based, and ResNet-ViT-mixing. b, The architecture of the fusion module. The figure shows the transformer-based fusion module, enabling case-level information fusion. c, The knowledge enhancement strategy. We first pre-train a text encoder with extra medical knowledge with contrastive learning, leveraging synonyms, descriptions, and hierarchy. Then we view the text embedding as a natural classifier to guide the diagnosis classification.

The model checkpoint can also be found in HuggingFace repository RadDiag.

Setup

Environment

To Install the python environments:

pip install -r requirements.txt

Data Preparation

• Download files from HuggingFace.
• Modify all the paths (files in DataPath, all the python files and bash files) to your own.
• The DataPath folder should be like:

src
|   DataPath
|   |    train.json
|   |    test.json
|   |    aug.json
|   |    disorder_label_dict.json
|   |    icd10_label_dict.json
|   |    xxxCheckpoint
|   |    |    pytorch.model.bin 
|   |    |    ......

Quickstart Demo

Here we use a demo to show in detail the directory structure and how to run the model for inference, as an example for reproducing the subsequent experiments.

To run this demo, please refer to RP3D_Demo_Instruction.md

Eval

1. cd .../src/Eval
2. add data file to ./DataPath
3. replace the relative path in ./Model/ and ./eval.py, ./eval.sh, etc with your absolute path.
4. set checkpoint in ./eval.sh
5. Run in terminal: bash eval.sh

For more info, please refer to readme files in src/Eval/

Train

From Scratch

1. cd .../src/Train
2. add data file to ./DataPath/
3. replace the relative path in ./Model/ and ./train.py, ./train.sh, etc with your absolute path.
4. Run in terminal: bash train.sh

Load Checkpoint

1. cd .../src/Train
2. add data file to ./DataPath/
3. replace the relative path in ./Model/ and ./train.py, ./train.sh, etc with your absolute path.
4. set checkpoint in ./train.sh
5. Run in terminal: bash train.sh

For more info, please refer to readme files in src/Train/

Finetune

1. cd .../src/Finetune
2. add data file to ./DataPath/
3. replace the relative path in ./Model/ and ./train.py, ./train.sh, etc with your absolute path.
4. set checkpoint in ./train.sh
5. Run in terminal: bash train.sh

For more info, please refer to readme files in src/Finetune/

Zeroshot

1. cd .../src/Zeroshot
2. add data file to ./DataPath/
3. replace the relative path in ./Model/ and ./predict.py, ./predict.sh, eval.py, etc with your absolute path.
4. set checkpoint in ./eval.sh
5. run in terminal: bash predict.sh
6. run in terminal: python eval.py

For more info, please refer to readme files in src/Zeroshot/

Benchmark

Classification results on Disorders and ICD-10-CM levels. In this table Fusion Module and Knowledge Enhancement strategy are both used. We report the results on Head/Medium/Tail class sets separately.

ROC curves on Disorders and ICD-10-CM, including head/medium/tail parts respectively. The shadow in the figure shows the 95\% CI~(Confidence interval) and FM, KE are short for Fusion Module and Knowledge Enhancement.

Comparison

The AUC Score Comparison on Various External Datasets. SOTA denotes the best performance of former works (pointed with the corresponding reference) on the datasets. Scratch means using our model but training from scratch. Ours means to use our checkpoint to fintune.

In zero-shot assessment, we evaluate the transferring ability of our final model on normal/abnormal diagnosis for the external dataset.

Acknowledgment

We sincerely thank all the contributors who uploaded the relevant data in our dataset online. We appreciate their willingness to make these valuable cases publicly available.

Citation

@article{zheng2023large,
          title={Large-scale Long-tailed Disease Diagnosis on Radiology Images},
          author={Zheng, Qiaoyu and Zhao, Weike and Wu, Chaoyi and Zhang, Xiaoman and Zhang, 
            Ya and Wang, Yanfeng and Xie, Weidi},
          journal={arXiv preprint arXiv:2312.16151},
          year={2023}
}