[论文解读] Diffusion Models for Implicit Image Segmentation Ensembles

Julia Wolleb, Robin Sandkühler+3 more|arXiv (Cornell University)|Dec 6, 2021

Medical Image Segmentation Techniques被引用 71

一句话总结

论文训练扩散概率模型以生成脑肿瘤分割掩模，并利用随机采样形成隐式集成，生成不确定性图并提高 BRATS2020 的分割性能。

ABSTRACT

Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.

研究动机与目标

• 激励医学影像分割中的不确定性量化并提高病变分割性能。
• 提出一种基于扩散模型的方法，在训练和采样时利用输入图像作为图谱先验。
• 通过随机扩散采样实现分割掩模的隐式集成生成。
• 通过多份分割样本的像素级方差提供不确定性图。

提出的方法

• 通过将噪声添加到真实掩模的方式训练 DDPM 以生成分割掩模，同时通过拼接将输入的 MR 图像作为先验引入。
• 修改 DDPM 的采样方程以以图像先验 X_t = b 〃irst(x_b,t) 为条件并生成一个均值/方差分割集成。
• 使用线性噪声时间表，T=1000 步，以及 U-Net 架构来预测去噪步长 epsilon_theta。
• 通过对多份分割掩模采样并计算其方差来得到像素级不确定性图。
• 通过对多份采样掩模取平均并进行阈值化来获得最终分割。
• 在 BRATS2020 上与 nnU-Net 和 SegNet 进行基准测试，并报告 Dice、Jaccard 和 HD95 指标。

实验结果

研究问题

• RQ1 扩散模型能否通过将相应图像作为先验来适配医学图像的语义分割？
• RQ2 通过随机扩散采样生成多份分割样本是否比单份预测带来更好的分割指标？
• RQ3 在 BRATS2020 上，基于扩散的分割在准确性和不确定性可视化方面与最先进方法(nnU-Net、SegNet)相比如何？

主要发现

Dice 系数	HD95 距离	Jaccard 系数	空
我们 (1 次采样运行)	0.866 [0.892]	6.052	0.795 [0.819]	31
我们 (5 次运行的集成)	0.881 [0.909]	5.178	0.819 [0.845]	34
nnU-Net (5 折交叉验证的集成)	0.891 [0.905]	5.004	0.831 [0.845]	17
SegNet (1 次运行)	0.839 [0.867]	7.190	0.761 [0.786]	34
贝叶斯 SegNet (5 次运行的集成)	0.838 [0.841]	13.707	0.747 [0.749]	3

• 单样本 DDPM 分割在 Dice 0.866 和 Jaccard 0.795，HD95 为 6.052。
• 5 次扩散样本的集成将 Dice 提升至 0.881，Jaccard 提升至 0.819，HD95 为 5.178。
• nnU-Net 达到 Dice 0.891、Jaccard 0.831，HD95 5.004。
• SegNet 得到 Dice 0.839、Jaccard 0.761，HD95 7.190。
• 贝叶斯 SegNet（MC dropout，5 次样本）Dice 0.838，Jaccard 0.747，HD95 13.707。
• 集成均值图提供改进的分割性能并实现详细的不确定性可视化。