【Diffusion分割】MedSegDiff-v1:Diffusion模型进行医学图像分割

MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

摘要：

        最近的许多研究发现，Diffusion Probabilistic Model DPM 在图像去模糊、超分辨率和异常检测等各种其他视觉任务中也很有用。受 DPM 成功经验的启发，我们提出了 MedSegDiff 模型，这是第一个基于 DPM 的通用医学图像分割任务模型。为了增强 DPM 在医学图像分割中的分步区域注意力，我们提出了动态条件编码，为每个采样步骤建立状态自适应条件。此外，我们还提出了特征频率解析器（FF-Parser），以消除这一过程中高频噪声成分的负面影响。我们在三种不同图像模式的医疗分割任务中验证了 MedSegDiff 的有效性，包括眼底图像上的视杯分割、核磁共振图像上的脑肿瘤分割和超声图像上的甲状腺结节分割。实验结果表明，MedSegDiff 在性能上远远超过了最先进的（SOTA）方法，证明了所提模型的通用性和有效性。

1. Introduction

        受最近 DPM 成功的启发，我们针对医学图像分割任务提出了一种独特的基于 DPM 的分割模型。据我们所知，我们是第一个为一般医学图像分割提出基于 DPM 的模型的人。我们注意到，在医学图像分割任务中，病变/器官往往是模糊的，很难从背景中区分出来。在这种情况下，自适应校准过程对于获得准确的结果至关重要。根据这一思路，我们在 vanilla 条件 DPM 的基础上提出了动态条件编码，设计出了名为 MedSegDiff 的模型。在迭代采样过程中，MedSegDiff 的每一步都以图像为条件，以学习分割图。为了实现自适应区域注意力，我们将当前步骤的分割图整合到每一步的图像先验编码中。具体来说，我们以多尺度的方式将当前步骤的分割掩码与图像先验的特征层融合在一起。这样，被破坏的当前步骤掩码就能动态增强条件特征，提高重建精度。为了消除损坏掩模中的高频噪声，我们提出了特征频率解析器（FF-Parser）来过滤傅立叶空间中的特征。在每个跳转连接路径上都采用 FF-解析器进行多尺度整合。 

2. Method

        按照 DPM 的标准实现方式，我们采用 UNet 作为学习网络。图 1 举例说明了这一点。为了实现分割，我们使用原始图像先验函数ε作为阶跃估计函数的条件：

        其中，EI t 是原始图像的条件特征嵌入，Ex t 是当前步骤中分割图的特征嵌入。这两个嵌入相加后通过 UNet 解码器 D 进行重构。步骤索引 t 与添加的嵌入和解码器特征相整合，并使用共享学习查找表进行嵌入，如（Ho 等人，2020 年）所述。 

2.1. Dynamic Conditional Encoding

        在大多数条件 DPM 中，条件先验将是一种独特的特征嵌入。然而，众所周知，医学图像分割因对象模糊而具有挑战性，病变或组织往往难以从背景中区分出来。而核磁共振成像或超声波图像等低对比度图像模式则进一步加剧了这一问题。为了解决这个问题，我们为每一步提出了动态条件编码。我们发现，原始图像包含准确的分割目标信息，但很难与背景区分开来。同时，当前步骤的分割图包含增强的目标区域，但并不完全准确。 

        将当前步骤的分割信息 xt 整合到条件原始图像编码中，以实现互补。具体来说，我们在特征层面上实现整合。在原始图像编码器中，我们用当前步骤的编码特征增强其中间特征。条件特征图 mk I 的每个尺度都与相同形状的下一个编码特征 mkx 融合，其中 k 是层的索引。在该机制中，我们首先对两个特征图进行层归一化处理，然后将它们相乘，得到一个亲和图。然后，我们将亲和图与条件编码特征相乘，以增强注意力区域，如下所示：

 其中，⊗ 表示元素相乘，LN 表示层归一化。我们在中间两级应用了这一操作，其中每一级都是按照 ResNet34 实现的卷积级。

        这种动态条件编码策略有助于 MedSegDiff 动态定位和校准分割。但由于整合了被噪声干扰的 xt 嵌入，它可能会引入额外的高频噪声。为了解决这个问题，我们进一步提出了 FF-Parser 来限制特征中的高频成分。

2.2. 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 全局调整特定频率的成分。因此，可以对其进行训练，以限制自适应整合的高频成分。 

2.3. Training and Architecture

        MedSegDiff 是按照 DPM 的标准流程进行训练的（Ho 等人，2020 年）。具体来说，损失可以表示为：

        在每次迭代中，随机抽取一对原始图像 Ii 和分割标签 Si 进行训练。迭代次数取自均匀分布，ε取自高斯分布。

        MedSegDiff 的主要架构是经过改进的 ResUNet（Yu 等人，2019 年），我们使用 ResNet 编码器和 UNet 解码器来实现这一架构。我们沿用了（Nichol 和 Dhariwal，2021 年）的详细网络设置。I 和 xt 都使用两个独立的编码器进行编码。每个编码器由包含多个残差块的卷积阶段组成。每个阶段的残差块数量与 ResNet34 相同。每个残差块由两个卷积块组成，每个卷积块由一个组归一化层和一个 SiLU（Elfwing 等人，2018 年）激活层组成，然后是一个卷积层。残差块通过一个线性层、一个 SiLU 激活层和另一个线性层接收时间嵌入。然后将结果添加到第一个卷积块的输出中。得到的 EI 和 Ext 相加后被发送到最后一个编码阶段。连接一个标准卷积解码器来预测最终结果。