MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

基于扩散概率模型的医学图像分割

来源：Proceedings of Machine Learning Research   227:1623–1639, 2023

摘要：扩散概率模型(Diffusion Probabilistic Model, DPM)是近年来计算机视觉领域研究的热点之一。它的图像生成应用程序，如Imagen、Latent Diffusion Models和Stable Diffusion，已经展示了令人印象深刻的生成能力，在社区中引发了广泛的讨论。此外，最近的许多研究发现DPM在各种其他视觉任务中很有用，包括图像去模糊、超分辨率和异常检测。受DPM成功的启发，我们提出了MedSegDiff，这是第一个用于一般医学图像分割任务的基于DPM的模型。为了增强DPM在医学图像分割中的分步区域注意力，我们提出了动态条件编码，该编码为每个采样步骤建立状态自适应条件。此外，我们提出了特征频率解析器(FF-Parser)来消除高频噪声成分在此过程中的负面影响。我们验证了MedSegDiff在三种不同图像模式下的医学分割任务中的有效性，包括眼底图像的视杯分割、MRI图像的脑肿瘤分割和超声图像的甲状腺结节分割。我们的实验结果表明，MedSegDiff在性能上明显优于最先进的SOTA方法，证明了所提出模型的泛化和有效性。

、

介绍：最近，扩散概率模型(DPM)作为一种强大的生成模型而受到欢迎(Ho et al .， 2020)。这些模型能够生成具有高多样性和高合成质量的图像。大型扩散模型，如DALL-E2 (Ramesh等人，2022)、Imagen(撒哈拉等人，2022a)和Stable diffusion (Rombach等人，2022)，已经展示了令人难以置信的生成能力(Zhao和Shi, 2021;Goodfellow et al, 2020)。扩散模型最初应用于不存在标签的领域。然而，最近的研究表明，它们对于标签是唯一的问题也有效，例如超分辨率(撒哈拉等人，2022b)，去模糊(Whang等人，2022)和分割(Amit等人，2021)。

受到 DPM 最近成功的启发，我们提出了一种独特的基于 DPM 的医学图像分割任务分割模型。据我们所知，我们是第一个提出基于 DPM 的通用医学图像分割模型的人。我们注意到，在医学图像分割任务中，病变/器官通常是模棱两可的，难以与背景区分开来。在这种情况下，自适应校准过程对于获得准确的结果至关重要。按照这个想法，我们提出了基于普通条件 DPM 的动态条件编码来设计所提出的模型，名为 MedSegDiff。在迭代采样过程中，MedSegDiff 使用图像先验来调节每个步骤以学习分割图。为了实现自适应区域注意，我们在每个步骤中将当前步骤的分割图集成到图像先验编码中。具体而言，我们以多尺度方式在特征级别将当前步骤分割掩码与图像先验融合。 这允许损坏的当前步骤掩码动态增强条件特征并提高重建精度。为了消除损坏掩模中的高频噪声，我们提出了特征频率解析器（FF-Parser）来过滤傅里叶空间中的特征。在每个跳跃连接路径上采用 FF-Parser 进行多尺度集成。

文章贡献;

•我们提出了第一个基于dpm的一般医学图像分割模型。

•我们引入了动态条件编码策略来实现分步注意。

•我们提出了FF-Parser方法来有效地消除高频噪声成分。

•我们的模型在具有不同图像模式的三种不同的医学分割任务上实现了最先进的性能。

我们基于(Ho et al, 2020)中提出的扩散模型设计了我们的模型。扩散模型是由两个阶段组成的生成模型:正向扩散阶段和反向扩散阶段。在前向过程中，一个分割标签x0经过一系列的T步被高斯噪声逐渐破坏。在反向过程中，训练神经网络通过反转噪声过程来恢复原始数据。这可以表示为:

其中θ是逆向过程的一组参数。从高斯噪声分布出发，pθ(xT) = N (xT;0, In×n)，其中I为原始图像，反向过程将潜在变量分布pθ(xT)变换为数据分布pθ(x0)。为了与正演过程保持对称，反向过程逐级恢复噪声图像，得到最终清晰的分割结果。

遵循DPM的标准实现，我们采用UNet作为学习网络。图1显示了一个示例。为了实现分割，我们使用原始图像先验条件阶跃估计函数λ，其给出如下:

