SPM: 一种即插即用的形状先验模块!

奥比中光3D视觉开发者社区

于 2023-04-25 08:30:00 发布

阅读量819

点赞数

文章标签：人工智能深度学习计算机视觉

原文链接：https://mp.weixin.qq.com/s/vh7ntRNacGvbcX9mpsKtxQ

版权

作者：CVHub
编辑：3D视觉开发者社区

本文主要的核心点在于探索形状先验(shape priors)对分割性能的影响。

导读

按笔者的理解，这是因为在医学图像中，不同的器官或病灶通常具有特定的形状和结构，这些形状和结构信息对于分割模型来说非常关键，因此先前的许多工作尝试利用形状先验来设计分割模型，以获得具有解剖形状信息的更好掩模(mask)。换句话理解就是引入形状先验可以帮助分割模型在分割过程中更好地考虑和利用目标物体的形状信息，从而提高分割性能。

为此，本文集中探讨了三种带有形状先验的分割模型：

🔹基于图谱的模型(atlas-based models)

🔹基于统计的模型(statistical-based models)

🔹基于U-Net的模型(UNet-based models)

作者认为，前两种方法的泛化能力较差，而 UNet-based 模型由于相比于前两者泛化性能要好，但由于它是倾向于使用隐式形状先验，这在不同形状的器官上缺乏良好的可解释性和泛化能力。综上所述，本文提出了一种新的形状先验模块(Shape Prior Module, SPM)，它可以显示地引入形状先验，以促进 UNet-based 模型的分割性能。

最后，为了评估 SPM 的有效性，作者在三个具有挑战性的公共数据集上进行了实验。结果表明，作者提出的模型达到了最先进的性能水平。此外，SPM 在不同的经典 CNNs 网络和最新的基于 Transformer 的骨干网络上展现了出色的泛化能力，可作为不同数据集分割任务的即插即用结构。

so，讲到这里，可能还有小伙伴没理解什么叫“隐式”，什么又叫“显示”，下面讲下个人理解：

隐式形状先验通常是通过在模型中加入先验信息，例如特定的损失函数或正则化项来实现的。这些隐式的形状先验通常难以解释，因为它们是通过一些特殊的方式集成到模型中的，而不是直接考虑目标物体的形状信息。例如，在基于 UNet 的模型中，可以通过使用 Dice 损失函数来强制模型更加注重目标物体的轮廓信息，从而隐式地考虑了形状先验信息。

相反，显式形状先验则直接将形状先验信息作为输入提供给模型。例如，在本文中，作者提出了一个新的形状先验模块，它明确地将形状先验信息作为输入，并利用这些信息来引导模型更好地分割目标物体。这种显式的形状先验可以更好地解释和调整，因为它们直接考虑了目标物体的形状和结构信息。

方法

本文的主要贡献是在U-Net网络基础上设计了一个形状先验模块(SPM)，通过引入可学习形状先验，为每个类别施加解剖形状约束来增强网络的表示能力。与DETR类似，作者设定了 n 个（分割类别数目）可学习原型，类比于DETR中Transformer解码器中的对象查询。

SPM的输入包括原始跳跃特征和原始形状先验，经过“特征提纯”后会生成对应的增强跳跃特征和形状先验。最终，通过这些增强后的特征和先验，模型会生成更加精准的分割掩膜。

与DETR不同，SPM会与多尺度特征进行交互，而不仅仅是来自编码器最深层的特征。因此，在跳跃连接之前的分层编码特征在经过SPM处理后将获得更多的形状信息。增强形状先验由两个部分组成:

🔹全局形状先验

🔹局部形状先验

它们分别由自更新块(self-update)和交叉更新块(cross-update block)生成。下面我们就围绕这两个模块简单唠唠。

Self-update block (SUB)

为了建模类别之间的关系，本文提出了自更新块(参照上图第一行)来生成具有 N 个通道之间交互的形状先验。公式就不展开了，从流程图上可以看出，输入特征是来自编码层对应层级输出的长跳跃连接，作者首先整了两个Resblock做下特征提取并进行维度压缩。其次，便是借鉴ViT中自注意力机制(Self-Attention)的启发，通过引入来自深层的特征，构建对应的QKV。如此一来，网络便可以有效的建模远距离的上下文依赖，从而获得与全局区域相关的丰富纹理信息。

Cross-update block (CUB)

正所谓顾此失彼，试想一下，当一个人看得越“广”，其对目标的聚焦程度显然越“低”。换句话说就是，虽然SUB带来了全局上下文信息，但最关键的形状和轮廓信息缺忽略了，这是由于缺乏相应的归纳偏置来建模局部视觉结构和定位具有不同尺度的对象。

