点击下方卡片,关注
「AI视界引擎」
公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
素描标注在大规模医学数据集中显著降低了复杂解剖结构所需的密集标注的成本和劳动需求。然而,当前的素描监督学习方法在有效传播稀疏标注标签到密集分割 Mask 以及准确分割物体边界方面能力有限。
为了解决这些问题,作者提出了一种渐进协作学习框架,该框架结合了新的算法和MedSAM基础模型,在训练过程中提升信息质量。
(1)通过一种新算法丰富 GT 素描分割标签,并传播素描以估计物体边界。
(2)通过融合Med-SAM和提出的Sparse Mamba网络的特征嵌入来优化由Med-SAM指导的训练,从而增强特征表示。这种增强的表示也促进了Med-SAM解码器在丰富素描上的微调。
(3)在推理阶段,作者引入了一个Sparse Mamba网络,它通过跳过采样程序取代传统的顺序像素处理方法,能够高效地捕捉局部和全局依赖关系。作者在ACDC、CHAOS和MSCMRSeg数据集上的实验验证了该框架的有效性,并且在与九种最先进的方法进行比较时表现更优。
作者的代码可在SparseMamba-PCL.git中获得。
1 Introduction
医学图像分割在计算机辅助诊断中具有巨大的潜力,因为它能够识别感兴趣的区域,如器官、病灶或肿瘤,从而帮助临床医生制定更有效的治疗计划和随访策略[19,12]。然而,高度准确的分割需要大量的标注数据集,生成这些数据集既耗时又成本高昂[30,13]。为了解决这一挑战,研究行人已经探索了基于素描监督的方法[5,27,21],这些方法依赖于稀疏标注而非逐像素标签进行模型训练。然而,结构信息对于精确分割至关重要,但在以往的工作中,很少有研究致力于将其整合到学习框架中。
从潦草 Token 中学习是一个精细的过程,需要高性能模型和复杂的学习策略。近期在医学图像分割领域的高 performances 模型取得了进展,主要依赖于卷积神经网络(CNNs)、Transformer [18,?,4] 和 Mambas [24,15,2,22]。Mambas 因其能够捕捉全局依赖性而备受关注,解决了 CNN 的局限性,且与 Transformer 的二次复杂度不同,Mambas 维持了线性的计算复杂度。Mambas 使用二维选择性扫描 (SS2D) 来从多个方向处理图像,并利用选择性状态空间模型 (S6) 来处理每个方向上的序列,以捕捉全局依赖性。然而,多方向扫描可能导致特征冗余并掩盖重要片段,因为每个片段可能被多次扫描,这使得难以确定每个片段的重要性 [20]。
使用基础模型在弱监督数据上训练一个分割模型可以显著增加训练框架的先验知识,正如[28]所演示的那样。在此基础上,作者利用Med-SAM对模型进行训练,Med-SAM是Segment Anything Model(SAM)的一种变体,该变体已经在医学领域数据上进行了微调[3]。
然而,有效地利用Med-SAM进行弱监督训练仍然具有挑战性:
(1)MedSAM的表现很大程度上依赖于输入 Prompt ,而在弱监督设置中,这些 Prompt 往往不理想[11];(2)在训练过程中静态使用Med-SAM而不对其进行微调限制了编码器的特征提取能力,同时也限制了解码器在专家任务上的分割性能;
(3)Med-SAM在目标边界处的分割精度较差。
为了解决这些问题,作者提出了一种新颖的scribble监督医学图像分割框架,通过逐步协作学习实现。本框架的贡献如下:
渐进式协作学习(PCL)增强了协作学习以提高分割性能。它通过细化粗略的分割 Mask 为精确的边界框 Prompt ,结合两个不同编码器的嵌入以获得更丰富的特征表示,并在训练过程中迭代微调Med-SAM解码器以集成专家知识。
SparseMamba 采用修改后的跳样算法替代了传统 Mambas 中按顺序处理图像 Patch 的流程,以增强模型捕获全局依赖关系的能力。
Scribble-传播的目标边界估计器(SPOBE)利用图像和草图中的边缘线索生成辅助监督信号,从而提高沿目标边界的分割准确性。
作者在ACDC、CHAOS和MSCMRSeg数据集上进行了大量实验。作者的框架优于9种最新的scribble监督分割方法。
2 The Proposed Framework
SparseMamba-PCL 是一种弱监督医学图像分割框架,它将目标边界估计器、Med-SAM 引导的训练算法和稀疏 Mamba 网络(图1l)协同结合在一起。首先,作者的框架提取目标边界像素作为辅助监督信号,并将其与 Token 相结合生成增强的 Token (图1(a))。在训练过程中,作者同时使用稀疏 Mamba 网络和 Med-SAM 处理输入图像(图1(b) 和图1(c))。稀疏 Mamba 网络生成粗略的分割 Mask ,从中提取边界框以 Prompt Med-SAM。同时,Med-SAM 将其编码器的图像嵌入与来自稀疏 Mamba 网络的嵌入合并,然后与边界框 Prompt 合并。最后,Med-SAM 的 Mask 解码器生成精炼的分割 Mask 。在训练过程中,这些精炼的 Mask 用于优化稀疏 Mamba 网络的权重。此外,通过使用增强的 Token 同时训练稀疏 Mamba 网络和 Med-SAM 来进一步增强监督。
2.1 Scribble-Propagated Object Boundary Estimator
可靠地识别目标目标的边界对于准确分割至关重要。然而,潦草标注极少能提供高质量的边界信息,这使得弱监督分割变得非常具有挑战性。为了缓解边界信息缺乏的问题,作者提出了一种潦草标注传播的目标边界估计器(SPOBE),该方法通过迭代过程利用边缘检测器和潦草标注来识别目标目标的边界。
首先,对输入图像应用边缘检测器并获得完整的边缘图
. 然而,该边缘图包含许多与给定目标物体边界无关的边缘(见图2(c))。为了区分边界的边缘和噪声边缘,作者设计了一种迭代方案。图2展示了单次迭代的过程。在第一次迭代中,作者使用大小为
的方形 Kernel 初始化一个与输入图像相同大小的计数图,并对其进行膨胀,从而得到擦除 Token (图2(a) 和 (b))。计数图和类
的擦除 Token 分别记为
和
. 在第
次迭代中,类似的表示方法被用于表示
和
。通过
来识别初始边界边缘,如图2(c) 和2(d)所示,其中
表示逻辑与操作。在第二次迭代中,作者使用更大的方形 Kernel (大小为
)对擦除 Token 进行膨胀,以获得
。如果以
为中心、大小为
的
贴图中的总值小于预先设定的阈值
,则计数图
;否则
. 这个阈值限制了每次迭代中可以添加到目标边界图中的边缘像素数量,以减少 Kernel 尺寸增大时误判噪声边缘像素的风险。第二次迭代中识别出的边界边缘为
. 该过程重复
次,并且 Kernel 尺寸逐渐增大。图2(e) 展示了最终目标边界的示例。
2.2 Sparse Mamba Network
图1(b)展示了SparseMamba,这是一种编码器解码器网络,在编码器中包含Sparse Mamba块以捕获局部和全局依赖关系。解码器使用残差块和转置卷积来保留细节和分辨率。U-Net风格的 Short-Cut 连接在编码器和解码器之间融合层级特征。SparseMamba还集成了一个双注意力模块 6 ,用于建模空间和通道依赖关系。解码器输出经过卷积层和Softmax激活函数后,预测分割概率图。
稀疏Mamba块(SMB)。如图1(d)所示,输入特征尺寸为
,先通过两个残差块,然后扩展并转置为
,其中
。特征在两条并行分支中进行处理:第一条分支使用线性层和SiLU激活将特征扩展至
;第二条分支应用线性层、1D卷积、SiLU以及Sparse SS2D。输出通过哈达玛乘积合并,并投影回
,再 Reshape 为
。如图3(a)所示,SS2D从图像的两个起始位置(左上角和右下角)以及两种空间方向(垂直和水平)扫描每个图像,从而产生四个独立的扫描操作。每个方向序列由一个S6块处理以捕捉全局依赖关系。然而,多向扫描引入了冗余性,因为每个小块被多次扫描,这可能掩盖了重要的小块,并导致确定它们重要性的困难。为此,Sparse SS2D(图3(b))采用跳过采样[20]策略,确保每个小块仅被扫描一次。这种策略减少了冗余性并增强了空间关系的表现力。
2.3 Progressive Collaborative Learning
渐进协作学习(PCL)是作者提出的一种训练框架,利用Med-SAM的解剖学先验引导SparseMamba的预测并精调网络权重。首先,如算法1所示,SparseMamba和Med-SAM的两个编码器独立地计算图像嵌入。这些嵌入随后相加以形成一个单一的融合嵌入,该融合嵌入结合了两个编码器的空间和结构信息,从而提供更丰富和多样化的特征表示。接下来,作者使用SparseMamba的输出,粗略分割 Mask
,来提取用于Med-SAM的边界框 Prompt 。Med-SAM的表现高度依赖于准确的 Prompt ,因为含糊的 Prompt 会导致由于医学图像的结构复杂性而产生次优分割[11]。在早期训练阶段,当
尚未充分训练时,作者将SparseMamba的边界框与从素描中提取的边界框结合起来,以增加 Prompt 的稳定性。然后,这些 Prompt 通过Med-SAM的 Prompt 编码器处理以生成 Prompt 嵌入。最后,融合的图像嵌入与 Prompt 嵌入合并为单个表示,Med-SAM解码器利用此表示生成具有更高准确性和边界精度的精细分割 Mask
。作者使用这些输出对SparseMamba和Med-SAM解码器进行微调。对于SparseMamba,作者使用Dice损失计算
和
之间的差异$L_{\mathrm{Dice}}(y_{1},y_{2})=1-\frac{2\left
}{|y_{1}|+|y_{2}|}
。
此
外
,
作
者
计
算
了
增
强
后
的
素
描
E_{S}
与
y_{1}
以
及
E_{S}
与
y_{2}
之
间
的
部
分
交
叉
熵
损
失
L_{pC E}(\boldsymbol{y},s)=-\sum_{c}\sum_{i\in\omega_{s}}\log(y_{i}^{c})
。
这
两
个
项
通
过
权
重
因
子
\lambda=0.5
平
衡
,
得
到
总
损
失
L_{\mathrm{total}}$(方程1)。
这用于优化SparseMamba。另一方面,为了确保Med-SAM在训练过程中能够获取任务特定的知识,作者使用
对Med-SAM的解码器进行微调。
3 Experiments
3.1 Experimental Settings
数据集与评估指标。作者方法在三个公开数据集上进行了评估:(1) ACDC [19]:包含来自100名患者的 cine-MRI 图像,人工标注了右心室(RV)、左心室(LV)和心肌(MYO)。该数据集分为70个训练集、15个验证集和15个测试集。(2) MSCMRseg [29]:包含45名心肌病患者晚期钆增强(LGE) MRI 扫描,并对RV、LV和MYO进行了标注,分为25个训练集、5个验证集和20个测试集。(3) CHAOS [10]:包含20个具有肝脏、肾脏和脾脏的腹部T1加权MRI图像,分为70%训练集、15%验证集和15%测试集。对于涂鸦标注,ACDC数据集使用了手动创建的涂鸦[19],而MSCMRSeg和CHAOS数据集则使用ITK-Snap进行1像素宽涂鸦的标注[5]。所有结果均基于5折交叉验证,评估指标为Dice系数和95%Hausdorff距离(HD95)。
实施细节。图像和标注被调整至相同分辨率,均为 (256 \times 256) 像素。在训练过程中,每个图像都被归一化到 ([0,1]) 的范围内,并进行了随机旋转、随机镜像和随机噪声处理。作者使用带有权重衰减 (10^{-4})、动量 (0.9) 和批量大小为 16 的 PolyLR 调度策略的 SGD 对模型进行优化,在总计 90,000 个迭代次数内进行。测试时,作者使用 SparseMamba 的输出进行预测。所有实验均在 PyTorch 中实现,并在 NVIDIA 2080Ti GPU 上训练,实验条件一致。
3.2 Experimental Analysis
表1将SparseMamba与基于Transformer和Mamba的方法在三个数据集中进行了比较。Transformer通常能够提高Dice分数,但HD95结果表现出不稳定现象,在三个数据集上分别达到了最高的HD95值(Trans UNet),以及在两个数据集上达到了最低的HD95值(SwinTrans)。相比之下,Mambas在Dice和HD95方面表现出一致的良好性能。SparseMamba在三个数据集上均实现了最高的Dice分数,在ACDC和MSCMRSeg上产生了第二低的HD95值,在CHAOS数据集上则产生了最低的HD95值。这些结果表明,SparseMamba在分割准确性和边界精度方面相较于现有方法有所提升。
表2展示了消融研究,并将SparseMamba-PCL与九种最新的scribble监督方法在三个数据集上的表现进行了对比。作者的方法在所有数据集上均实现了最高的Dice分数,并且在CHAOS数据集上具有最低的HD95值,在ACDC和MSCMRSeg数据集上则是第二低的HD95值,这证实了其在分割精度和边界细化方面的有效性。Baseline+SPOBE和Baseline+PCL相比Baseline(SparseMamba)有所改进,证明了边界 Aware 监督和SAM引导学习的益处。图4(l)比较了SparseMamba-PCL与其他scribble监督方法的分割性能,SparseMamba-PCL显示出更加平滑的边缘,能够精确地界定物体边界,而其他方法则表现出锯齿状或模糊的边缘。此外,示例还展示了SparseMamba-PCL在ACDC、CHAOS和MSCMRSeg数据集上实现的一致性分割质量,突显了其在多个医学领域中的适应能力。这种适应能力和精确的分割对于准确的容积分析和临床决策至关重要,因为即使是细微的边界不准确性也可能导致诊断错误。总之,SparseMambaPCL架构在多种医学图像数据集上为分割指标提供了稳定且 robust 的改进。
参考
[0]. SparseMamba-PCL: Scribble-Supervised Medical Image Segmentation via SAM-Guided Progressive Collaborative Learning .
点击上方卡片,关注
「AI视界引擎」
公众号