赞
踩
SurgicalSAM: Efficient Class Promptable Surgical Instrument Segmentation
SurgicalSAM:高效的可分类手术器械细分
2023年 发布在Arxiv
Paper Code
Segment Anything Model(SAM)是一个强大的基础模型,它彻底改变了图像分割。为了将SAM应用于手术器械分割,常见的方法是定位器械的精确点或盒,然后将它们作为SAM的提示以zeroshot方式使用。然而,我们观察到这种幼稚流水线的两个问题:(1)自然对象和手术器械之间的域间隙导致SAM的泛化较差;以及(2)SAM依赖于精确的点或框位置进行准确分割,需要大量的手动指导或性能良好的专业检测器进行及时准备,这导致复杂的多阶段流水线。为了解决这些问题,我们引入了SurgicalSAM,这是一种新型的SAM端到端高效调整方法,可以有效地将特定于医疗器械的信息与SAM的预训练知识相结合,以提高泛化能力。具体来说,我们提出了一个轻量级的基于原型的类提示编码器的调整,直接从类原型生成提示嵌入,并消除了显式提示的使用,以提高鲁棒性和更简单的管道。此外,为了解决手术器械类别之间的低类间方差,我们提出了对比原型学习,进一步提高了类原型的区分度,以获得更准确的类提示。在EndoVis 2018和EndoVis 2017数据集上进行的大量实验结果表明,SurgicalSAM实现了最先进的性能,同时只需要少量的可调参数。源代码可在https://github.com/wenxi-yue/SurgicalSAM上获得。
Keywords:隐式提示,类原型,基础模型,轻量级
手术器械分割是手术视觉中的一项重要任务,其目的是在手术场景中精确地分割出手术器械。现有的SIS(Surgical instrument segmentation)深度学习方法通过设计和训练具有特定任务组件的专家模型取得了令人印象深刻的结果。然而,这些方法通常需要训练完整的模型参数集(即,完全训练)使用SIS数据集,导致效率低下。此外,由于SIS数据集的规模有限,经过训练的模型往往表现出低于标准的泛化性能。Segment Anything Model(SAM)作为可解释细分的开创性基础模型,最近获得了极大的关注。将SAM用于下游医疗任务对于提高培训效率和利用强大的预先培训知识具有很大的希望。目前的研究主要采用SAM在零拍摄的方式用于医学图像分割。然而,SAM预训练中缺乏足够的医疗数据以及自然对象和医疗目标之间的巨大领域差距阻碍了SAM对医疗任务的直接推广。许多研究报告了SAM在零样本医学图像分割中的低于标准的性能。
具体而言,手术器械在专业外观、复杂的解剖背景和高类别间相似性方面与自然物体有显著差异。我们在SIS上评估了三种基本的零样本SAM策略:(1)MaskTrackRCNN或Mask2Former作为边界框检测器与SAM一起,(2)Track Anything和(3)PerSAM,分别代表基于检测,基于跟踪和基于参考的框架。这些方法证明了较差的结果,其中基于检测和基于跟踪的方法描绘了不正确的轮廓,并且基于参考的方法错误地识别了仪器类别。这进一步突出了弥合自然外科领域差距的挑战,并强调了SAM调整的必要性。此外,SAM的性能依赖于显式提示的精确位置。我们证实了这一点,通过一个提示的鲁棒性研究SIS通过引入各种规模和位置抖动的地面实况边界框作为提示SAM和记录的预测mAP。如图2所示,我们的研究证明了SAM对提示抖动的敏感性:即使在所提供的边界框提示中存在微小偏差,也会显着损害分割精度。因此,现有的零样本SAM框架通常涉及复杂的多级流水线,需要精确的手动指导或性能良好的专业检测器来提供准确的点或边界框以进行准确的提示。这种复杂性进一步限制了SAM在外科领域的直接应用。
为了解决上述挑战,本文提出了SurgicalSAM,这是一种端到端的方法,可以通过有效调整SAM来有效缓解外科-自然领域的差距。SurgicalSAM与现有管线的比较如下图所示。我们提出了一个轻量级的基于原型的类提示编码器,它采用一个仪器类作为提示,并通过与图像嵌入交互来学习类原型,以直接为掩码解码器生成提示嵌入。通过调整基于原型的类提示编码器和掩码解码器,手术知识与SAM的预训练知识相结合,有效地缓解了领域差距。此外,我们直接从类提示中生成潜在提示嵌入并消除显式点和边界框的使用的策略进一步解决了与显式提示相关的鲁棒性差的问题,并保持了端到端的流水线。在SurgicalSAM中,类原型在有效地从图像中提示感兴趣的仪器方面起着至关重要的作用。然而,不同的手术器械类别往往表现出高相似性和低的类间差异,因此提出了一个很大的挑战。为了解决这个问题,我们进一步提出了对比原型学习,利用对比损失来获得有区别的学习类原型。该方法增强了细粒度仪器类别之间的区分,从而导致更准确的类别提示和改进的分割结果。
在这项工作中,我们解决了手术器械分割的任务,通过有效的调整SAM类可扩展的方式。具体地,给定具有空间分辨率H ×W的手术图像 I ∈ R H × W × 3 I ∈ R^{H×W×3} I∈RH×W×3,图像c中仪器的类别作为提示,我们的目标是预测图像的类别c掩码,表示为 M ( c ) M(c) M(c): M ( c ) = S u r g i c a l S A M ( I , c ) M^{(c)} = SurgicalSAM(I, c) M(c)=SurgicalSAM(I,c)(即给定图像I和提示c)
SurgicalSAM由三个核心组件组成,如上图所示:图像编码器、基于原型的类提示编码器和掩码解码器。与SAM类似,图像编码器EI首先提取输入图像的嵌入为
F
I
∈
R
h
×
w
×
d
F_I ∈ R^{h×w×d}
FI∈Rh×w×d,其中h × w表示图像嵌入的形状,d表示嵌入通道的数量。然后,我们基于原型的类提示编码器
E
C
P
E_{CP}
ECP利用类原型B来激活图像嵌入,并利用所获得的以提示类c为条件的激活特征来生成提示嵌入,包括密集提示嵌入
T
D
(
c
)
T^{(c)}_D
TD(c)和稀疏提示嵌入
T
S
(
c
)
T^{(c)}_S
TS(c)。最后,通过掩码解码器
D
M
D_M
DM使用图像嵌入和提示嵌入来预测mask
M
(
c
)
M^{(c)}
M(c)。上述过程可以表示为:其中
T
O
T_O
TO表示可学习的输出tokens。
基于原型的类提示编码器:
基于原型的类提示编码器利用图像和类原型之间的相似性来创建提示嵌入。具体地,如上图所示,计算图像嵌入和类原型之间的空间相似性以激活图像内的类特定区域,从而产生类激活特征以生成用于掩码解码器的提示嵌入。此外,受SAM中前景和背景点提示的利用的启发,我们建议不仅采用提示类的原型,而且集成所有类原型,以结合积极和消极的线索。这种策略为模型提供了更鲁棒的先验,以有效区分具有高相似性的仪器类别。
具体地,基于原型的类提示编码器
E
C
P
E_{CP}
ECP建立在原型库
B
=
c
o
n
c
a
t
(
B
(
k
)
,
k
∈
1
,
2
,
.
,
C
)
∈
R
C
×
d
B = concat({B^{(k)}},k∈{1,2,.,C})∈ R^{C×d}
B=concat(B(k),k∈1,2,.,C)∈RC×d,由每个类的代表原型组成,其中C是类的总数。给定一个图像I,其图像嵌入为FI,我们构造一个相似矩阵
S
=
c
o
n
c
a
t
(
S
(
k
)
,
k
∈
1
,
2
,
.
,
C
)
∈
R
C
×
h
×
w
S = concat({S^{(k)}},k∈{1,2,.,C})∈ R^{C×h×w}
S=concat(S(k),k∈1,2,.,C)∈RC×h×w来表示图像与所有类的原型的空间相似性。它是通过计算每个空间位置的图像嵌入和每个类原型之间的点积生成的:
然后采用相似性矩阵作为空间注意力来激活类特异性区域,从而产生所有类的类激活特征
F
I
C
=
c
o
n
c
a
t
(
F
I
(
k
)
k
∈
1
,
2
,
.
.
.
,
C
)
∈
R
C
×
h
×
w
×
d
F^C_I =concat({F^{(k)}_I }k∈{1,2,...,C}) ∈ R^{C×h×w×d}
FIC=concat(FI(k)k∈1,2,...,C)∈RC×h×w×d:
F
I
(
k
)
=
F
I
◦
S
(
k
)
+
F
I
,
f
o
r
k
∈
1
,
2
,
.
.
.
,
C
F^{(k)}_I = F_I ◦ S^{(k)} +F_I, for k ∈ {1, 2, ..., C}
FI(k)=FI◦S(k)+FI,fork∈1,2,...,C
其中 ◦ 和+分别表示逐元素乘法和加法,
F
I
(
k
)
∈
R
h
×
w
×
d
F^{(k)}_I ∈ R^{h×w×d}
FI(k)∈Rh×w×d表示类k的类激活特征。
类激活的功能是用来制做密集和稀疏提示嵌入。在SAM中,密集提示嵌入来自前景掩膜,为分割对象提供积极的线索。模仿这一点,我们利用积极类的类激活特征,即,提示类c,用于编码稠密提示嵌入
T
D
(
c
)
∈
R
h
×
w
×
d
T^{(c)}_D ∈ R^{h×w×d}
TD(c)∈Rh×w×d。这是通过两层多层感知器(MLP)实现的:
其中fD和gD是具有中间维度rD的两个线性投影函数。
另一方面,SAM中的稀疏提示嵌入从正信息(前景点和边界框)和负信息(背景点)编码。受此启发,我们使用所有类的类激活功能生成稀疏提示嵌入,这些类包括积极的提示类和消极的非提示类。然后通过一对正嵌入和负嵌入来区分正类和负类。具体来说,
F
I
C
F^C_I
FIC首先被馈送到两层MLP中以获得正不可知的稀疏提示嵌入,
T
ˆ
S
C
=
c
o
n
c
a
t
(
ˆ
T
S
(
k
)
,
k
∈
1
,
2
,
.
.
.
,
C
)
∈
R
C
×
n
×
d
Tˆ^C_S=concat({ ˆT^{(k)}_S },k∈{1,2,...,C}) ∈ R^{C×n×d}
TˆSC=concat(ˆTS(k),k∈1,2,...,C)∈RC×n×d
其中fS和gS是两个具有中间维度rS的线性投影函数,n表示每个类的稀疏令牌的数量,并且
T
ˆ
S
(
k
)
∈
R
n
×
d
Tˆ^{(k)}_S ∈ R^{n×d}
TˆS(k)∈Rn×d表示由类k激活的正不可知稀疏提示嵌入。然后,分别将一对正嵌入和一对负嵌入λ+ ∈ Rd和λ− ∈ Rd添加到嵌入中对应于正类(类c)和负类(类c以外),导致最终的稀疏提示嵌入
T
S
(
c
)
∈
R
C
×
n
×
d
T^{(c)}_S ∈ R^{C×n×d}
TS(c)∈RC×n×d是正感知的。
然后 T S ( c ) T^{(c)}_S TS(c)被整形为Cn×d,并与 T D ( c ) T^{(c)}_D TD(c)一起被馈送到掩码解码器中用于掩码预测。
原型对比学习:
本文方法依赖于判别类原型精确的仪器类别识别和准确的类区域激活。然而,在具有高度相似的仪器外观的手术场景中获得准确的类原型是具有挑战性的。为了增强原型的区分度,以获得更准确的类提示,我们提出了对比原型学习,以在调整框架期间获得优化的类原型,如上图所示。具体来说,我们提出了由infoNCE损失引起的原型对比损失,其中类原型被视为锚点,训练图像中基于SAM的类嵌入被视为样本。给定图像嵌入FI,将类别c的真实二进制掩码G©处理为分辨率h ×w,并用于通过对前景特征进行平均来提取类别c的基于SAM的类别嵌入
v
(
c
)
∈
R
d
v(c)∈ Rd
v(c)∈Rd:
为此,原型对比损失表示为:
其中τ是指用于调制相似性的温度参数,并且B(k)是类别k的原型。可以说
L
P
C
L
L_{PCL}
LPCL增强了类k(锚)的原型与类k(正样本)的基于SAM的类嵌入之间的相似性,同时抑制了类k(锚)的原型与类k以外的类(负样本)的基于SAM的类嵌入之间的相似性。这导致更多的区别性原型表示和通过SAM调整增强的外科领域知识注入。
损失函数:
SurgicalSAM训练效率高。在调整期间,大图像编码器被冻结,并且仅更新基于轻量级原型的提示编码器和掩码解码器的参数。调优是端到端的,由两个项组成的损失函数监督:分割的骰子损失和原型学习的原型对比损失:
其中mi和gi分别是图像的像素i处的预测logit值和gt二进制值。
对EndoVis 2017和EndoVis 2018的数据进行预处理。对于基于原型的提示编码器,中间尺寸rD和rS均设置为128,并且对于EndoVis 2018和EndoVis 2017,每个类的标记数n分别设置为2和4。对于原型对比损耗,使用0.07的温度τ。在训练方面,我们使用SAM的ViT-H版本的预训练权重初始化图像编码器、掩码解码器以及SurgicalSAM的正嵌入和负嵌入(λ+和λ−)。我们模型的图像编码器和正负嵌入保持冻结,而提示编码器和掩码解码器的权重更新。我们采用了Adam优化器,EndoVis 2018和EndoVis 2017的学习率分别为0.001和0.0001。为了减少计算量,我们在训练中采用预先计算的图像嵌入,采用批量大小为32。我们的模型使用PyTorch实现,并在Nvidia Tesla V100 16 GB GPU上进行训练和评估。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。