为了解决这个限制，CUB 应运而生。受到卷积核固有的局部性和尺度不变性的归纳偏置的启发，基于卷积的 CUB 为 SPM 注入归纳偏置，以获得更细致的局部形状信息。此外，基于编码器中卷积特征具有定位区分性区域的显著潜力的事实，作者尝试在原始跳跃特征和形状先验之间进行交互。

综上所述，所提出的形状先验模块通过引入全局和局部特征对原始跳跃特征和形状先验进行了增强。全局形状先验能够通过自注意力块模拟类间关系，包含粗糙的形状先验和丰富的纹理信息以建立全局上下文。局部形状先验通过卷积引入归纳偏差，从而能够捕捉具有不同尺度的局部视觉结构和定位目标。此外，SPM还通过与全局形状先验的交互来增强原始跳跃特征，从而促进生成具有较好形状表示和全局上下文的特征，最终获得更准确的预测分割掩模。

实验

性能指标

下面分别展示所提方法与其他主流的分割算法在BraTS 2020、VerSe 2019和ACDC的表现：

超过的部门就不讲了哈，问就是因为引入了“形状先验模块”带来的性能提升。直接看上面这张图，nnUnet和nnFormer在Myo指标上表现更加优异，这是因为ACDC这个MRI数据集具备较大的体素空间，这将加剧标签不平衡的影响，而这两者由于采用有效的重采样策略将改善心肌组织的不平衡分布，从而使模型对该区域的关注度更高。

效果展示

图 4：不同模型在 BraTS 2020 上的预测分割掩码

图 5：VerSe 2019 上不同模型的预测分割掩码

图 6：不同模型在 ACDC 上的预测分割掩码

可视化分析

本文集中探讨了形状先验和跳跃特征的这两部分的定性结果，显然，这两者是可以互补的，这在许多的语义分割方法中也多次提及到，感兴趣的小伙伴也可以关注微信公众号 CVHub 添加小编微信 cv_huber 领取语义分割大礼包，笔者已对整个技术脉络进行梳理和深刻剖析。

上图展示了跳跃特征对明确形状先验的影响。其中:

案例（a）展示了从不同阶段生成的明确形状先验。具体来说，形状先验由 N 个通道注意力图组成，其中 N 是分割类别的数量，每行表示来自每个阶段的形状先验。我们可以发现，随着自上而下的过程，形状先验对于地面真实区域呈现出更准确的激活图。特别是，在第一阶段中错误激活的区域将在 SPM 的第二和第三阶段中被抑制。在可视化结果中，存在一种称为反向激活的现象，这意味着除了 GT 区域之外的所有区域都被激活。

案例（b）中则展示了形状先验的最后一个阶段和最后一个通道的典型例子。作者声称，这种现象是由全局形状先验造成的，它为整个区域带来了全局上下文和丰富的纹理信息，甚至包括背景区域。实质上，通过反向注意力简单地定位 ROI，其中 ROI 用清晰的轮廓突出显示。从这个角度来看，反向激活类似于正向激活。

接下来，我们一起将形状先验分解为来自 SUB 和 CUB 的两个组成部分，即全局形状先验和局部形状先验：

从图8可以观察到，得益于自注意力模块，全局形状先验具有全局的感受野，包含上下文和纹理。然而，SUB 的结构缺乏归纳偏差来模拟局部视觉结构。在这里，我们可以发现，全局形状先验负责对 GT 区域进行粗定位。而由 CUB 生成的局部形状先验可以通过引入卷积核提供更精细的形状信息，这些卷积核具有局部归纳偏差。

文章的最后，我们再来讨论下形状先验对跳跃特征的影响。上图可视化比较了原始跳跃特征和增强跳跃特征之间的差异。显而易见，增强的跳跃特征通过SPM的处理，强化了整个肿瘤区域，同时通过全局形状先验的引入，为整个区域提供了丰富的纹理信息。最后，通过对心脏 CT 图像的案例解释了特征细化的过程。

总结

本文详细讨论了三种具有形状先验的分割模型，包括基于图谱的模型、统计形状模型和基于 UNet 的模型。为了增强 UNet-based 模型上形状先验的可解释性，本文提出了一个形状先验模块 SPM ，可以明确地引入形状先验以促进不同数据集上的分割性能。所提方法在 BraTS 2020、VerSe 2019 和 ACDC 三个主流的数据集上均实现了最先进的性能。此外，根据定量和定性实验结果，SPM 在不同的骨干网络上表现出良好的泛化能力，可作为一种即插即用的结构。