其中EI t为原始图像的条件特征嵌入，Ex t为当前步骤分割映射的特征嵌入。这两个嵌入被添加并通过UNet解码器D进行重建。步骤索引t与添加的嵌入和解码器特征集成，并使用共享学习查找表进行嵌入，如(Ho et al, 2020)中所述。

具体实现过程;

在大多数条件DPM中，条件先验将是唯一的特征嵌入。然而，医学图像分割被认为是具有挑战性的，因为物体模糊，病变或组织往往难以从背景中区分出来。低对比度的图像模式，如MRI或超声图像，进一步加剧了这个问题。为了解决这个问题，我们提出了每个步骤的动态条件编码。我们观察到原始图像包含准确的分割目标信息，但难以与背景区分。同时，当前步分割图包含增强的目标区域，但并不完全准确。这促使我们进行整合将当前步分割信息xt转换成有条件的原始图像编码进行相互补全。具体来说，我们在功能级别上实现集成。在原始图像编码器中，我们利用当前步长编码特征来增强其中间特征。

条件特征映射mk I的每个尺度与相同形状的xt编码特征mk x融合，其中k为层的索引。融合是使用一种类似于注意力的机制a来实现的。在这种机制中，我们首先对两个特征映射应用层归一化，然后将它们相乘得到一个亲和力映射。然后，我们将亲和映射与条件编码特征相乘，以增强关注区域，如下所示:

其中⊗表示逐元素乘法，LN表示层归一化。我们将此操作应用于中间两个阶段，其中每个阶段都是遵循ResNet34实现的卷积阶段。

这种动态条件编码策略有助于MedSegDiff动态定位和校准分割。但由于对被噪声破坏的文本嵌入进行积分，可能会引入额外的高频噪声。为了解决这个问题，我们进一步提出了高频解析器来约束特征中的高频成分。

FF-Parser

我们将FF-Parser引入到特征集成路径中，以解决xt嵌入集成带来的高频噪声问题。FF-Parser设计用于约束xt特性中与噪声相关的组件。主要思想是学习一个参数化的关注映射，应用于傅里叶空间特征。给定一个解码器特征映射m∈R H×W×C，我们首先沿着空间维度进行二维FFT(快速傅立叶变换)，可以表示为:

式中F[·]为二维FFT。接下来，我们通过将一个参数化的关注映射a∈C H×W×C乘以m来调制m的频谱:

其中⊗表示元素相关产物。最后，我们采用逆FFT将M '逆回空间域:

FF-Parser可以看作是频率滤波器的一种可学习版本，在数字图像处理中得到了广泛的应用(Pitas, 2000)。与空间注意力不同，FF-Parser全局调整特定频率的组成部分。因此，可以训练它来约束高频分量进行自适应集成。

MedSegDiff的主要架构是一个改进的ResUNet (Yu et al .， 2019)，我们使用ResNet编码器和UNet解码器来实现。我们遵循(Nichol和Dhariwal, 2021)的详细网络设置。I和xt都使用两个单独的编码器进行编码。每个编码器由包含多个残差块的卷积阶段组成。每个阶段的剩余块数量遵循ResNet34。每个残块由两个卷积块组成，每个卷积块由群归一化层和SiLU (Elfwing et al, 2018)激活层组成，然后是卷积层。

残差块通过一个线性层、一个SiLU激活和另一个线性层接收时间嵌入。然后将结果添加到第一个卷积块的输出中。将得到的EI和Ext加在一起，发送到最后的编码阶段。连接一个标准的卷积解码器来预测最终结果。

数据集介绍;

我们使用不同的图像模式对三种不同的医学任务进行了实验:眼底图像的光学杯分割，MRI图像的脑肿瘤分割，超声图像的甲状腺结节分割。我们在分别包含1200、2000和8046个样本的ence2数据集(Fang等人，2022)、brat -2021数据集(Baid等人，2021)和DDTI数据集(Pedraza等人，2015)上评估了我们的方法在青光眼、甲状腺癌和黑色素瘤诊断方面的性能。在这些数据集中，分割和诊断标签都是公开可用的，我们按照各自数据集的默认设置将数据分成训练/验证/测试集。

结论;

本文提出了一种基于dpm的通用医学图像分割方案MedSegDiff，该方案结合了动态条件编码和FF-Parser两种新技术来提高分割性能。我们对三种不同图像模式的医学图像分割任务进行了对比实验，结果表明我们的模型优于以往的SOTA方法。作为第一个基于dpm的医学图像分割应用，我们相信MedSegDiff将成为该领域未来研究的重要基准.