DiffTumor

<h1 align="center">DiffTumor</h1> <div align="center">

<a href="https://twitter.com/bodymaps317"> <img src="https://img.shields.io/twitter/follow/BodyMaps?style=social" alt="Follow on Twitter" /> </a><br/> Subscribe us: https://groups.google.com/u/2/g/bodymaps

</div>

We have shown that many types of tumors in different abdominal organs look very similar if they are small (less than 2cm) which implies that we can train the AI to detect tumors in the liver or kidney but only training the diffusion model on tumors in the pancreas (Q. Chen et al. CVPR 2024). These studies have been validated both by studies of radiologists (by challenging them to distinguish between a synthetic tumor and a real tumor) and by comprehensive tests of the AI algorithms trained using stimulated data.

<p align="center"><img width="100%" src="figures/fig_teaser.png" /></p>

Paper

<b>Towards Generalizable Tumor Synthesis</b> <br/> Qi Chen¹, Xiaoxi Chen², Haorui Song³, Zhiwei Xiong¹, Alan L. Yuille³, Wei Chen³ and Zongwei Zhou^3,* <br/> ¹University of Science and Technology of China, <br/> ²Shanghai Jiao Tong University, <br/> ³Johns Hopkins University <br/> CVPR, 2024 <br/> paper | code | slides | huggingface

We have documented common questions for the paper in Frequently Asked Questions (FAQ).

We have summarized publications related to tumor synthesis in Awesome Synthetic Tumors .

We have released videos for Visual Turing Test. Check to see if you could tell which is real tumor and which is synthetic tumor.

0. Installation

git clone https://github.com/MrGiovanni/DiffTumor.git
cd DiffTumor

See installation instructions to create an environment and obtain requirements.

1. Train Autoencoder Model

You can train Autoencoder Model on AbdomenAtlas 1.0 dataset by your own. The release of AbdomenAtlas 1.0 can be found at https://huggingface.co/datasets/AbdomenAtlas/AbdomenAtlas_1.0_Mini.

cd STEP1.AutoencoderModel
datapath=<your-datapath> (e.g., /data/bdomenAtlasMini1.0/)
gpu_num=1
cache_rate=0.05
batch_size=4
dataset_list="AbdomenAtlas1.0Mini"
python train.py dataset.data_root_path=$datapath dataset.dataset_list=$dataset_list dataset.cache_rate=$cache_rate dataset.batch_size=$batch_size model.gpus=$gpu_num

We offer the pre-trained checkpoint of Autoencoder Model, which was trained on AbdomenAtlas 1.1 dataset (see details in SuPreM). This checkpoint can be directly used for STEP2 if you do not want to re-train the Autoencoder Model. Simply download it to STEP2.DiffusionModel/pretrained_models/AutoencoderModel.ckpt

cd STEP2.DiffusionModel/pretrained_models/
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/AutoencoderModel/AutoencoderModel.ckpt

2. Train Diffusion Model

In our study, Diffusion Model focuses on the tumor region generation (simple texture and small shape). Early-stage tumors appear similar in the three abdominal organs, enabling models to effectively learn these characteristics from minimal examples. If you want to train Diffusion Model that synthesize early tumors, you need to first process the data to filter out the early tumors labels. We take the example of training Diffusion Model for early-stage liver tumors.

Download the public dataset MSD-Liver (More datasets can be seen in installation instructions).

wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/Task03_Liver.tar.gz
tar -zxvf Task03_Liver.tar.gz

We offer the preprocessed labels for early-stage tumors and mid-/late- stage tumors.

wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/preprocessed_labels.tar.gz
tar -zxvf preprocessed_labels.tar.gz

<details> <summary style="margin-left: 25px;">Preprocess details</summary> <div style="margin-left: 25px;">

1. Download the dataset according to the installation instructions.
2. Modify data_dir and tumor_save_dir in data_transfer.py.
3. python -W ignore data_transfer.py

</div> </details> Start training.

