MedSegDiff：基于扩散模型的医学图像分割

MedSegDiff是一个基于扩散概率模型（DPM）的医学图像分割框架。该算法在我们的论文《MedSegDiff：基于扩散概率模型的医学图像分割》和《MedSegDiff-V2：基于Transformer的扩散医学图像分割》中有详细阐述。

扩散模型通过不断添加高斯噪声来破坏训练数据，然后学习如何通过逆转这个噪声过程来恢复数据。训练完成后，我们可以通过将随机采样的噪声传递给学习到的去噪过程来生成数据。在这个项目中，我们将这一思想扩展到医学图像分割。我们利用原始图像作为条件，从随机噪声生成多个分割图，然后对它们进行集成以获得最终结果。这种方法捕捉了医学图像中的不确定性，并在多个基准测试中优于先前的方法。

快速概览

MedSegDiff-V1	MedSegDiff-V2

新闻

[置顶] 加入我们的Discord讨论组，提问并与他人交流。
22-11-30. 本项目仍在快速更新中。查看TODO列表以了解下一步计划。
22-12-03. 修复BraTs2020的bug。添加示例案例。
22-12-15. 修复多GPU分布式训练问题。
22-12-16. DPM-Solver ✖️ MedSegDiff完成 🥳 现在可以在MedsegDiff中使用DPM-Solver。通过设置--dpm_solver True来享受其超快的采样速度（1000步 ❌ 20步 ⭕️）。
22-12-23. 修复DPM-Solver的一些bug。
23-01-31. MedSegDiff-V2即将推出 🥳 。请先查看我们的论文《MedSegDiff-V2：基于Transformer的扩散医学图像分割》。
23-02-07. 优化BRATS采样工作流程。添加用于处理原始3D BRATS数据的数据加载器。
23-02-11. 修复3D BRATS数据训练bug，[issue 31]。
23-03-04. 论文《MedSegDiff：基于扩散概率模型的医学图像分割》已正式被MIDL 2023接收 🥳
23-04-11. 基于v2框架的新版本已发布 🥳。它比之前的版本更准确、更稳定、更具域适应性，同时不会占用太多资源。我们还修复了许多小问题，如requirement.txt和isic csv文件。非常感谢所有报告问题的人，你们真的帮了我们很多忙 🤗。另外，默认情况下将运行新版本。如果想运行之前的版本，请添加"--version 1"。
23-04-12. 为isic数据集添加了简单的评估文件（script/segmentation_env）。用法：python scripts/segmentation_env.py --inp_pth *保存预测图像的文件夹* --out_pth *保存真实标签图像的文件夹*
23-12-05. 论文《MedSegDiff-V2：基于Transformer的扩散医学图像分割》已正式被AAAI 2024接收 🥳

要求

pip install -r requirement.txt

示例案例

从皮肤图像中进行黑色素瘤分割

从 https://challenge.isic-archive.com/data/ 下载 ISIC 数据集。"data" 文件夹下的数据集结构应如下所示：

data
|   ----ISIC
|       ----Test
|       |   |   ISBI2016_ISIC_Part1_Test_GroundTruth.csv
|       |   |   
|       |   ----ISBI2016_ISIC_Part1_Test_Data
|       |   |       ISIC_0000003.jpg
|       |   |       .....
|       |   |
|       |   ----ISBI2016_ISIC_Part1_Test_GroundTruth
|       |           ISIC_0000003_Segmentation.png
|       |   |       .....
|       |           
|       ----Train
|           |   ISBI2016_ISIC_Part1_Training_GroundTruth.csv
|           |   
|           ----ISBI2016_ISIC_Part1_Training_Data
|           |       ISIC_0000000.jpg
|           |       .....
|           |       
|           ----ISBI2016_ISIC_Part1_Training_GroundTruth
|           |       ISIC_0000000_Segmentation.png
|           |       .....

训练时，运行：python scripts/segmentation_train.py --data_name ISIC --data_dir *输入数据目录* --out_dir *输出数据目录* --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 1e-4 --batch_size 8
采样时，运行：python scripts/segmentation_sample.py --data_name ISIC --data_dir *输入数据目录* --out_dir *输出数据目录* --model_path *保存的模型* --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --num_ensemble 5
评估时，运行 python scripts/segmentation_env.py --inp_pth *保存预测图像的文件夹* --out_pth *保存真实标签图像的文件夹*

