FDiff-Fusion: Denoising diffusion fusion network based on fuzzy learning for 3D medical image segmentation

摘要：

        近年来，去噪扩散模型在图像分割建模中取得了令人瞩目的成就。凭借其强大的非线性建模能力和优越的泛化性能，去噪扩散模型逐渐被应用到医学图像分割任务中，为该领域带来了新的视角和方法。然而，现有方法忽略了分割边界的不确定性和区域的模糊性，导致分割结果不稳定、不准确。为解决这一问题，本文提出了一种基于模糊学习的三维医学图像分割去噪扩散融合网络（FDiff-Fusion）。通过将去噪扩散模型集成到经典的 UNet 网络中，该模型可以有效地从输入的医学图像中提取丰富的语义信息，从而为医学图像分割提供出色的像素级表示。由于医学图像分割普遍存在边界不确定和区域模糊的问题，本文在 U-Net 网络的跳转路径上设计了一个模糊学习模块。该模块为输入的编码特征设置多个模糊成员函数来描述特征点之间的相似度，并对模糊成员函数应用模糊规则，从而增强了模型对不确定边界和模糊区域的建模能力。此外，为了提高模型分割结果的准确性和鲁棒性，我们在测试阶段引入了迭代注意力特征融合方法，在注意力模块的全局上下文信息中加入局部上下文信息，对每个去噪时间步的预测结果进行融合。最后，为了验证 FDiff-Fusion 的有效性，我们在 BRATS 2020 脑肿瘤数据集和 BTCV 腹部多器官数据集上将其与现有的高级分割网络进行了比较。结果表明，在这两个数据集上，FDiff-Fusion 显著提高了 Dice 分数和 HD95 距离，证明了它在医学图像分割任务中的优越性。

1. Introduction：

        在医学图像处理领域，去噪一直是一个基本而关键的问题。传统的去噪方法通常依赖于滤波器或变换域处理，如小波变换、均值滤波和高斯滤波。然而，这些方法在处理复杂噪声和保留图像细节方面效果有限 。近年来，随着深度学习技术的快速发展，去噪扩散模型逐渐受到广泛关注。其基本原理是利用深度神经网络模型，通常是 U-Net 网络模型，通过训练将含有噪声的输入图像映射到相应的干净图像空间。与传统的医学图像分割方法相比，去噪扩散模型具有更强的非线性建模能力和更优越的泛化性能，能更有效地捕捉图像中的复杂结构和细节信息。在 Wu 等人提出的 MedSegDiff 中，为了增强扩散概率模型（DPM）在医学图像分割中的分步区域注意力，他们提出了动态条件编码，为每个采样步骤建立状态自适应条件。此外，他们还进一步提出了特征频率解析器，以消除过程中高频噪声成分的负面影响。在 MedSegDiff 的基础上，Wu 等人提出了一种新颖的基于变换器的条件 U-Net 框架和一种新的频谱空间变换器，以模拟噪声和语义特征之间的相互作用。这种改进的架构产生了一种新的基于扩散的医学图像分割方法，称为 MedSegDiff-V2 [24]，大大提高了 MedSegDiff 的分割性能。Xing 等人[25]提出了用于医学体素分割的 Diff-UNet 方法。他们将扩散模型集成到标准的 U 型结构中，有效地提取了输入体素的语义信息，从而为医学体素分割提供了出色的像素级表示。为了增强扩散模型预测结果的鲁棒性，他们还在推理过程中引入了基于逐步不确定性的融合模块，以合并扩散模型在每一步的输出结果。Wolleb 等人[26]利用扩散模型来解决二维医学图像分割问题，在测试过程中，他们通过将每一步扩散的输出结果相加进行融合来增强分割的鲁棒性。 

        尽管去噪扩散模型在图像分割领域取得了巨大成功，但医学图像分割仍存在一些挑战和问题有待解决。例如，医学图像的个体差异较大，经常会受到各种形式的噪声干扰，包括来自电子设备的信号干扰、成像过程中的运动伪影以及基于物理特性的噪声。此外，医学图像中的分割边界不确定和分割目标区域模糊现象也很常见。上述模型难以有效处理医学图像中不确定的分割边界和模糊区域。

        针对这些问题，我们提出了一种基于模糊学习的去噪扩散融合网络，用于三维医学图像分割，命名为 FDiff-Fusion。这种分割方法通过引入模糊学习模块（FLM），增强了 U-Net 网络对医学图像中不确定分割边界和模糊区域的建模能力。此外，还设计了一个基于迭代注意机制（IAF）的融合模块，用于融合去噪过程中多个时间步骤的预测结果，从而在测试过程中获得更准确的分割结果。所提出的方法旨在提高医学图像分割的准确性和鲁棒性，更有效地处理边界不确定和区域模糊的情况。这对提高医学图像分割技术和临床诊断的准确性至关重要。

        1.针对医学图像分割中分割边界的不确定性和区域的模糊性，我们引入了模糊学习模块。该模块为输入编码特征设置多个模糊成员函数，以描述特征点之间的相似性。通过对这些模糊成员函数应用模糊规则，增强了模型对分割边界不确定性和区域模糊性的建模能力。

        2.为了提高模型预测结果的准确性和鲁棒性，我们提出了一个基于迭代注意力机制的融合模块。该模块将每个时间步骤分割图像的局部上下文信息纳入注意力模块的全局上下文信息中，从而优化最终的分割结果。

