ResNet50 还能再战 !结合 SAM 超越 Baseline,解决分割任务难题 !

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

最近引入的Segment-Anything Model(SAM)有潜力大大加速分割模型的发展。然而,直接将SAM应用于手术图像存在关键限制,包括:

(1)在测试时需要特定于图像的提示,从而阻止了完全自动化的分割;

(2)由于自然图像与手术图像之间的领域差距较大,导致效果不佳。

在这项工作中,作者提出了CycleSAM,一种用于单次手术场景分割的方法,该方法在测试时使用训练图像- Mask 对来自动识别测试图像中对应于每个目标类的点,然后可以用来提示SAM生成目标 Mask 。

为了产生高保真的匹配,作者引入了一种新颖的空间循环一致性约束,该约束强制测试图像中的点 Proposal 重新匹配到训练图像中目标前景区域内的点。

然后,为了解决领域差距,作者没有直接使用SAM中的视觉特征,而是采用了一种在自监督方式下在手术图像上预训练的ResNet50编码器,从而保持了高标签效率。

作者在两个不同的手术语义分割数据集上评估了CycleSAM的单次分割性能,全面超越了 Baseline 方法,并达到了完全监督性能的50%。

1 Introduction

SAM [(8)]是一种强大的通用分割模型,在自然图像领域取得了优异的结果。几项工作成功地将SAM应用于医学图像分割,以及手术工具场景分割。然而,SAM的标准应用需要用户在测试时提供点、框或 Mask 的提示,这阻碍了完全自动化的分割。

一些SAM适配方法已经被提出,以在测试时绕过用户提示。和通过自定义提示编码器增强 SAM,前者在测试时使用 GT 类别名称作为提示,而后者通过从训练图像中提取每个目标的前景特征构建类别原型提示。然而,这些方法通常以完全监督的方式应用,因此仍然需要大量的分割标注,这在手术领域极其昂贵;此外,它们仅在传统的语义分割模型(如DeepLabv3+和UNet)上实现了适度的性能提升。

相比之下,作者主要感兴趣的是利用SAM进行标签高效的分割。作者的关键洞察得到了[(17; 13)]的支持,即尽管自然图像和手术图像之间存在领域偏移,如果提供良好的提示,SAM仍然可以生成准确的 Mask 。这种提示和 Mask 生成的解耦可以使SAM在标签高效的环境中比传统分割方法取得更大的增益,因为即使在提示生成的部分准确性也可以转化为准确的 Mask 。为此,在本研究中,作者关注一次性手术场景分割的问题,与只考虑后者的[(22)]不同,作者同时解决复杂的解剖结构以及手术工具。

作者构建的方法CycleSAM是Personality [(23)]的扩展,它使用特征匹配在测试图像中识别与训练图像中每个目标类别匹配的点。为了获得更高保真的匹配,作者提出了CycleSelect模块,用于改进特征点匹配,包括:(1)一个空间循环一致性 Mask 函数,抑制那些在训练图像的目标前景区域不重新匹配的点;(2)多尺度特征匹配;(3)对点提示采样过程的一系列改进,包括多次前景点选择和相关负点选择。然后,认识到基于特征点的匹配高度依赖于具有区分性的底层特征空间,作者建议将SAM中的ViT特征提取器与MoCov2 [(2)]和DINO [(1)]自监督特征提取器进行增强,这些提取器在[(15)]中的领域内手术图像上进行预训练。最后,为了调整点特征,作者通过强制前景和背景点特征之间的嵌入级对比损失来训练线性特征投影器。

作者在两个具有挑战性的手术场景分割数据集上评估了作者的方法,分别是Endoscopes-Seg50 和CaDIS ,选择它们是因为它们包含了解剖结构以及工具。在实验中,作者展示了CycleSAM在一次性分割设置中优于传统专业方法和基于SAM的分割方法。

此外,仅用一个标记的训练图像,CycleSAM 就能达到接近完全监督性能的 50%,从而大大缩小了性能差距。

总之,作者的贡献如下:

  1. 作者提出了CycleSAM,一种基于特征匹配的1次适应SAM到复杂手术场景分割的方法。
  2. 作者引入了一个新颖的特征匹配模块 CycleSelect,它使用空间循环一致性在图像间大大提高了点对应。
  3. 作者在标签高效的手术场景分割方面全面超越了现有的SAM适配方法。

2 Methods

在本节中,作者提供了对CycleSAM的详细描述,它接收一个参考(训练)图像 ,相应的参考 Mask ,测试图像 ,图像编码器 ,以及SAM模型 ,并为每个目标类别 输出一个二值 Mask 。作者首先描述CycleSelect,这是一个模块,它基于密集特征匹配,为每个目标类别 在 中生成一组点提示 ;然后,每个 可以被传递到SAM Mask 解码器 次,以生成每个类别的 Mask 。接着,作者描述了CycleSAM的架构和训练过程,包括(1)特征相似性预热阶段,以细化匹配特征空间,以及(2)对PerSAM [23]的微调阶段,该阶段学习类别特定的三个 Mask , 和 预测的重加权,这些 Mask 由 预测得出。

