Diffusion Models for Implicit Image Segmentation Ensembles
隐式图像分割集成的扩散模型
摘要
扩散模型在图像生成建模方面表现出令人印象深刻的性能。本文提出了一种新的基于扩散模型的语义分割方法。通过改进训练和采样方案,我们证明了扩散模型可以对医学图像进行病灶分割。为了生成特定于图像的分割,我们在地面真值分割上训练模型,并在训练和采样过程的每一步中使用图像作为先验。利用给定的随机采样过程,我们可以生成一个分割掩码的分布(确定每个像素或每个区域在图像或空间中是否属于特定的类别或对象。)。此属性允许我们计算分割的像素不确定性映射,并允许隐式的分割集合,从而提高分割性能。我们在BRATS2020数据集上评估了我们的方法用于脑肿瘤分割。与最先进的分割模型相比,我们的方法产生了良好的分割结果,此外,还提供了详细的不确定性图。
1. 介绍
语义分割是医学图像分析中一个重要且探索得很好的领域(Rizwan I Haque and Neubert, 2020)。利用机器学习对医学图像中的病变进行自动分割已经显示出良好的性能(Isensee et al ., 2021),并且已经准备好用于临床应用以支持诊断(Sharrock et al ., 2021)。在医学应用中,测量给定预测的不确定性是非常有趣的,特别是当用于放射治疗等进一步治疗时。
在这项工作中,我们专注于BRATS2020脑肿瘤分割挑战(Menze et al, 2014;Bakas et al, 2017,2018)。该数据集为每个患者提供了四种不同的MR序列(即t1加权、t2加权、FLAIR和对比度增强的t1加权),以及基于像素的地面真值分割。在附录A中可以找到一个示例图像。
我们提出了一种基于去噪扩散概率模型(DDPM)的新分割方法(Ho et al ., 2020),该方法可以提供生成的分割掩模的不确定性映射。BRATS2020数据集图像的工作流程概述如图1所示。我们在分割蒙版上训练一个DDPM,并将原始脑磁共振图像作为图像添加,从而诱导出解剖信息。作为抽样
ddpm在每个采样步骤中都有一个随机元素,我们可以为相同的输入图像和相同的预训练模型生成许多不同的分割掩码。这种分割集合允许我们计算逐像素的方差图,从而可视化生成的分割的不确定性。此外,在平均映射中集成分割,提高了分割性能。
我们将自己与最先进的分割算法进行比较,并直观地将我们的方差图与常见的不确定性图进行比较。代码将在https://github.com/JuliaWolleb/Diffusion-based-Segmentation上公开。
相关工作
在医学图像分割中,常用的方法是应用U-Net (Ronneberger et al ., 2015)或SegNet (Badrinarayanan et al ., 2017)来预测每个输入图像的分割掩码。这种方法已