2. Preliminaries：

2.2. Fuzzy learning theory：

        模糊学习是一种基于模糊集理论的机器学习方法，旨在处理数据中的模糊性或不确定性[27]。模糊集理论用于描述事物之间的模糊关系，并为不确定性建模。以下是模糊学习的一些关键概念和方法： 

        Fuzzy Set Theory：模糊集合理论引入了模糊性的概念，用来描述事物之间的不确定性和模糊性。在这一理论中，一个元素可以部分属于一个集合，而不是完全属于或不属于，这主要通过成员函数来表示。

        Fuzzy Inference:模糊推理通过引入模糊集和模糊逻辑运算来处理模糊问题。模糊集是元素的成员度在 0 和 1 之间的函数，表示元素属于该集的程度。模糊逻辑运算，包括模糊与、模糊或、模糊非运算，用于处理模糊命题之间的关系。其基本思想是根据模糊规则进行推理，其中模糊规则由条件部分和结论部分组成。条件部分是一组模糊命题，结论部分是相应的模糊结论。模糊推理引入了一个模糊推理引擎，它根据模糊规则和输入的模糊命题进行推理，从而得到模糊结论 。

        Edge Detection:模糊学习也可用于医学图像的边缘检测任务。传统的医学图像分割算法往往难以有效处理模糊边界或复杂背景情况，而模糊学习方法可以更好地识别模糊边缘，提高边缘检测的准确性。 

        Fuzzy Membership Functions:在一些传统的图像分割算法中，可以引入模糊成员函数来描述像素与不同组织或结构之间的成员关系。这些模糊成员函数可以通过模糊学习方法进行学习和优化，从而更准确地描述图像中的模糊性和不确定性。

2.3. Denoising diffusion model：

        去噪扩散模型包括一个正向过程和一个反向过程。在正向过程中，噪声会逐渐施加到图像上，直到图像被破坏成完全高斯噪声的图像。然后，在反向过程中，使用深度神经网络模型来学习从高斯噪声图像中还原原始图像的过程 。具体来说，如图 1 所示，实线表示反向过程，虚线表示正向过程。在正向过程中，高斯噪声被逐渐添加到原始图像 x0 中，每一步得到的图像 xt 只与上一步得到的图像 xt-1 有关，直到 T 步的图像 xT 变成纯高斯噪声图像，这可以看作是一个马尔可夫过程。反向过程则是逐步去除高斯噪声，从给定的高斯噪声图像 xT 开始，一步一步进行，直到恢复到原始图像 x0。具体来说，在前向过程中，对于原始图像 x0，在马尔可夫过程的 T 个时间步上进行变分推理，以学习训练数据分布。这个过程可以表示为 ：

其中，t 代表马尔可夫链的第 t 个时间步长。αt = 1 - βt 是噪声控制器，用于控制第 t 个时间步中添加的噪声强度，且 αt ∈ (0, 1) 和 α1 > ⋯ > αT 。在反向过程中，深度神经网络（通常是 U-Net 网络）经过训练，学会去除添加的高斯噪声，从而将纯噪声图像 xT 还原为原始图像 x0。这一过程可表示为：

 