默认情况下，样本将保存在 ./results/ 目录下

从 MRI 中进行脑肿瘤分割

从 https://www.med.upenn.edu/cbica/brats2020/data.html 下载 BRATS2020 数据集。您的数据集文件夹结构应如下所示：

data
└───training
│   └───slice0001
│       │   brats_train_001_t1_123_w.nii.gz
│       │   brats_train_001_t2_123_w.nii.gz
│       │   brats_train_001_flair_123_w.nii.gz
│       │   brats_train_001_t1ce_123_w.nii.gz
│       │   brats_train_001_seg_123_w.nii.gz
│   └───slice0002
│       │  ...
└───testing
│   └───slice1000
│       │  ...
│   └───slice1001
│       │  ...

训练时，运行：python scripts/segmentation_train.py --data_dir （放置数据文件夹的位置）/data/training --out_dir 输出数据目录 --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --lr 1e-4 --batch_size 8
采样时，运行：python scripts/segmentation_sample.py --data_dir （放置数据文件夹的位置）/data/testing --out_dir 输出数据目录 --model_path 保存的模型 --image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16 --diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False --num_ensemble 5

其他示例

...

在自己的数据集上运行

在其他数据集上运行MedSegDiff很简单。只需按照 ./guided_diffusion/isicloader.py 或 ./guided_diffusion/bratsloader.py 的方式编写另一个数据加载器文件即可。如果遇到任何问题，欢迎提出问题。如果您能贡献您的数据集扩展，我们将不胜感激。与自然图像不同，医学图像因不同任务而有很大差异。扩展方法的泛化能力需要每个人的努力。

超参数和训练建议

要训练一个精细模型，即论文中的MedSegDiff-B，请将模型超参数设置为：

--image_size 256 --num_channels 128 --class_cond False --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16

扩散超参数设置为：

--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False

要加速采样：

--diffusion_steps 50 --dpm_solver True

在多个GPU上运行：

--multi-gpu 0,1,2（举例）

训练超参数设置为：

--lr 5e-5 --batch_size 8

并在采样时设置 --num_ensemble 5。

在大多数数据集上，训练约100,000步就会收敛。请注意，虽然在后期大部分步骤中损失不会下降，但结果的质量仍在不断提高。这种过程在其他DPM应用中也有观察到，比如图像生成。希望有聪明人能告诉我为什么🥲。

我很快会发布它在较小批量大小下的性能（适合在24GB GPU上运行），以满足比较需求🤗。

释放其全部潜力的设置是（MedSegDiff++）：

--image_size 256 --num_channels 512 --class_cond False --num_res_blocks 12 --num_heads 8 --learn_sigma True --use_scale_shift_norm True --attention_resolutions 24

然后用批量大小 --batch_size 64 训练它，并用集成数量 --num_ensemble 25 进行采样。

成为MedSegDiff的一部分！作者就是您！

欢迎为MedSegDiff做出贡献。任何能够提高性能或加速算法的技术都将不胜感激🙏。我正在撰写MedSegDiff V2，目标是发表在Nature期刊/CVPR等刊物上。我很乐意将贡献者列为我的共同作者🤗。

待办事项

致谢

代码大量借鉴自openai/improved-diffusion、WuJunde/ MrPrism、WuJunde/ DiagnosisFirst、LuChengTHU/dpm-solver、JuliaWolleb/Diffusion-based-Segmentation、hojonathanho/diffusion、guided-diffusion、bigmb/Unet-Segmentation-Pytorch-Nest-of-Unets、nnUnet、lucidrains/vit-pytorch

引用

请引用

@inproceedings{wu2023medsegdiff,
  title={MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model},
  author={Wu, Junde and FU, RAO and Fang, Huihui and Zhang, Yu and Yang, Yehui and Xiong, Haoyi and Liu, Huiying and Xu, Yanwu},
  booktitle={Medical Imaging with Deep Learning},
  year={2023}
}

@article{wu2023medsegdiff,
  标题={MedSegDiff-V2：基于扩散模型和Transformer的医学图像分割},
  作者={吴俊德 和 季伟 和 付华柱 和 徐敏 和 金月明 和 徐彦武}
  期刊={arXiv预印本 arXiv:2301.11798},
  年份={2023}
}