SLIViT: a general AI framework for clinical-feature diagnosis from limited 3D biomedical-imaging data
SLIViT is a data-efficient deep-learning framework that accurately measures disease-related risk factors in volumetric biomedical imaging, such as magnetic resonance imaging (MRI) scans, optical coherence tomography (OCT) scans, ultrasound videos, and Computed Tomography (CT) scans.
Below you may find step-by-step instructions on how to pre-train, fine-tune, and evaluate SLIViT. Please refer to our manuscript for further details and feel free to reach out regarding any concerns/issues you are experiencing with SLIViT.
Before using SLIViT, please make sure to have an appropriate conda environment with the relevant packages properly installed. Pre-trained and fine-tuned models can be downloaded from here. Our model is implemented using PyTorch, and this repository supports PyTorch GPU implementation.
Clone the repository and create conda enviroment:
git clone https://github.com/cozygene/SLIViT
conda create --name slivit python=3.8
Activate conda enviroment and install required packages:
conda activate slivit
conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia
cd SLIViT
pip install -r requirements.txt
Download the pre-trained backbone, fine-tuned SLIViT and move the files to the Checkpoints
folder.
![](Figs/Kermany_OCTs.png)
Download the Kermany dataset
python bb_train.py --dataset kermany --data_dir /path/to/data --meta_csv ./Dsets/kermany_meta.csv --pathologies Drusen,CNV,DME,Normal --out_dir /output/dir/to/save_pretrained_model/
![](Figs/ChestMNIST_Xrays.png)
python bb_train.py --dataset chestmnist --out_dir /output/dir/to/save_pretrained_model/
Supported 2D image formats : tiff
, jpeg
, png
, bmp
To commence pre-training with 2D Medical Scans, split the data into train
, validation
, test
sets. Create three
folders named train
, val
, test
, and move the scans accordingly.
Generate the meta.csv
file for 2D pre-training data as illustrated below:
F_Name | Path | Pathology-1 | Pathology-2 | Pathology-3 | Pathology-4 |
---|---|---|---|---|---|
CNV-6116901-21.jpeg | /train/CNV/CNV-6116901-21.jpeg | 1.0 | 0.0 | 0.0 | 0.0 |
DME-4616882-33.jpeg | /test/DME/DME-4616882-33.jpeg | 0.0 | 1.0 | 0.0 | 0.0 |
CNV-7907754-23.jpeg | /test/CNV/CNV-7907754-23.jpeg | 1.0 | 0.0 | 0.0 | 0.0 |
NORMAL-3757443-31.jpeg | /val/NORMAL/NORMAL-3757443-31.jpeg | 0.0 | 0.0 | 0.0 | 1.0 |
NORMAL-6434323-2.jpeg | /train/NORMAL/NORMAL-6434323-2.jpeg | 0.0 | 0.0 | 0.0 | 1.0 |
NORMAL-910422-8.jpeg | /val/NORMAL/NORMAL-910422-8.jpeg | 0.0 | 0.0 | 0.0 | 1.0 |
DRUSEN-8086850-28.jpeg | /train/DRUSEN/DRUSEN-8086850-28.jpeg | 0.0 | 0.0 | 1.0 | 0.0 |
In the above table, F_Name
denotes the file name for 2D medical scan files, Path
indicates the directory to these
files, and Pathology-1
, Pathology-2
, Pathology-3
, and Pathology-4
represent binary classes for the respective
pathologies.
After creating meta.csv
file, ConvNeXt backbone can be trained with following bash script:
python bb_train.py --dataset custom --meta_csv /path/to/meta.csv --pathologies Pathology-1,Pathology-2,Pathology-3,Pathology-4 --out_dir /output/dir/to/save_pretrained_model/ --b_size 16 --gpu_id 0 --n_cpu=32
--dataset
is the dataset for pre-training (kermany
,chestmnist
,custom
)--data_dir
is the root directory for the createdtrain
,validation
,test
set folders.--meta_csv
is the directory to the createdmeta.csv
.--pathologies
is a comma-separated list of pathologies for pre-training.--out_dir
is the output directory for saving the pre-trained backbone.--b_size
denotes the batch size for training.--gpu_id
specifies the GPU ID for training.--n_cpu
indicates the number of CPUs for data loading.--n_epochs
indicates the number of epochs to train.--patience
is the number of epochs to await without improvement before initiating early stopping
python slivit_train.py --dataset3d nodulemnist --task classification --bbpath ./Checkpoints/convnext_bb_kermany.pth --nObb_feat 4 --nslc 28 --depth 5 --dim 64 --heads 10 --out_dir /output/dir/to/save_finetuned_slivit/
The UKBB MRI dataset available here. After downloading the data, replace the paths in ./Dsets/ukbb_meta.csv
with corresponding paths to the downloaded scans.
python slivit_train.py --dataset3d ukbb --meta_csv ./Dsets/ukbb_meta.csv --task regression --bbpath ./Checkpoints/convnext_bb_kermany.pth --nObb_feat 4 --depth 5 --dim 256 --nslc 36 --heads 36 --out_dir /output/dir/to/save_finetuned_slivit/
The EchoNet Ultrasound videos are available here. After downloading the data, replace the paths in ./Dsets/ultrasound_meta.csv
with corresponding paths to the downloaded videos.
python slivit_train.py --dataset3d ultrasound --meta_csv ./Dsets/ultrasound_meta.csv --bbpath ./Checkpoints/convnext_bb_kermany.pth --nObb_feat 4 --nslc 32 --depth 5 --dim 256 --heads 32
Supported 3D volume formats : 2D slice files as tiff
or dcm
Generate the meta.csv
file for 3D fine-tuning data as illustrated below:
pid | path | Pathology | Split |
---|---|---|---|
1754715 | /scratch/avram/MRI/liver/dcm/with_annotations/5628597 | 1.4107 | train |
4501528 | /scratch/avram/MRI/liver/dcm/with_annotations/4501528 | 3.6272 | valid |
2161445 | /scratch/avram/MRI/liver/dcm/with_annotations/2161445 | 14.128 | valid |
5364532 | /scratch/avram/MRI/liver/dcm/with_annotations/5364532 | 1.3523 | train |
2653392 | /scratch/avram/MRI/liver/dcm/with_annotations/2653392 | 2.6253 | test |
1754715 | /scratch/avram/MRI/liver/dcm/with_annotations/1754715 | 5.2235 | train |
4986672 | /scratch/avram/MRI/liver/dcm/with_annotations/4986672 | 3.1356 | test |
In the above table, pid
denotes the unique patient ID for 3D medical volume, Path
indicates the directory to folder which contains 2D slice files for the correspoding volume, and Pathology
represent pathology score or binary pathology class.
After creating meta.csv
file, SLIViT can be fine tuned with following bash script:
python slivit_train.py --dataset3d custom --meta_csv /path/to/generated/meta.csv --bbpath /path/to/pretrained/convnext_bb.pth --task regression --pathology Pathology --out_dir /output/dir/to/save_pretrained_model/
--dataset3d
is the dataset for 3D fine-tuning (nodulemnist
,ukbb
,ultrasound
,custom3d
)--meta_csv
is the path to the createdmeta.csv
file--pathology
is pathology for 3D fine-tuning--bbpath
is the path to pre-rained ConvNeXt backbone--nObb_feat
is the number of classes the backbone was pre-trained on ( Kermany:4
, ChestMNIST:14
)--task
is the 3D Fine-tuning task (classification or regression)--nslc
is the number of slices to use for 3D Fine-tuning--depth
is the Vision Transformer depth--heads
is the number of heads for multihead attention--dim
specifies the dimension for encoding transformer input--mlp_dim
denotes the multi-layer perceptron dimension for ViT--dropout
is the dropout rate--emb_dropout
is the embeding dropout rate--patience
is the number of epochs to await without improvement before initiating early stopping
python slivit_test.py --dataset3d nodulemnist --checkpoint ./Checkpoints/slivit_noduleMNIST --depth 5 --dim 64 --nslc 28 --mlp_dim 64 --heads 10
python slivit_test.py --dataset3d ukbb --meta_csv ./Dsets/ukbb_meta.csv --checkpoint ./Checkpoints/slivit_ukbb --metric r2 --pathology PDFF --depth 5 --dim 256 --nslc 36 --heads 36
python slivit_test.py --dataset3d ultrasound --meta_csv ./Dsets/ultrasound_meta.csv --checkpoint ./Checkpoints/slivit_ultrasound --pathology EF_b --depth 5 --dim 256 --nslc 32 --heads 32 --mlp_dim 256
python slivit_test.py --dataset3d custom --meta_csv /path/to/generated/meta.csv --bbpath /path/to/finetuned/convnext_bb.pth --task TaskType --pathology Pathology
--dataset3d
is the dataset for 3D fine-tuning (nodulemnist
,ukbb
,ultrasound
,custom
)--meta_csv
is the path to the createdmeta.csv
file--pathology
is pathology for 3D fine-tuning--nObb_feat
is the number of classes the backbone was pre-trained on ( Kermany:4
, ChestMNIST:14
)--task
is the 3D Fine-tuning task (classification or regression)--metric
is the score metric for evaluation (roc-auc
,pr-auc
,r2
)--checkpoint
is the path to fine-tuned slivit--nslc
is the number of slices to use for 3D Fine-tuning--depth
is the Vision Transformer depth--heads
is the number of heads for multihead attention--dim
specifies the dimension for encoding transformer input--mlp_dim
denotes the multi-layer perceptron dimension for ViT
The 2D OCT dataset was downloaded from https://www.kaggle.com/datasets/paultimothymooney/kermany2018. The 3D OCT B-scan data are not publicly available due to institutional data use policy and concerns about patient privacy. However, they are available from the authors upon reasonable request and with permission of the institutional review board. The echocardiogram dataset is available at https://echonet.github.io/dynamic/index.html#dataset. The MRI dataset is available at https://www.ukbiobank.ac.uk. The 3D CT, the 2D CT, and the 2D X-ray datasets are available at https://medmnist.com/.
We kindly request users to cite the corresponding paper upon using our code, checkpoints, or conclusions in any capacity. Proper credit supports and recognizes the original creators' efforts.
@article{avram2023slivit,
title={SLIViT: a general AI framework for clinical-feature diagnosis from limited 3D biomedical-imaging data},
author={Avram, Oren and others},
journal={InReview preprint https://www.researchsquare.com/article/rs-3044914/latest},
year={2023}
}
Parts of the code are taken from the vit-pytorch repository. The figures depicting Kermany and the X-ray dataset are sourced from [1] and [2] respectively.
[1] Kermany DS, Goldbaum M, Cai W, Valentim CCS, Liang H, Baxter SL, McKeown A, Yang G, Wu X, Yan F, Dong J, Prasadha MK, Pei J, Ting MYL, Zhu J, Li C, Hewett S, Dong J, Ziyar I, Shi A, Zhang R, Zheng L, Hou R, Shi W, Fu X, Duan Y, Huu VAN, Wen C, Zhang ED, Zhang CL, Li O, Wang X, Singer MA, Sun X, Xu J, Tafreshi A, Lewis MA, Xia H, Zhang K. Identifying Medical Diagnoses and Treatable Diseases by Image-Based Deep Learning. Cell. 2018 Feb 22;172(5):1122-1131.e9. doi: 10.1016/j.cell.2018.02.010. PMID: 29474911.
[2] Wang, X., Peng, Y., Lu, L., Lu, Z., Bagheri, M., & Summers, R. (2017). ChestX-Ray8: Hospital-Scale Chest X-Ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases. In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). IEEE.