其中，pθ(⋅) 是深度学习模型，θ 是其参数。反向处理的主要目的是通过去除添加的高斯噪声来获得清晰的分割结果。

 

3. FDiff-Fusion 

        传统的医学影像分割方法通常是直接将医学影像作为输入，通过训练模型来预测相应的分割标签图。然而，与这种直接输入方法不同，去噪扩散模型采用的是另一种策略。去噪扩散模型的学习过程并不直接预测分割标签图，而是侧重于学习高斯噪声的去除过程。具体来说，在训练阶段，扩散模型将医疗图像和带有高斯噪声的分割标签图作为输入，然后学习使用深度学习模型去除医疗图像中的高斯噪声，从而生成清晰的分割结果。这种方法的关键在于通过去噪过程提高分割结果的质量，而不是直接对医学图像进行分割和预测。然而，尽管去噪扩散模型在图像分割领域取得了不俗的成绩，但在处理医学图像中的分割边界不确定性和区域模糊问题上仍存在挑战。为了克服这些挑战，提高分割结果的准确性和鲁棒性，我们提出了一种基于模糊学习的三维医学图像分割去噪扩散融合网络，命名为 FDiff-Fusion，如图 2 所示。

 3.1. FDiff-Fusion model

        FDiff-Fusion 模型包括两个过程，即正向过程和反向过程。在正向过程中，原始医学图像 x0∈R C×D×H×W 逐渐被高斯噪声污染，得到纯噪声图像 xT∈R C×D×H×W 。

        由于前向过程中的图像 xt 只取决于前一步 xt-1，因此这一过程可视为马尔可夫过程，满足以下条件：

        随后，利用重参数技术，我们通过公式（3）对 xt 进行采样，即生成一个随机变量 εt ∼ N(0, 1)、 αt = 1 - βt，由此可得:

详细推导过程可看这篇一文讲透AI绘画的Diffusion技术原理：从数学到编程（2） - 知乎 (zhihu.com) 

        如图 2 所示，在测试阶段，提供三维医学图像 I∈R D×H×W×C 作为去噪扩散模型每个时间步的条件嵌入，其中 D、H、W 和 C 分别代表三维医学图像的深度、高度、宽度和通道。首先，在三维医学图像 I 中加入 T 步高斯噪声，得到纯噪声图像 xT。然后，对纯噪声图像 xT 进行分步去噪操作。以第 t 个时间步为例，将噪声图像 xt 和原始医学图像 I 沿信道串联，并输入 U-Net 网络的编码器，得到多尺度特征 ft∈ R C×D×H×W 。同时，为了更好地将原始三维医学图像作为条件嵌入，将 I 输入深度卷积神经网络（CNN），以获得多尺度条件特征 fI∈ R C×D×H×W 。由于 ft 和 fI 包含相同数量和大小的特征，因此将相应尺度的特征相加以获得融合特征。然而，由于三维医学图像中存在分割目标边界不确定和区域模糊的问题，如果简单地将 U-Net 网络编码器提取的多尺度特征和 CNN 提取的多尺度特征相加作为每个时间步的编码特征，将难以解决分割目标边界不确定和区域模糊的问题。为解决这一问题，我们提出了模糊学习模块（FLM）。一方面，原始三维医学图像包含精确的分割目标，但很难确定其分割边界。另一方面，当前时间步的分割图包含增强的分割目标，但缺乏准确性。因此，我们在 U-Net 网络的编码器和解码器之间的连接路径上使用模糊学习模块（FLM），对编码特征进行模糊特征处理。随后，U-Net 编码器处理过的多尺度特征通过跳转连接输入解码器进行处理，从而得到第 t 个时间步的预测结果 xt-1。最后，为了提高模型分割的准确性和鲁棒性，通过迭代关注将每个时间步得到的分割图进行融合，得到最终的分割结果 ̂ x0 ∈ R C×D×H×W 。

        与测试阶段只有三维医学图像作为模型输入不同，本文在训练阶段同时将三维医学图像 I ∈ R D×H×W ×C 及其对应的分割标签 x0 ∈ R C×D×H×W 输入模型。在这里，三维医学图像作为条件信息嵌入，而相应的分割标签则用于添加高斯噪声 T 步，从而得到高斯噪声图像 xT ∈ RC×D×H×W。经过 T 步去噪操作后，模型预测出分割标签图像 ̂ x0 ∈ RC×D×H×W，然后将其与分割标签图像一起用于损失计算。FDiff-Fusion 结合了传统图像分割网络中常用的均方误差（MSE）损失、二元交叉熵（BCE）损失和骰子损失进行训练。因此，该模型的总损失 Ltotal 定义为：

