MedSegDiff： Medical Image Segmentation with Diffusion Probabilistic Model

简单不看版本：

有错误欢迎指正，谢谢各位大佬。

这是作者的第一版本文章，总的来说比较简单。总共提出两点改进：

1、由于医学图像较为特殊，病变组织很难与背景相区别，尤其是低分辨率的图像。另外作者认为原图中有很多目标的信息，但是很难分割，而扩散模型中的任意t时刻的分割图中有较为增强的分割目标信息，但不准确。基于这两点，作者提出了将两者融合互补的ideal。作者提出了一个动态条件编码器dynamic condition encoding，在每一步的时候都将两幅feature map进行融合。首先假设扩散模型已经生成了一张t时刻的feature map，需要训练神经网络来恢复图像，这个时候扩散模型中的feature map传递到原本的encoder中，同时将label（论文中说是raw image但我想应该不是原始图片，应该是label，作者的框架图片中也使用的是label）传递到作者提出来的动态条件编码器中，在每一层结束后进行融合。具体融合方式请看下文。

但是！在融合过程中作者发现扩散模型中任意t时刻的feature map虽然包含有一定的目标信息，但同时也会带来一些高频噪声high-frequency noise。所以有了下面这一点。

2、作者提出了FF-Parser来抑制扩散模型中的feature map带来的高频噪声，主要思想是使用了傅里叶变换来完成的。具体请看下文。

MedSegDiff： Medical Image Segmentation with Diffusion Probabilistic Model

论文：2211.00611v1.pdf (arxiv.org)

摘要：

扩散概率模型(DPM)是近年来计算机视觉研究的热点之一。它的图像生成应用，如Imagen、潜扩散模型和稳定扩散模型，已经显示出令人印象深刻的生成能力，引起了广泛的讨论。最近的许多研究还发现，它在许多其他视觉任务中也很有用，比如图像去模糊、超分辨率和异常检测。受DPM成功的启发，我们提出了第一个基于DPM的通用医学图像分割任务模型，我们将其命名为MedSegDiff。为了提高医学图像分割中DPM的分步区域注意，提出了动态条件编码方法，为每一步采样建立状态自适应条件。我们进一步提出了特征频率分析器(FF-Parser)，以消除高频噪声成分在此过程中的负面影响。我们在三种不同图像形态的医学分割任务上验证了MedSegDiff，分别是眼底图像上的光学杯分割、MRI图像上的脑肿瘤分割和超声图像上的甲状腺结节分割。实验结果表明，MedSegDiff模型在性能上优于现有的SOTA方法，具有较好的泛化性和有效性。

介绍：

医学图像分割是将医学图像划分为有意义区域的过程。分割是许多医学图像分析应用的基本步骤，如诊断、手术规划和图像引导手术。这很重要，因为它可以让医生和其他医疗专业人员更好地理解他们正在看的东西。它还可以更容易地比较图像和跟踪随时间的变化。近年来，人们对医学图像自动分割方法越来越感兴趣。这些方法有可能减少人工分割所需的时间和精力，并提高结果的一致性和准确性。随着深度学习技术的发展，越来越多的研究成功地将基于神经网络(NN)的模型应用于医学图像分割任务，从流行的卷积神经网络(CNN)[1]到最近的视觉变形器(ViT)[2]、[3]。

最近，扩散概率模型(DPM)[4]作为一种强大的生成模型而受到欢迎，它能够生成具有高多样性和合成质量的图像。最近的大型扩散模型，如DALL-E2 [5]， Imagen [6]，Stable Diffusio[7]表现出来不可思议的生成效果。扩散模型最初应用于没有绝对ground truth的领域。然而，最近的研究表明，它对于地面真实值唯一的问题也很有效，如超分辨率[8]和[9]去模糊。

