【医学分割】基于显式形状先验学习的医学图像分割

Learning With Explicit Shape Priors for Medical  Image Segmentation

        医学图像分割是医学图像分析和手术规划的一项基本任务。近年来，基于 UNet 的网络在医学图像分割领域大行其道。然而，卷积神经网络（CNN）存在感受野有限的问题，无法模拟器官或肿瘤的长程依赖性。此外，这些模型严重依赖于最终分割头的训练。而现有的方法无法同时很好地解决上述局限性。因此，在本文的工作中，提出了一个新颖的形状先验模块（SPM），它可以明确地引入形状先验来提高基于 UNet 模型的分割性能。显式形状先验包括全局形状先验和局部形状先验。前者具有粗略的形状表示，可为网络提供全局建模能力。后者具有更精细的形状信息，可作为额外的指导，减轻分割头对可学习原型的严重依赖。为了评估 SPM 的有效性，在三个具有挑战性的公共数据集上进行了实验。提出的模型达到了最先进的性能。此外，SPM 可以作为一种即插即用的结构，融入经典的 CNN 和基于变换器的骨干网络，从而为不同数据集上的分割任务提供便利。

INTRODUCTION

Limitations

        然而，目前基于 UNet 的医学图像分割模型存在以下局限性1 由于卷积核的固有特性，CNN 的感受野有限，无法利用器官或组织之间的远距离和全局空间关系，因此无法实现精细的形状表示。因此，一些隐式注意力模块被用来扩大模型的感受野，在本文中被称为隐式形状模型。2⃝ 分割掩码主要基于由分割头解释的最终可学习原型的训练。具体来说，考虑到一个有 N 个语义类别的分割任务，需要学习 N 个类别原型来进行像素分类。每个类别只学习到一个可学习的原型，这就使用了有限的表示能力，因此不足以描述丰富的类内差异。在这种情况下，基于 UNet 的模型在提取器官或肿瘤的精确形状信息方面面临重大挑战。特定的损失函数旨在将明确的形状先验或解剖约束整合到分割框架中，而不是 Dice 损失或交叉熵损失，它们可以提取与感兴趣区域相关的足够的结构信息，包括形状和拓扑结构。然而，这些损失函数是针对特定任务的，无法在不同数据集上轻松扩展。此外，有人提出了显式形状模型，以增强模型的形状表征能力，并将形状先验作为额外的输入。

Implicit Shape Models

        为了解决 CNN 受限感受野的限制，之前的研究尝试为 U 形结构引入隐式解剖形状先验，即隐式形状模型。这些形状先验可以通过隐式注意力模块注入网络结构。从理论上讲，引入注意力模块 M 是为了加强编码器 E 生成的深度特征，使其更加专注于具有特定形状的前景区域。然后，解码器 D 将增强的特征进行合并，以获得相应图像 X 的更精确的掩码 Y。

        更具体地说，注意力模块可分为两类。第一类属于基于卷积的注意力模块。BB-UNet 通过在训练前生成的边界框（BB）过滤器来增强跳过的特征。虽然边界框滤波器可以提供特定器官的形状信息，但获取边界框滤波器需要人工干预。此外，Attention UNet 采用注意力门（AGs）来增强对特定任务有益的突出特征。AGs 还能抑制来自无关区域的冗余特征激活，这些区域在一定程度上具有形状先验性。SE 块通过对通道之间的相互依赖性建模，自适应地重新校准通道特征响应。CBAM 采用通道注意模块和空间注意模块来提高对特定区域形状的表征能力。然而，基于卷积的注意模块堆叠无法有效扩大有效感受野。而且，它们对长期依赖性建模的能力仍然有限。

        与此不同，第二类注意力模块是基于自我注意力机制，它通过查询、关键和值向量为全局上下文建模提供了一种可行的方法。许多基于 Transformer 的模型被提出来，其中包含各种类型的自我注意，用于对医学图像的长距离依赖性建模。TransUNet  将二维 UNet 与预先训练好的视觉变换器 (ViT) 结合起来，通过叠加每个切片的预测值来解决体积图像分割问题。SwinUNETR 采用基于移动窗口的注意力来提取三维斑块的特征，然后通过残差卷积块合并多尺度编码特征，以获得最终的掩码。然而，与 CNN 不同的是，这些基于 Transformer 的模型需要大量数据资源进行训练，无法简单、精细地学习数据源内部的形状先验信息等归纳偏差。

Explicit Shape Models

        为了减轻对最终可学习原型训练的严重依赖，先验方法试图在分割框架中引入额外的形状信息，称之为显式形状先验。与上文提到的隐式形状先验不同，显式形状先验具有很强的可解释性，能为感兴趣区域（ROI）提供粗略的定位。这些研究可分为两类，包括基于图集的模型和以高斯混合模型（GMM）为代表的统计形状模型。

        具体来说，第一种范式基于图集，其本质是通过源图像和目标图像之间的配准变换进行标签传播。然后将此变换应用于源地面真值（GT），就能获得目标图像的地面真值。显然，由于数据源和成像噪声的限制，非刚性配准无法完美地处理分割任务。因此，更可行的办法是以非局部方式建立匹配关系。此外，采用源图像中一组候选图像的加权组合（称为配准基础）来获得更精细的分割掩码也是有益的。配准基的分割掩码可作为形状先验，促进目标图像的分割。整个模型可以用公式：