3.2. Fuzzy learning module

        为解决三维医学图像中的边界不确定性和区域模糊问题，我们提出了模糊学习模块（FLM）来解决这一问题。具体来说，针对 U-Net 网络提取的特征，沿着 U-Net 网络的跳转路径设计一个模糊残差连接进行处理。模糊残差连接分为模糊学习模块和残差连接，其中模糊学习模块包括模糊成员函数和模糊规则。FLM的目标是模拟特征图与每个像素的语义类别之间的复杂规则。如图 3 所示，对于模糊学习模块输入编码特征的每个通道，对特征图中的每个特征点应用 M 个模糊成员函数，以便将编码器提取的特征图转换为模糊值，模糊值为浮点数。每个通道的模糊成员函数保持不变，但不同通道的模糊成员函数可能不同。每个模糊成员函数为每个通道中的每个特征点分配一个模糊类别标签，描述其与其他特征点的相似性。模糊成员函数层采用高斯函数作为模糊成员函数： 

        其中，Fx,y,z,c 表示特征图谱通道 C 中的特征点，(x,y,z) 表示特征点的相应位置.μk,c和σk,c是第k个高斯成员函数的均值和标准偏差，μk,c和σk,c是随机初始化和可学习的，F′x,y,z,k,c表示通道C中位置(x,y,z)处特征点的第k个输出模糊类别标签。随后，对每个特征点的所有模糊类别标签应用模糊规则 "AND"，以获得编码特征的最终模糊特征信息 F ′ x,y,z,c： 

        式（18）根据每个模糊成员函数的输出应用模糊规则 "AND"，得到每个特征点的模糊特征信息 F′ x,y,z,c。因此，U-Net 编码器提取的特征图经模糊学习模块处理后，可得到模糊特征图 F∈ R C×D×H×W。 

         作为网络的一部分，模糊学习模块可被视为由高斯成员函数的多个参数描述的模糊层。值得注意的是，参数 μk,c 和 σk,c 在同一通道中保持不变，但在不同通道中可能会有所不同。这是因为同一通道中的特征点是通过相同的卷积核提取的，具有相似的特征，而不同通道中的特征点是通过不同的卷积核获得的，具有不同的特征。此外，为了实现模糊学习模块，模糊成员函数的数量（用 M 表示）是唯一需要预设的超参数。

         最后，如图 3 左侧所示，模糊学习模块的输出是一个模糊张量 F，它与 U-Net 处理的特征张量 ft 大小相同。因此，如公式（19）所示，两者通过残差连接进行融合。

        其中，U-Net 和模糊张量的输出特征图经过批量归一化（BN）处理，以限制其动态范围。随后，采用简单的加法运算来整合模糊逻辑信息。

        通过使用高斯成员函数对编码特征进行模糊化，该模块能够使用非排他性的模糊成员值来处理医学图像中的噪声和不确定性，这比确定性特征更加稳健。

3.3. Iterative attention fusion 

        在传统的医学图像分割任务中，最终预测结果被用作模型的最终分割结果。在 FDiff-Fusion 的反向过程中，每个时间步都会产生一个预测结果。然而，每个时间步骤的预测结果在处理边界和细节方面可能会有所不同，而且可能包含不同程度的不确定性和噪声。因此，为了提高分割模型的准确性和鲁棒性，我们通过注意力机制将反向过程中每个时间步获得的分割图进行整合，从而得到最终的分割结果。为此，我们设计了一个基于迭代注意机制的融合模块，以优化模型的预测结果。传统的融合方法通常采用简单的加法或加权融合。相比之下，我们引入了迭代注意机制，以实现对逆向过程中获得的每个分割图的深度融合。