cd STEP2.DiffusionModel/
vqgan_ckpt=<pretrained-AutoencoderModel> (e.g., /pretrained_models/AutoencoderModel.ckpt)
fold=0
datapath=<your-datapath> (e.g., /data/10_Decathlon/Task03_Liver/)
tumorlabel=<your-labelpath> (e.g., /data/preprocessed_labels/)
python train.py dataset.name=liver_tumor_train dataset.fold=$fold dataset.data_root_path=$datapath dataset.label_root_path=$tumorlabel dataset.dataset_list=['liver_tumor_data_early_fold'] dataset.uniform_sample=False model.results_folder_postfix="liver_early_tumor_fold"$fold""  model.vqgan_ckpt=$vqgan_ckpt

We offer the pre-trained checkpoints of Diffusion Model, which were trained for early-stage and mid-/late- stage tumors for liver, pancreas and kidney, respectively. This checkpoint can be directly used for STEP3 if you do not want to re-train the Diffusion Model. Simply download it to STEP3.SegmentationModel/TumorGeneration/model_weight

<details> <summary style="margin-left: 25px;">Checkpoints</summary> <div style="margin-left: 25px;">

| Tumor | Type | Download | | ---- | ---- | ---- | | liver | early | link | | liver | mid&late | link | | pancreas | early | link | | pancreas | mid&late | link | | kidney | early | link | | kidney | mid&late | link |

</div> </details>

3. Train Segmentation Model

Download healthy CT data

<details> <summary style="margin-left: 25px;">from Huggingface</summary> <div style="margin-left: 25px;">

(More details can be seen in the corresponding huggingface repository).

mkdir HealthyCT
cd HealthyCT
huggingface-cli download qicq1c/HealthyCT  --repo-type dataset --local-dir .  --cache-dir ./cache
cat healthy_ct.zip* > HealthyCT.zip
rm -rf healthy_ct.zip* cache
unzip -o -q HealthyCT.zip -d /HealthyCT

</div> </details> <details> <summary style="margin-left: 25px;">from Dropbox</summary> <div style="margin-left: 25px;">

wget https://www.dropbox.com/scl/fi/j8di09jm1s798ofnwlkf6/HealthyCT.tar.gz?rlkey=ujuc82109eceld1vmwwu24z26
mv HealthyCT.tar.gz?rlkey=ujuc82109eceld1vmwwu24z26 HealthyCT.tar.gz
tar -xzvf HealthyCT.tar.gz

</div> </details>

Prepare Autoencoder and Diffusion Model. Put the pre-trained weights to STEP3.SegmentationModel/TumorGeneration/model_weight

cd STEP3.SegmentationModel/TumorGeneration/model_weight/
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/AutoencoderModel/AutoencoderModel.ckpt
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/DiffusionModel/liver_early.pt
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/DiffusionModel/liver_noearly.pt
cd ../..

Start training.

cd STEP3.SegmentationModel

healthy_datapath=<your-datapath> (e.g., /data/HealthyCT/)
datapath=<your-datapath> (e.g., /data/10_Decathlon/Task03_Liver/)
cache_rate=1.0
batch_size=12
val_every=50
workers=12
organ=liver
fold=0

# U-Net
backbone=unet
logdir="runs/$organ.fold$fold.$backbone"
datafold_dir=cross_eval/"$organ"_aug_data_fold/
dist=$((RANDOM % 99999 + 10000))
python -W ignore main.py --model_name $backbone --cache_rate $cache_rate --dist-url=tcp://127.0.0.1:$dist --workers $workers --max_epochs 2000 --val_every $val_every --batch_size=$batch_size --save_checkpoint --distributed --noamp --organ_type $organ --organ_model $organ --tumor_type tumor --fold $fold --ddim_ts 50 --logdir=$logdir --healthy_data_root $healthy_datapath --data_root $datapath --datafold_dir $datafold_dir

# nnU-Net
backbone=nnunet
logdir="runs/$organ.fold$fold.$backbone"
datafold_dir=cross_eval/"$organ"_aug_data_fold/
dist=$((RANDOM % 99999 + 10000))
python -W ignore main.py --model_name $backbone --cache_rate $cache_rate --dist-url=tcp://127.0.0.1:$dist --workers $workers --max_epochs 2000 --val_every $val_every --batch_size=$batch_size --save_checkpoint --distributed --noamp --organ_type $organ --organ_model $organ --tumor_type tumor --fold $fold --ddim_ts 50 --logdir=$logdir --healthy_data_root $healthy_da