点击下方卡片，关注 「AI视界引擎」 公众号
医学图像分割越来越依赖于深度学习技术，然而，具有前景的性能通常伴随着高昂的标注成本。本文介绍了Weak-Mamba-UNet，这是一个创新性的弱监督学习（WSL）框架，它利用了卷积神经网络（CNN）、视觉 Transformer （ViT）以及前沿的视觉Mamba（VMamba）架构进行医学图像分割，特别是在处理基于涂鸦的标注时。
所提出的WSL策略结合了三种不同的架构，但都是对称的编码器-解码器网络：基于CNN的UNet用于详细局部特征提取，基于Swin Transformer的SwinUNet用于全面的全局上下文理解，以及基于VMamba的Mamba-UNet用于有效的长距离依赖建模。该框架的关键概念是一种协作与交叉监督机制，它使用伪标签以促进跨网络的迭代学习和细化。
Weak-Mamba-UNet的有效性在带有处理过的涂鸦标注的公开MRI心脏分割数据集上得到了验证，其性能超过了仅使用UNet或SwinUNet的类似WSL框架。这突出了在稀疏或不精确标注的场景中的潜力。源代码已公开可用。
1 Introduction
医学图像分割对于医疗图像分析及有效治疗计划制定具有重要意义，用于医疗目的的深度学习网络，例如UNet。UNet以其对称的U型编码器-解码器架构和积分跳跃连接而闻名，已成为基础的分割 Backbone 网络。这些跳跃连接有效地保留了关键的空间信息，将编码器和解码器层的特征合并，以提高网络的性能。编码器降低输入以提取高级特征，解码器则利用这些特征来重建图像，从而提高分割性能。UNet的进步导致了各种增强型网络的产生，这些网络旨在解决CT和MRI扫描中复杂解剖结构的分割问题。
近期的发展已经引入了创新架构，如Transformer和Mamba，这两种架构都擅长捕捉全局上下文信息。Transformer通过多头的自注意力机制实现这一点，而Mamba则以其基于状态空间模型（SSM）的计算效率著称。这些架构已被应用于一系列计算机视觉任务中，导致了如Vision Transformer，Swin Transformer，nnFormer，ScribFormer，以及UNetr等基于Transformer的发展，以及基于Mamba网络的Vision Mamba，UMamba，Segmamba，MambaUNet，VM-UNet，和Semi-MambaUNet等的发展。
深度学习方法的有效性往往取决于大量准确标注数据集的可用性，这在医学图像分析领域可能难以获得。为了解决获取如像素级分割 Mask 这类详细标注所涉及的高成本和时间问题，研究转向了半监督学习（SSL）和弱监督学习（WSL）。SSL集中于使用小规模的像素级标注数据来训练网络，而WSL则采用更简单的标注形式，如边界框、勾选标记和点，为在有限信号监督下训练分割网络提供了一种可行方法。在这些方法中，基于涂鸦的标注因其对专家的高效率和便捷性而特别受到关注，简化了标注过程，同时未显著降低监督的质量。图1展示了MRI扫描的例子，传统的密集标注和基于涂鸦的标注。
紧随Transformer和Mamba架构在计算机视觉任务中的近期成功，以及对有限标注数据的关注，本文提出了Weak-Mamba-UNet。所提出的弱监督学习框架（WSL）在一个专为基于涂鸦的医疗图像分割设计的多视角交叉监督学习方案中整合了卷积、Transformer和Mamba架构。据作者所知，这是首次尝试利用Mamba架构对带有涂鸦标注的医疗图像进行分割。Weak-Mamba-UNet的贡献主要有三个方面：
- 基于Mamba的分割网络与WSL结合用于基于涂鸦标注的医疗图像分割的整合。
- 开发一种新颖的多视图交叉监督框架，该框架能够在有限信号监督的条件下，实现三种不同架构：CNN，ViT和Mamba的协同操作。
- 在公开可用的预处理数据集上，对Weak-Mamba-UNet进行的基于涂鸦实验演示，展示了Mamba架构提高CNN和ViT在弱监督学习（WSL）任务中性能的能力。
2 Methodology
Weak-Mamba-UNet的框架如图2所示。在本研究中，二元组表示基于涂鸦的标注训练数据集，而二元组表示密集标注的测试数据集。这里，对应于一个高度为和宽度为的2D灰度图像。涂鸦标注分别表示对应于右心室（RVC）、左心室（LVC）、心肌（MYO）、背景以及未标注像素的区域。
三个分割网络分别表示为 , 和 ，在图2中分别以绿色、蓝色和黄色突出显示。对于输入 的一个分割网络的预测表示为 ，其中 代表网络参数。这三个网络的预测可以结合起来形成一个密集的伪标签 。
整体损失由基于涂鸦的部分交叉熵损失 和密集信号伪标签骰子系数损失 组成。总的训练目标旨在最小化组合损失 ，其公式如下：
在这段文本中，"" 表示三个网络中的每一个。所有的数学符号在图2中定义。最终的评估是在测试集上衡量预测标签 与真实密集标签 之间的的一致性。
Scribble-Supervised Learning
为了解决由稀疏信号涂鸦监督所带来的挑战，作者采用了一种修改后的交叉熵（CE）函数，该函数只关注已标注的像素，而忽略 未标注 的像素。这种方法产生了一种部分监督分割损失。
具体来说，作者引入了部分交叉熵（pCE），它仅在训练网络时利用涂鸦标注，记作 。这如方程式2中所示：
其中 表示给定像素的索引，而 代表被标记了涂鸦的像素集合。变量 指示类别索引（本研究中为4）， 和 分别表示网络对于第 个像素属于第 个类别的真实值和预测概率。 被用于所有三个网络 、 和 ，并被表示为 ，其中 。
Multi-View Cross-Supervised Learning
受到Cross Pseudo Supervision（CPS），Cross Teaching和Multi-view Learning的启发，这些方法旨在促进在不同网络扰动下的连贯性正则化，作者提出的多视角交叉监督学习框架将Mamba-UNet与原始UNet和Swin UNet相结合。每个网络遵循U型编码器-解码器架构。
具体来说，UNet使用具有核的2层CNN，并执行4级下采样和上采样。Swin-UNet使用2个Swin Transformer块，而Mamba-UNet则整合了2个Visual Mamba块。SwinUNet和MambaUNet都执行3级下采样和上采样，并在ImageNet上进行预训练。这个设置引入了三种截然不同的架构视角，每种视角都单独初始化以确保观点多样性。
为了促进网络之间的相互增强，制定了一个复合伪标签 ，将稀疏标签信息转换为密集信号标签，如下方程所示：
其中 、 和 分别是平衡基于CNN的UNet、基于ViT的SwinUNet和基于Mamba的MambaUNet贡献的权重因子。这些因子在每次迭代中随机生成，并遵循 的条件，这一做法引入了一种受[27, 15]启发的数据扰动元素。这种方法确保了从每个网络中整合多样化的视角，增强了生成的伪标签的鲁棒性和泛化能力。
一旦提供了伪标签，可以通过基于Dice系数的损失 实现密集信号监督，具体表述为：
Dice损失函数（）被用于所有三个网络：、 和 ，并且分别表示为 ，其中 。
3 Experiments
数据集： Weak-Mamba-UNet以及各种 Baseline 方法的性能评估是基于一个公开可用的MRI心脏分割数据集进行的[1]。scribble标注是从原始密集标注中衍生出来的，与之前的研究保持一致。
为了在评估过程中保持一致性，所有图像都被调整到了统一的像素分辨率。实验在一个配备有Nvidia GeForce RTX 3090 GPU和Intel Core i9-10900K CPU的Ubuntu 20.04系统上进行，使用了PyTorch框架。整个实验运行平均耗时4个小时。作者对Weak-Mamba-UNet以及所有其他 Baseline 方法进行了30,000次迭代训练，每次批处理大小为24。
优化是通过随机梯度下降（SGD）进行的，初始学习率为0.01，动量设置为0.9，权重衰减为0.0001。网络在每次200次迭代后在验证集上进行评估，只在验证性能提升时保存网络权重。
Baseline 分割网络与弱监督学习框架： 图2展示了弱监督学习Mamba-UNet的框架，其中包括三个分割的主干网络。为了确保公平比较，作者还采用了图3：在假设5%的数据作为标注数据时的示例分割结果。
基于CNN的UNet 和基于Swin ViT的SwinUNet 被用作不同弱监督学习(WSL)框架的分割 Backbone 网络。所评估的WSL Baseline 框架包括部分交叉熵(pCE) ，不确定性感知的自集成与转换一致的平均教师模型(USTM) ，Mumford，门控条件随机场(Gated CRF) 。在这些框架中，SwinUNet 和 UNet 都被用作分割 Backbone 网络。
结果： 为了评估Weak-Mamba-UNet相对于其他弱监督学习（WSL） Baseline 方法的表现，作者采用了一系列全面的评价指标。在相似性度量中，数值越高表示性能越好（），作者包含了Dice系数（Dice）、准确率（Acc）、精确度（Pre）、灵敏度（Sen）和特异性（Spe）。
在差异性度量中，数值越低越理想（），作者考虑了95%的Hausdorff距离（HD）和平均表面距离（ASD）。鉴于数据集关注的是4类分割任务，作者报告了所有类别上这些指标的平均值。在ACDC数据集上，作者定量比较的结果详细列于表1中，突出了最佳性能结果的关键观察点，并用下划线进行了强调。
值得注意的是，采用SwinUNet架构的弱监督学习方法（pCE-SwinUNet和USTM-SwinUNet）通常优于基于UNet框架的方法（pCE-UNet和USTM-UNet）。例如，pCE-SwinUNet在DSC和HD上的得分分别为0.7620和54.6531，超过了pCE-UNet，这突显了在WSL框架内采用先进算法的重要性。
然而，像Weak-Mamba-UNet那样优化多个独立算法的整合，可以带来更加令人印象深刻的结果。图3通过三个说明性样本切片及其实际标签展示了作者提出方法的功效。这些示例说明了传统的pCE和USTM框架可能导致错误预测，而作者的新型多模型组合方法有效地解决了这些问题，实现了更优的分割效果。
消融研究：表2中呈现的消融研究展示了所提出的三阶段学习（WSL）框架与不同分割 Backbone 网络组合的贡献。从表2中可以看出，包含SwinUNet的WSL框架表现较差，这表明尽管独立的SwinUNet算法性能能够超越UNet，但多个Multi-SwinUNet模型之间缺乏差异化。
值得注意的是，Mamba-UNet可以通过学习更远距离的特征依赖性来增强多个Mamba-UNet模型之间的特征多样性，以展现出色的性能。最后，作者提出的WSL框架在大多数分割指标上实现了最优结果，这证明多种不同类型的独立算法可以通过不同 Level 的特征信息相互补充，以提高网络的分割性能。
4 Conclusion
Weak-Mamba-UNet通过在涂鸦监督学习框架内整合CNN、ViT和VMamba的特征学习能力，显著降低了标注工作所需成本和资源。采用的多视角交叉监督学习方法增强了不同网络架构的适应性，使它们能够相互受益。关键的是，本研究证明了在有限信号监督下，新颖的Visual Mamba网络架构在医学图像分割中的有效性。
这项研究的积极成果不仅突显了网络在分割任务中的高准确性，还强调了其在医学图像分析中更广泛应用的潜力，尤其是在资源受限的环境下。
参考
[1].Weak-Mamba-UNet:Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation.
点击上方卡片，关注 「AI视界引擎」 公众号