其中，Xtrain 和 Ytrain、Xtest 和 Ytest 分别代表训练数据源和测试数据源中的图像和 GT，Xb 和 Yb 代表包含 m 对源图像和 GT 的配准基组，ωi 是加权系数。此外，R 指的是 X i b 和 Xtest 之间的配准变换。T 是变换矩阵，用于对 Yi b 注册基中的每个 GT 进行注册变换。

        对于基于图集的模型，在推理过程中存在很大的计算成本。此外，注册基础的选择对模型的稳健性也很重要。具体来说，基向量应覆盖整个数据集的分布特性。然而，由于存在形状变化，采用固定的模板形状无法覆盖所有生物物体。因此，从训练数据集中以统计方式收集大量形状前验对于提高模型的分割性能和鲁棒性至关重要。

        第二种分割模式是统计形状模型。一种代表性的方法是高斯混合模型（GMM），它通过一组可学习的高斯分布完成从图像空间 I 到标签空间 L 的连续映射。这些高斯分布可被视为数据集的明确形状先验。在训练过程中，采用期望最大化（EM）算法迭代更新可学习的高斯分布和分割掩码。在推理过程中，利用学习到的形状先验作为独立核。整个模型如下式所示： 

其中，K 表示从训练过程中生成的高斯分布，Ki 表示 K 中的每个元素，N 表示高斯核的数量，也就是语义类别的数量。G 是一个映射函数，它将 n 个高斯概率值分配给每个像素，每个值都由内核 Ki 生成。然而，GMM 对噪声和动态背景仍然很敏感。此外，EM 算法的初始设置对最终解决方案至关重要。

        其他一些相关工作扩大了统计形状模型的影响。点分布模型 (PDM) 设计用于表示形状的平均几何形状，以及从形状训练集中推断出的一些几何变化统计模式。主动形状模型（ASM）是物体形状的统计模型，它通过迭代变形来适应新图像中的物体示例。形状受点分布模型的约束，以在标注示例的训练集中看到的方式变化。不过，ASM 仅使用形状约束，并不能利用所有可用信息。因此，有人提出了主动外观模型（AAM），用于将物体形状和外观的统计模型与新图像进行匹配。具体来说，该算法利用当前估计的外观和形状与目标图像之间的差异来驱动优化过程。此外，约束局部模型结合了全局形状模型和局部纹理模型，用于划分每个地标点。这种方法能有效地模拟局部区域的形状变形。然而，这些明确的形状模型无法嵌入到深度分割网络中，对未见过的数据集的泛化能力仍然较差。 

Contributions

        同时，为了解决这两个局限性，本文加入了可学习的显式形状先验，以增强基于 UNet 模型的形状表示。在物体检测领域，DETR 引入了一组可学习的物体查询，然后推理物体与全局图像上下文的关系，直接输出最终预测结果。受这一设计的启发，本文设计了可学习形状先验，它实际上是一个 N 信道向量，每个信道都包含特定区域类别的丰富形状信息。形状先验基于自我注意生成，它赋予分割模型以全局感受野。同时，可学习的形状先验可以用更丰富的形状信息来增强编码的深度特征，然后驱动网络生成更好的遮罩，这可以减轻对可学习原型的严重依赖。此外，编码特征还有助于形状先验的迭代更新。基于这一理论，提出了由自更新块（SUB）和交叉更新块（CUB）组成的形状先验模块（SPM）。

        首先，SUB 用于生成特定数据集的全局形状先验。基于自注意机制，通过计算形状先验的每个通道对之间的相似性，将可学习的形状先验全局化，以模拟类间关系。在这里，形状先验中的每个通道都对应于特定解剖结构的形状表征。因此，SUB 在模拟数据集的长程依赖性方面发挥着至关重要的作用，从而缓解了基于卷积的隐式注意力模块的缺点。

        其次，SUB 的结构缺乏还原偏差，无法对解剖学的精细形状表征进行建模。这就是为什么 CUB 被设计用来描述局部形状先验的原因。来自编码器的卷积特征具有很强的识别局部结构的能力。因此，卷积特征和全局形状先验之间的相互作用将输出具有更精细形状表征的局部先验，从而减轻基于自我注意的隐式形状增强模块所面临的收敛性挑战。同时，形状先验可以利用丰富的全局上下文（包括纹理和结构信息）来丰富卷积特征。

        第三，如图 1 所示，可学习的形状先验被明确纳入分割模型，以获得更好的性能和良好的可解释性。而且，它们能够减轻对原型学习的严重依赖。与其他显式形状模型相比，本文提出的 SPM 在不同数据集上具有更强的泛化能力。具体来说，由于固定模板无法覆盖整个数据集的分布特性，可学习的 N 通道形状先验比从训练集中选取的图集更稳健。此外，与 GMM、ASM 等统计形状模型（SSM）相比，使用 SPM 的深度分割网络对背景噪声和优化过程中的初始解不那么敏感。

