We have document for common questions for code and common questions for paper.

CLIP-Driven Universal Model

Paper

This repository provides the official implementation of Universal Model.

CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection
Rank First in Medical Segmentation Decathlon (MSD) Competition
Jie Liu¹, Yixiao Zhang², Jie-Neng Chen², Junfei Xiao², Yongyi Lu²,
Yixuan Yuan¹, Alan Yuille², Yucheng Tang³, Zongwei Zhou²
¹City University of Hong Kong, ²Johns Hopkins University, ³NVIDIA
paper | code | slides | poster | talk | blog

⏳ Dataset Link

💡 Preparation

Main Requirements

connected-components-3d
h5py==3.6.0
monai==0.9.0
torch==1.11.0
tqdm
fastremap

python3 -m venv universal
source /data/zzhou82/environments/universal/bin/activate

git clone https://github.com/ljwztc/CLIP-Driven-Universal-Model.git
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
pip install 'monai[all]'
pip install -r requirements.txt
cd pretrained_weights/
wget https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/swin_unetr.base_5000ep_f48_lr2e-4_pretrained.pt
cd ../

Dataset Pre-Process

Download the dataset according to the dataset link and arrange the dataset according to the dataset/dataset_list/PAOT.txt.
Modify the ORGAN_DATASET_DIR value in label_transfer.py (line 51) and NUM_WORKER (line 53)
python -W ignore label_transfer.py

Current Template

Index	Organ
1	Spleen
2	Right Kidney
3	Left Kidney
4	Gall Bladder
5	Esophagus
6	Liver
7	Stomach
8	Aorta
9	Postcava
10	Portal Vein and Splenic Vein
11	Pancreas
12	Right Adrenal Gland
13	Left Adrenal Gland
14	Duodenum
15	Hepatic Vessel
16	Right Lung
17	Left Lung
18	Colon
19	Intestine
20	Rectum
21	Bladder
22	Prostate
23	Left Head of Femur
24	Right Head of Femur
25	Celiac Truck
26	Kidney Tumor
27	Liver Tumor
28	Pancreas Tumor
29	Hepatic Vessel Tumor
30	Lung Tumor
31	Colon Tumor
32	Kidney Cyst

How expand to new dataset with new organ?

Set the following index for new organ. (e.g. 33 for vermiform appendix)
Check if there are any organs that are not divided into left and right in the dataset. (e.g. kidney, lung, etc.) The RL_Splitd in label_transfer.py is used to processed this case.
Set up a new transfer list for new dataset in TEMPLATE (line 58 in label_transfer.py). (If a new dataset with Intestine labeled as 1 and vermiform appendix labeled as 2, we set the transfer list as [19, 33])
Run the program label_transfer.py to get new post-processing labels.
More details please take a look at common questions

📦 Training

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -W ignore -m torch.distributed.launch --nproc_per_node=8 --master_port=1234 train.py --dist True --data_root_path /mnt/zzhou82/PublicAbdominalData/ --num_workers 12 --num_samples 4 --cache_dataset --cache_rate 0.6 --uniform_sample

📦 Validation

CUDA_VISIBLE_DEVICES=0 python -W ignore validation.py --data_root_path /mnt/zzhou82/PublicAbdominalData/ --start_epoch 10 --end_epoch 40 --epoch_interval 10 --cache_dataset --cache_rate 0.6

📦 Test

CUDA_VISIBLE_DEVICES=0 python -W ignore test.py --resume ./out/epoch_61.pth --data_root_path /mnt/zzhou82/PublicAbdominalData/ --store_result --cache_dataset --cache_rate 0.6

📒 To do

Code release
Dataset link
Support different backbones (SwinUNETR, Unet, DiNTS, Unet++)
Model release
Pesudo label release
Tutorials for generalizability, transferability, and extensibility

🛡️ License

This project is under the CC-BY-NC 4.0 license. See LICENSE for details.

🙏 Acknowledgement

A lot of code is modified from monai.

📝 Citation

If you find this repository useful, please consider citing this paper:

@article{liu2023clip,
  title={CLIP-Driven Universal Model for Organ Segmentation and Tumor Detection},
  author={Liu, Jie and Zhang, Yixiao and Chen, Jie-Neng and Xiao, Junfei and Lu, Yongyi and Landman, Bennett A and Yuan, Yixuan and Yuille, Alan and Tang, Yucheng and Zhou, Zongwei},
  journal={arXiv preprint arXiv:2301.00785},
  year={2023}
}

Xinya-liuliu/CLIP-Driven-Universal-Model