受到DPM最近成功的启发，我们设计了一个独特的基于dpm的医学图像分割模型。据我们所知，在一般医学图像分割的背景下，我们首次提出了基于dpm的图像分割模型。我们注意到，在医学图像分割任务中，病变/器官通常是模糊的，很难从背景中区分出来。在这种情况下，自适应校准过程是获得精细结果的关键。按照这种思路，我们在专有dpm之上提出动态条件编码来设计模型，命名为MedSegDiff。注意，在迭代采样过程中，MedSegDiff以图像先验条件约束每一步，以便从中学习分割映射。对于自适应区域注意，我们将当前步骤的分割图整合到每一步的图像预先编码中。具体实现是将当前步分割掩码与先验图像在特征层上进行多尺度融合。这样，损坏的当前步掩码有助于动态增强条件特征，从而提高重建精度。在此过程中，为了消除已损坏的给定掩码中的高频噪声，我们进一步提出了特征频率解析器(FF-Parser)来过滤给定掩码中的特征傅里叶空间。在每个跳跃连接路径上采用FF-Parser，实现多尺度集成。我们在三个不同的医学分割任务上验证了MedSegDiff，眼球分割，脑肿瘤分割和甲状腺结节分割。这些任务的图像有不同的形态，分别是眼底图像、脑CT图像、超声图像。MedSegDiff在不同模式下的所有三个任务上都优于之前的SOTA，这表明了所提出的方法的泛化和有效性。简而言之，本文的贡献是:

1. 第一个提出基于DPM的医学图像生成模型。
2. 针对step-wise attention.问题提出动态条件编码策略
3. 提出FF-Parser用于消除高频噪声
4. 在数据集上表现优秀。

理论

我们基于[4]DDPM中提到的扩散模型来设计我们的模型。扩散模型是由正向扩散阶段和反向扩散阶段两个阶段组成的生成模型。在正向过程中，对分割标签X0通过一系列增加噪声的步骤T逐渐添加高斯噪声。在反向过程中，通过反转噪声过程来训练神经网络恢复原始数据。

1. 动态条件编码

在大多数条件DPM中，条件先验将是唯一的给定信息。然而，医学图像分割因目标模糊而臭名昭著。病变或组织通常很难与其背景相区分，尤其是低分辨率的图像，比如MRI。只给定一张静态图片I作为先验信息，每一步是很难学习的。为了解决这个问题，作者为每一步提出了dynamic condition encoding。我们注意到，一方面，原始图像包含了精确的分割目标信息，但很难从背景中区分出来，另一方面，当前步分割图包含了增强的目标区域，但不准确。这促使我们将当前步长分割信息Xt整合到条件原始图像编码中进行互补。我们在特征级别进行集成。在原始图像编码器中，我们利用当前步编码特征增强了原始图像编码器的中间特征。每个尺度的conditional feature map 都和相同特征的Xt编码特征融合。该操作应用于中间两个阶段，其中每个阶段都是ResNet34之后实现的卷积阶段。这种策略有助于MedSegDiff动态地定位和校准分割。尽管该策略是有效的，但另一个具体的问题是， Xt嵌入将导致额外的高频噪声。为了解决这个问题，我们提出了FF-Parser来约束特征中的高频噪声。

B、FF-Parser

主要是通过傅里叶变化抑制Xt中带进来的高频噪音，构建了一个和m同样大小的可学习矩阵A。

C、训练和架构

该网络使用DDPM同样的训练方式。唯一不同的是损失函数，

在每一轮的迭代中，标签和原始图像被随机抽取一组进行训练。迭代数字是随机从均匀分布中得到的。MedSegDiff的主要架构是一个经过改进的Re- sUNet[11]，我们在一个UNet解码器之后使用一个ResNet编码器来实现它。具体网络设置参照[12]。I和X t用两个单独的编码encoder。

编码器由三个卷积阶段组成。每个阶段包含几个残差块。每一阶段的剩余块数按照ResNet34设置。每个残差块由两个卷积块组成，每个卷积块由组范数和SiLU[13]活动层和一个卷积层组成。残差块通过线性层、SiLU激活和另一线性层接受时间嵌入。