访问官网
Github
Huggingface
文档
论文可泛化的肿瘤合成
我们已经证明,在不同腹部器官中,许多类型的小肿瘤(小于2cm)看起来非常相似。这意味着我们可以仅通过训练胰腺肿瘤的扩散模型,就能让人工智能检测肝脏或肾脏的肿瘤(Q. Chen等人 CVPR 2024)。这些研究已经通过放射科医生的研究(通过挑战他们区分合成肿瘤和真实肿瘤)和使用模拟数据训练的人工智能算法的全面测试得到了验证。
论文
走向可泛化的肿瘤合成
Qi Chen1, Xiaoxi Chen2, Haorui Song3, Zhiwei Xiong1, Alan L. Yuille3, Wei Chen3 和 Zongwei Zhou3,*
1 中国科学技术大学,
2 上海交通大学,
3 约翰霍普金斯大学
CVPR, 2024
论文 | 代码 | 幻灯片 | huggingface
我们已经在常见问题解答(FAQ)中记录了有关该论文的常见问题。
我们已经在Awesome Synthetic Tumors 
中总结了与肿瘤合成相关的出版物。
我们已经发布了视觉图灵测试的视频。 看看你能否分辨出哪个是真实肿瘤,哪个是合成肿瘤。
0. 安装
git clone https://github.com/MrGiovanni/DiffTumor.git
cd DiffTumor
查看安装说明以创建环境并获取所需的依赖项。
1. 训练自编码器模型
你可以在AbdomenAtlas 1.0数据集上自行训练自编码器模型。AbdomenAtlas 1.0的发布可以在https://huggingface.co/datasets/AbdomenAtlas/AbdomenAtlas_1.0_Mini找到。
cd STEP1.AutoencoderModel
datapath=<你的数据路径> (例如, /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
我们提供了自编码器模型的预训练检查点,该模型在AbdomenAtlas 1.1数据集上进行了训练(详见SuPreM)。如果你不想重新训练自编码器模型,这个检查点可以直接用于第2步。只需将其下载到STEP2.DiffusionModel/pretrained_models/AutoencoderModel.ckpt
cd STEP2.DiffusionModel/pretrained_models/
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/AutoencoderModel/AutoencoderModel.ckpt
2. 训练扩散模型
在我们的研究中,扩散模型专注于肿瘤区域的生成(简单的纹理和小形状)。早期肿瘤在三个腹部器官中看起来相似,使模型能够从少量的例子中有效地学习这些特征。如果你想训练合成早期肿瘤的扩散模型,你需要首先处理数据以筛选出早期肿瘤标签。我们以训练早期肝脏肿瘤的扩散模型为例。
下载公开数据集MSD-Liver(更多数据集可以在安装说明中看到)。
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/Task03_Liver.tar.gz
tar -zxvf Task03_Liver.tar.gz
我们提供了早期肿瘤和中/晚期肿瘤的预处理标签。
wget https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/preprocessed_labels.tar.gz
tar -zxvf preprocessed_labels.tar.gz
预处理详情
开始训练。
cd STEP2.DiffusionModel/
vqgan_ckpt=<预训练的AutoencoderModel> (例如，/pretrained_models/AutoencoderModel.ckpt)
fold=0
datapath=<你的数据路径> (例如，/data/10_Decathlon/Task03_Liver/)
tumorlabel=<你的标签路径> (例如，/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
我们提供了针对肝脏、胰腺和肾脏的早期和中晚期肿瘤分别训练的扩散模型预训练检查点。如果你不想重新训练扩散模型，可以直接使用这些检查点进行第3步。只需将其下载到STEP3.SegmentationModel/TumorGeneration/model_weight
3. 训练分割模型
下载健康CT数据
准备自动编码器和扩散模型。将预训练权重放入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 ../..
开始训练。
cd STEP3.SegmentationModel
healthy_datapath=<你的数据路径> (例如，/data/HealthyCT/)
datapath=<你的数据路径> (例如，/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_datapath --data_root $datapath --datafold_dir $datafold_dir
Swin-UNETR
backbone=swinunetr 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
我们提供了分割模型（U-Net、nnU-Net和Swin UNETR）的预训练检查点，这些模型在肝脏、胰腺和肾脏的真实和合成肿瘤上进行了训练。
<details>
<summary style="margin-left: 25px;">U-Net</summary>
<div style="margin-left: 25px;">
| 肿瘤 | 下载 |
| ---- | ---- |
| 肝脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/unet_synt_liver_tumors.pt) |
| 胰腺 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/unet_synt_pancreas_tumors.pt) |
| 肾脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/unet_synt_kidney_tumors.pt) |
</div>
</details>
<details>
<summary style="margin-left: 25px;">nnU-Net</summary>
<div style="margin-left: 25px;">
| 肿瘤 | 下载 |
| ---- | ---- |
| 肝脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/nnunet_synt_liver_tumors.pt) |
| 胰腺 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/nnunet_synt_pancreas_tumors.pt) |
| 肾脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/nnunet_synt_kidney_tumors.pt) |
</div>
</details>
<details>
<summary style="margin-left: 25px;">Swin UNETR</summary>
<div style="margin-left: 25px;">
| 肿瘤 | 下载 |
| ---- | ---- |
| 肝脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/swinunetr_synt_liver_tumors.pt) |
| 胰腺 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/swinunetr_synt_pancreas_tumors.pt) |
| 肾脏 | [链接](https://huggingface.co/MrGiovanni/DiffTumor/resolve/main/SegmentationModel/swinunetr_synt_kidney_tumors.pt) |
</div>
</details>
## 4. 评估
```bash
cd SegmentationModel
datapath=/mnt/ccvl15/zzhou82/PublicAbdominalData/
organ=liver
fold=0
datafold_dir=cross_eval/"$organ"_aug_data_fold/
# U-Net
python -W ignore validation.py --model=unet --data_root $datapath --datafold_dir $datafold_dir --tumor_type tumor --organ_type $organ --fold $fold --log_dir $organ/$organ.fold$fold.unet --save_dir out/$organ/$organ.fold$fold.unet
# nnU-Net
python -W ignore validation.py --model=nnunet --data_root $datapath --datafold_dir $datafold_dir --tumor_type tumor --organ_type $organ --fold $fold --log_dir $organ/$organ.fold$fold.unet --save_dir out/$organ/$organ.fold$fold.unet
# Swin-UNETR
python -W ignore validation.py --model=swinunetr --data_root $datapath --datafold_dir $datafold_dir --tumor_type tumor --organ_type $organ --fold $fold --log_dir $organ/$organ.fold$fold.unet --save_dir out/$organ/$organ.fold$fold.unet
我们还在HuggingFace 🤗上提供了DiffTumor的singularity容器
inputs_data=/path/to/your/CT/scan/folders
outputs_data=/path/to/your/output/folders
wget https://huggingface.co/qicq1c/DiffTumor/resolve/main/difftumor_final.sif
SINGULARITYENV_CUDA_VISIBLE_DEVICES=0 singularity run --nv -B $inputs_data:/workspace/inputs -B $outputs_data:/workspace/outputs difftumor_final.sif
引用
@inproceedings{chen2024towards,
title={面向可泛化的肿瘤合成},
author={陈琦 and 陈晓曦 and 宋浩睿 and 熊志伟 and Yuille, Alan and 魏晨 and 周宗伟},
booktitle={IEEE/CVF计算机视觉与模式识别会议论文集},
pages={11147--11158},
year={2024}
}
致谢
本研究得到了Lustgarten胰腺癌研究基金会和McGovern基金会的支持。代码库修改自NVIDIA MONAI。论文内容受待审专利保护。