METHODOLOGY 

Unified Framework for Explicit Shape Models

        如图 1 所示，本文主要讨论了三种可以提供明确形状先验的分割范式。这些范式可以统一如下：

其中，I 表示作为分割框架输入的测试图像，O 表示模型的输出。S 表示以不同方式生成的明确形状先验，用于提高分割性能。P 指联合输入的模型预测过程。D 指对生成的 N 通道预测进行单次解码，N 是分割类别的数量。

        在这项工作中，所提出的范式为 U 型神经网络引入了可学习的显式形状先验 S。具体来说，S 被用作网络与图像相结合的输入。网络的输出是由 S 生成的预测遮罩和注意力图。然后，注意力图的通道可以提供地面实况区域的丰富形状信息。显式形状先验模型可描述如下： 

其中，F 代表推理过程中的前向传播，S 代表构建图像空间 I 和标签空间 L 之间映射的连续形状先验。训练完成后，可学习的形状先验值就固定下来了，在推理过程中，随着输入斑块的变化，可动态生成精炼的形状先验值。细化后的形状前验可以作为注意力图，定位感兴趣的区域，并抑制背景区域。此外，一小部分不准确的基本事实不会对 S 的学习产生重大影响，这揭示了所提出的范式的鲁棒性。 

Shape Prior Module

        1) 概述：如图 1 所示，所提出的模型是一个分层 U 型网络，由类似 ResNet 的编码器、基于 Resblock的解码器和形状先验模块（SPM）组成。SPM 是一个即插即用的模块，可以灵活地插入到其他网络结构中，以提高分割性能。在下面的章节中，将详细介绍 SPM，包括该模块的动机、详细结构和功能。

        为了摆脱对最终可学习原型的依赖，建议在基于 UNet 的网络中引入明确的形状先验，对每个类别施加解剖形状约束，以增强网络的表示能力。受 DETR 的启发，设计了 n 个（分割类别的数量）可学习原型，类似于 DETR 变换解码器中的对象查询。如图 2 所示，SPM 的输入是原始跳过特征 Fo 和原始形状先验 So，然后再细化为增强跳过特征 Fe 和增强形状先验 Se。具体来说，可学习的形状先验会在卷积编码特征的指导下生成具有足够形状信息的精炼注意力图。同时，编码特征将通过形状先验生成更准确的分割掩码。与 DETR 不同的是，SPM 将与多尺度特征进行交互，而不仅仅是编码器瓶颈处的特征。因此，跳转连接前的分层编码特征将通过 SPM 获得更丰富的形状信息。增强形状先验由两个部分组成，即全局形状先验和局部形状先验，分别由自更新块和交叉更新块生成。

        2) 自更新区块：远距离依赖建模：由于目标是引入能够定位目标区域的显式形状先验，所以形状先验的大小为 N × 空间维度。N 指的是类的数量，空间维度与斑块大小有关。为了缓解有限感受野的缺点，本研究考虑了形状前验内部的长程依赖性。具体来说，提出了自更新块（SUB）来模拟类间关系，并生成具有 N 个通道间交互作用的全局形状前验。受视觉转换器（ViT）的自我关注机制的启发，N 个类之间的自我关注 Smap 的亲和图由如下公式构建，它描述了形状前验各通道之间的相似性和依赖关系。 

        其中，Cmap 表示交叉注意阶段的亲和图，Qc 和 Kc 表示将 Fo 和 So 投射到查询和关键向量的卷积变换。Cmap 是一个 C × N 矩阵，用于评估 C 信道特征图 Fo 和 N 信道形状先验之间的关系。卷积特征图 Fo 的特定通道与形状先验的特定通道相关。然后，Cmap 作用于转换后的全局形状先验 SG 来完善 Fo，使其具有更准确的形状特征和丰富的全局纹理。 

其中，Vc 指的是将 So 投射到值向量中的卷积变换，Fe 代表增强的跳过特征。同时，局部形状先验 SL 是由下采样的 Fe 生成的，它具有模拟局部视觉结构（边缘或角落）的特性。

        总之，原始形状前验可通过全局和局部特征得到增强。全局形状先验可以模拟类间关系，从而在自我注意区块的基础上提供具有充分全局纹理信息的粗形状先验。局部形状前验则通过引入基于卷积的归纳偏差来显示更精细的形状信息。此外，原始跳过特征通过与全局形状前验的交互作用得到进一步增强，这将促进生成具有分辨形状表征和全局上下文的特征，进而获得更精确的预测掩码。