可泛化的肿瘤合成

我们已经证明,在不同腹部器官中,许多类型的小肿瘤(小于2cm)看起来非常相似。这意味着我们可以仅通过训练胰腺肿瘤的扩散模型,就能让人工智能检测肝脏或肾脏的肿瘤(Q. Chen等人 CVPR 2024)。这些研究已经通过放射科医生的研究(通过挑战他们区分合成肿瘤和真实肿瘤)和使用模拟数据训练的人工智能算法的全面测试得到了验证。

<p align="center"><img width="100%" src="https://yellow-cdn.veclightyear.com/835a84d5/d3fad6cd-7f9a-42b5-b3f2-fdc61ae96b96.png" /></p>

论文

<b>走向可泛化的肿瘤合成</b> <br/> Qi Chen¹, Xiaoxi Chen², Haorui Song³, Zhiwei Xiong¹, Alan L. Yuille³, Wei Chen³ 和 Zongwei Zhou^3,* <br/> ¹中国科学技术大学, <br/> ²上海交通大学, <br/> ³约翰霍普金斯大学 <br/> CVPR, 2024 <br/> 论文 | 代码 | 幻灯片 | 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

<details> <summary style="margin-left: 25px;">预处理详情</summary> <div style="margin-left: 25px;"> 1. 根据[安装说明](documents/INSTALL.md)下载数据集。 2. 修改[data_transfer.py](https://github.com/MrGiovanni/DiffTumor/blob/main/data_transfer.py)中的`data_dir`和`tumor_save_dir`。 3. `python -W ignore data_transfer.py`

开始训练。

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。论文内容受待审专利保护。