CycleSelect

CycleSelect算法,如算法1所示,基于一个关键假设:属于某一目标类别的点在不同图像中应该具有相似的特征表示。因此,给定一个图像- Mask 对作为参考,作者可以通过比较它们的基础特征表示,识别测试图像中与参考前景点相匹配的目标点。

picture.image

CycleSelect的第一步是使用图像编码器提取参考图像和测试图像的特征,得到特征图和;然后作者将这些特征图和参考 Mask 空间调整为大小。接下来,为了构建每个类别的目标特征,作者首先从前景区域均匀采样个点,并用这些点索引参考特征图,得到点特征;然后,作者计算所有前景点的特征平均值,得到一个前景总结特征,将其与连接起来,得到,其中,是特征通道维度。

计算完后,作者继续识别相似点特征。为此,作者计算目标特征与测试特征图之间的余弦相似度矩阵。最后,作者将与个目标特征的相似度汇总到一个全局相似度图中,其中最高值的索引对应于测试图像中最可能属于目标类别的点。

空间循环一致性 Mask 。虽然作者可以直接使用来采样点提示,但这常常导致错误的匹配,因为产生的特征图在点粒度上不够具有辨识性。为了获得更稳健的相似度矩阵,受到几项先前工作的启发[6; 7; 9; 16],作者引入了一个空间循环一致性约束,该约束强制匹配点重新匹配到属于参考目标内的点。直观地说,这种约束可以帮助过滤掉由噪声特征图和局部歧义导致的较差匹配,对于处理不含目标目标的测试图像尤为重要。

为了将这种循环一致性约束整合到作者的方法中,作者没有简单地从中移除点,而是通过三个步骤计算循环一致性 Mask :(1) 作者在所有测试点和参考点特征之间密集计算余弦相似性,(2) 作者找到与每个测试特征最相似的参考特征的索引,(3) 如果重匹配特征的索引不在参考 Mask 中,作者将中相应的值设置为1。然后作者可以将和原始相似性矩阵组合成最终的循环一致性相似性矩阵;为此,作者只需从被 Mask 的的元素中减去一个常数。

多尺度特征匹配 为了处理不同大小的目标,作者 Proposal 通过在不同尺度重复CycleSelect,产生一组相似性矩阵。然后作者可以将这些相似性聚合成一个单一的相似性矩阵,在此过程中,作者将每个矩阵调整到大小,就像作者对多个目标特征所做的那样。算法2详细介绍了多尺度匹配。

picture.image

点提示采样 PerSAM [23] 每张图像只使用一个点提示;而作者是从最终的相似性矩阵中选择前景和背景的前个点,从而得到更丰富的提示信息。然后,为了提高背景点的实用性,作者将相似度分数聚类成3个箱,从分数最低的箱中选择前个点。直观地说,这个过程将点分成了强烈正匹配、模糊点和强烈负匹配。通过从强烈负匹配中选择相似度最高的点,作者可以获得信息量大的排除点提示,同时确信所选择的点确实是负点。最后,作者用每个类别的可学习阈值对正匹配进行阈值处理;作者在训练期间计算这些阈值,作为相似性矩阵的加权平均值,如下所示:

按照这个过程,作者获得了个点提示,作者将这些提示传递给SAM Mask 解码器以得到 Mask ,和。

CycleSAM Architecture and Training

图1展示了CycleSAM的总体架构,其主要的组成部分是可训练的CycleSelect模块。CycleSelect包含四个不同的可训练组件:(1) 特征投影块 ,包括每个目标类别和特征尺度的两个线性特征投影器(总共个),(2) 特征金字塔网络(FPN)[10] 用于融合多尺度图像特征,(3) 每个类别的可训练尺度权重 用来重新加权CycleSelect中多尺度匹配产生的相似性矩阵 ,以及 (4) 类似于[23]中每个类别的可训练 Mask 权重 用来重新加权 预测的三种 Mask 。为了以一次学习的方式训练CycleSAM,作者简单地将参考图像 的一个副本作为测试图像,将这个副本称为 。作者分两个阶段进行训练:第一个阶段是特征相似性预热阶段,在这个阶段作者训练 , 和 ;第二个阶段是PerSAM [23]的 Mask 细化训练过程,在这个阶段作者只训练 。

picture.image

特征相似性预热 PerSAM [23] 直接使用来自 的特征图 和 ,在他们的案例中,这是来自 的ViT图像编码器,以构建一个目标特征 并计算相似性矩阵 。然而,在处理外科图像时,由于大量的领域偏移,提取的特征空间远不够具有辨识性。因此,作者 Proposal 改为使用冻结的SSL预训练的ResNet50编码器 [15] 作为 。接着,为了提高CycleSelect中多尺度匹配的质量,作者进一步使用可训练的FPN 处理SSL特征;因此,更新后的图像编码器变为 。最后,由于这样的SSL预训练编码器可能缺乏专门针对像素级分割训练的SAM ViT编码器的细粒度精度,作者在计算 之前,对测试特征图 和目标特征 应用非对称线性特征投影器:,其中 是 的输出通道数, 是最终投影的特征维度。作者对每个类别 和特征尺度 包含这些投影器,总共产生 个投影器 。