3.3.1. Multi-scale channel attention fusion

        多尺度通道注意力模块分为局部注意力和全局注意力两个分支。其主要功能是通过局部和全局特征加权来增强输入特征的表征能力。具体来说，MS-CAM 应用局部和全局注意力机制处理输入特征，生成注意力权重，然后使用这些权重对输入特征进行加权。如图 4 所示，在局部注意力分支中，给定具有 C 个通道的特征 x∈R C×D×H×W 后，将应用两层 1 × 1 × 1 卷积（Conv3d）和批归一化（BatchNorm3d）以及 ReLU 激活函数来处理输入特征。局部上下文 L(x) ∈ R C×D×H×W 的计算方法如下： 

其中，δ 表示 ReLU 激活函数，B 表示批量归一化。PWConv1 和 PWConv2 的核心大小分别为 c r ×C×1×1 和 C × C r × 1 × 1。 该分支捕捉局部空间信息，生成局部特征。 

        值得注意的是，L(x) 的形状与输入特征相同，从而保留并强调了低级特征中的微妙细节。同样，在全局注意力分支中，输入特征 x 通过全局平均池化层（GAP）在空间维度上被压缩为 1 × 1 × 1。然后，使用两层 1 × 1 × 1 卷积、批量归一化和 ReLU 激活函数来提取全局上下文 G(x) ∈ R C×D×H×W：

其中，g(x) 表示全局平均集合（GAP）。该分支捕捉全局空间信息，生成全局特征。给定全局上下文 G(x) 和局部上下文 L(x)，MS-CAM 的细化特征 x′∈ RC×D×H×W 可由下式求得：

 

其中，M(x) ∈ RC×D×H×W 表示 MS-CAM 生成的注意力权重，σ 为西格莫函数，⊕ 表示广播加法，⊗ 表示元素乘法。

3.3.2. Attention fusion 

        注意力融合采用多尺度通道注意力机制对输入特征进行加权，然后对加权特征进行元素加法，生成融合特征张量。这一过程可使模型自适应地关注输入特征中的重要信息，从而提高特征表示能力和模型的整体性能。 

        基于多尺度通道注意力模块，注意力融合可表示为：

其中，Z∈ R C×D×H×W 表示融合特征，⊕ 表示特征的逐元素求和，⊗ 表示逐元素相乘。如图 5(a)所示，AF 首先对两个输入张量进行逐元素求和，然后将其输入 MS-CAM，MS-CAM 模块的输出为注意力权重 M(X ⊕ Y)，虚线表示 1 - M(X ⊕ Y)。随后，两个输入张量分别与相应的注意力权重逐元素相乘，得到注意力加权输出。 

         注意力权重的范围介于 0 和 1 之间，1-M(X ⊕ Y) 的范围也是如此，从而允许网络在输入之间进行软选择或加权平均。

3.3.3. Iterative attention fusion

        进行简单的元素求对于输入特征 X 和 Y，对输入特征和可能会影响最终的融合权重。由于这仍然是一个特征融合问题，因此需要使用另一个注意力模块来融合输入特征。首先，对两个输入张量进行注意力融合，得到初始融合张量 X ⊕ Y。 

        然后，将初始融合张量输入 MS-CAM 模块，进行局部和全局注意力处理，生成注意力权重 M（X ⊕ Y）。最后，将两个输入张量与注意力权重逐元素相乘，并将逐元素相乘的结果逐元素相加，得到最终的融合张量 Z。如图 5(b) 所示，这种两阶段方法被称为迭代注意力融合。迭代注意力融合可表示如下：

 

5. Conclusion and outlook 

        虽然 FDiff-Fusion 在医学影像分割领域取得了巨大成功，但仍有许多方向需要进一步探索和改进。例如，我们的模型由于引入了复杂的去噪扩散模型、模糊学习模块和迭代注意力融合，在训练过程中需要大量的计算资源。这可能会限制其在计算资源有限的环境中的适用性。此外，我们的模型在训练过程中依赖于高质量的注释数据，这意味着在数据稀缺或质量不高的情况下，其性能可能会受到影响。未来的工作可以扩展到更多类型的医学图像任务，如肿瘤分割、器官分割和病变检测等。此外，我们还认为集合学习、半监督学习和无监督学习也可以作为未来研究工作的潜在方向，以进一步提高我们方法的性能和适用性。