为了训练这些层,作者引入了一个特征相似性对比损失 (见算法3),该损失最大化目标特征与前景测试特征之间的相似性,同时最小化目标特征与背景特征之间的相似性。注意,由于作者在训练期间的测试图像仅仅是参考图像,作者还有一个测试 Mask ,作者用它来分离前景和背景特征。最后,为了训练多尺度权重 ,作者简单地将预测的相似性矩阵 当作一个分割 Mask ,并计算与真实 Mask 之间的L1损失。

picture.image

算法3 特征相似性对比损失

** Mask 加权** 在预训练特征相似性组件(, , )之后,作者冻结它们并继续训练 。作者遵循与[23]相同的过程,使用CycleSAM预测 的 Mask ,并计算与真实 Mask 之间的损失(Dice损失和二元交叉熵损失)。在这里,作者继续对 进行非对称增强以改善泛化能力。

3 Experiments and Results

作者在3.1节首先描述作者的评估数据集和评价指标,然后呈现主要结果。接着,在3.2节中,作者提供一个消融研究,探讨CycleSAM改进的提示选择机制的影响。

Main Results

作者在两个不同的手术场景分割数据集Endoscapes-Seg50和CaDIS上评估了所有方法,并报告了不包括背景的mIoU和平均Dice得分(mDice)。Endoscapes-Seg50[12]是一个用于腹腔镜胆囊切除术中安全解剖的分割数据集,包含6个类别(5个解剖结构和一个工具类别);由于其包含了细粒度的解剖结构/窗口(例如,胆囊板,胆囊三角),因此具有独特的挑战性。同时,CaDIS[4]是一个用于白内障手术场景分割的数据集,与Endoscopes-Seg50一样,也包含了解剖结构和工具。CaDIS包含了三个不同粒度的分割任务:在这项工作中,由于作者已经采用了具有挑战性的1次射击设置,作者关注第一个任务,这是一个8类分割问题(4个解剖结构,1个工具类别,3个其他)。对于训练,作者随机选择了一张包含所有目标类别的图像,并将其用于所有方法。

表1展示了所有方法在两个数据集上的结果,分为四类:

(1)专业模型,作者考虑DeepLabV3Plus-ResNet50(DLV3P-R50)作为基础分割模型,并实验了不同的主干初始化(默认、SurgMoCov2和SurgDINO);

(2)基于SAM的方法,包括适应方法[22, 11]和自动提示生成方法[23, 20];

(3)CycleSAM,也用不同的初始化进行了测试;

(4)一个完全监督的DeepLabV3Plus模型,代表天花板的性能。对于SurgMoCov2和SurgDINO,作者使用了[15]提供的权重,对于Endoscopes-Seg50实验使用Cholec80预训练模型,对于CaDIS实验使用CaDIS预训练模型。作者还测试了使用原始SAM ViT主干以及来自组(1)的1次射击训练的DLV3P-R50模型来初始化的CycleSAM变体。

picture.image

作者最好的方法是使用训练效果最佳的组1模型DLV3P-SurgMoCov2初始化的CycleSAM。CycleSAM-DLV3P-SurgMoCov2不仅超过了所有 Baseline 方法,而且在两个数据集上均实现了约50%的完全监督mIoU和超过50%的mDice。有趣的是,作者发现最好的 Baseline 方法是组1的专业DeepLabV3Plus模型,最佳模型接近CycleSAM(Endoscopes-Seg50差距为5.40 mIoU,CaDIS差距为4.07 mIoU)。

另一方面,作者全面超越了现有的基于SAM的 Baseline 方法,这些方法在依赖现成的特征表示方面对于一次性手术场景分割非常无效。通过研究CycleSAM-ViT,它直接使用SAM ViT特征提取器作为,性能与其他基于SAM的模型相似,这一点得到了强调。此外,使用MedSAM[11]作为底层模型而不是SAM的PerSAM-F-MedSAM,其表现比原始PerSAM还要差,突显了当前将SAM适应到医疗领域的局限性。通过灵活地利用任何特定领域的特征提取器,CycleSAM在基于SAM的现有方法上取得了巨大的性能提升(Endoscopes-Seg50提升+15.44 mIoU,CaDIS提升+18.03 mIoU)。

Prompt Selection Ablation Study

表2展示了作者改进的点采样方法对PerSAM [23]默认方法的影响;通过包含5个不同的点作为提示并包含相关的负点,作者在Endoscopes-Seg50上的mIoU提高了8.54,几乎使性能翻了一番。

picture.image

4 Conclusion

作者介绍了CycleSAM,这是一种使用循环一致特征匹配的方法,用于将SAM一次性适应于手术场景分割。作者的方法不仅局限于已被深入研究的工具分割问题,还解决了两个具有挑战性的手术解剖分割任务。

通过利用手术领域SSL预训练的图像编码器,CycleSAM有效地解决了域偏移问题,在一次性设置中大幅超越了现有方法的表现。

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
评论
未登录
暂无评论