FocSAM 分割质量 SOTA 的同时显著降低计算成本！

用户5536580

发布于 2024-07-04 15:38:43

390

发布于 2024-07-04 15:38:43

分割任何模型（SAM）在分割模型领域标志着的一个重要里程碑，其突出特点是强大的零样本能力和处理多样化提示的能力。SAM遵循一个流水线，将交互式分割分为通过大型编码器进行图像预处理和通过轻量级解码器进行交互式推理，确保了高效的实时性能。然而，在这一流水线上，SAM在具有挑战性的样本上面临稳定性问题。这些问题源于两个主要因素。首先，图像预处理使得SAM无法在交互过程中动态地使用图像 Level 的放大策略来重新聚焦目标目标。其次，轻量级解码器难以充分地将交互信息与图像嵌入融合。为了解决这两个限制，作者提出了FocSAM，其流水线在设计上重新考虑了两个关键方面。首先，作者提出了动态窗口多头自注意力（Dwin-MSA），以动态地使SAM的图像嵌入重新聚焦于目标目标。 Dwin-MSA将注意力计算定位在目标目标周围，以最小的计算开销增强了与目标相关的嵌入。其次，作者提出了像素级动态ReLU（P-DyReLU），以充分整合初始点击中具有重大影响的交互信息，这些初始点击对整体分割结果有显著影响。实验上，FocSAM增强了SAM的交互式分割性能，使其达到现有最先进方法的分割质量，同时在CPU上的推理时间仅为此方法的约。代码可在https://github.com/YouHuang67/focsam获取。

1 Introduction

互动分割通过使用有限的手动标注提高扩大图像分割数据集的效率，避免了完全标记的繁琐工作。最近，Segment Anything Model (SAM) 在实时、高质量的互动分割方面表现出色，能够响应标注者的提示，如点击、边界框或粗略 Mask 。

SAM在处理各种提示的泛化能力和效率使其成为跨多种分割相关任务的通用工具。

本研究基于SAM 专注于基于点击的互动分割。

SAM 和同时期的InterFormer 开创了一种新的互动分割流程。该流程将强大的Vision Transformers (ViTs)作为图像编码器来预处理图像，生成适用于同一图像内所有目标的图像嵌入。

在互动过程中，这些图像嵌入和标注者提供的提示（例如，点击）被送入轻量级解码器以产生分割结果。这个流程将大型ViT的力量与现场互动分割所需的速度结合起来。

遵循这样的流程，SAM甚至使标注者能够在仅使用CPU的设备上进行实时、高质量的互动分割，有助于显著扩展图像分割标注。

然而，SAM的流程有两个局限性。首先，流程中的图像预处理使得无法高效实现图像级放大策略[46]，该策略在互动过程中动态地将模型聚焦于目标目标。

其次，由于需要实时响应，SAM的轻量级解码器难以充分融合互动信息与预处理图像嵌入，从而削弱了互动反馈对分割质量的正面影响。

因此，在具有挑战性的场景中，如几乎与背景融为一体的伪装目标，SAM面临稳定性问题。图1清晰地展示了SAM分割结果的不稳定性，紧随足够数量的先前点击（例如，9次点击）后的额外点击意外地触发了分割质量的显著下降，例如IoU从降到。这种不稳定性显著限制了SAM在更广泛的图像分割标注中的应用。

因此，作者提出了FocSAM来解决SAM的局限性。FocSAM的流程基于SAM并引入了一个额外的焦点优化器。这个优化器在目标的互动过程中调整SAM的图像嵌入，增加了可忽略的计算。调整带来了两个主要改进。首先，优化器使用初始分割结果，受到图像级放大的启发，将图像嵌入重新聚焦于包含目标目标的区域。其次，优化器充分融合了初始几次点击的嵌入，这些点击对最终分割结果影响很大[35]，进一步增强了与目标相关的嵌入。

为了以最小的计算开销实现FocSAM的焦点优化器，作者引入了动态窗口多头自注意力（Dwin-MSA）和像素级动态ReLU（P-DyReLU）。Dwin-MSA将图像嵌入划分为窗口，并在动态最小子集上执行有效的注意力计算，这些子集与先前预测的 Mask 相交的窗口嵌入，避免了在无关背景区域上的冗余计算。Dwin-MSA采用移动策略[39]来确保嵌入之间的长距离互动，保持动态效率。P-DyReLU作为Dwin-MSA的非线性激活，融合了初始几次点击的互动信息与图像嵌入。具体来说，P-DyReLU采用DyReLU[6]，并利用SAM解码器的点击融合 Query 嵌入来增强与目标相关的图像嵌入并抑制不相关的嵌入。

实验上，FocSAM在计算开销可忽略的情况下，展示了优于SAM的互动分割性能。在包括DAVIS[43]、SBD[19]、GrabCut[45]、Berkeley[26]、MVTec[2]和COD10K[10]的数据集上，FocSAM在点击次数（NoC）上与最先进的SimpleClick[36]相匹配，但FocSAM的CPU推理时间仅约为SimpleClick的。此外，当每幅图像的目标数量超过时，FocSAM的时间效率进一步提高，对于CPU推理，仅需要大约SimpleClick所需时间的。

作者的贡献总结如下：

作者引入了FocSAM，通过动态增强与目标相关的图像嵌入并深入整合互动信息到这些嵌入中，以提高SAM的性能。
FocSAM通过提出的Dwin-MSA和P-DyReLU实现，额外计算开销可忽略。
在包括DAVIS、SBD、GrabCut、Berkeley、MVTec和COD10K的数据集上，FocSAM在NoC上与最先进的SimpleClick相匹配，在CPU上的推理时间仅需SimpleClick的。

2 Related Work

Interactive Segmentation

深度网络被整合到交互式分割中[3, 12, 14, 45]，由DIOS[55]率先启动，随后在基于点击的方法如DEXTR[32, 41]，FCA-Net[34]，BRS[25]，以及f-BRS[46]中取得进展。接下来的方法专注于提高交互式分割的各个方面。SimpleClick[36]首次将大型的Vision Transformers[9]引入这个领域。

InterFormer[23]紧随其后，通过重用图像特征的新颖流程来减少模型冗余。SAM[28]也采用了这个流程，并实现了强大的零样本能力和多样的提示，导致了各种下游应用。

然而，SAM无法有效地采用图像级放大策略[46]和有效整合交互信息，这限制了其更广泛的应用。

作者引入了FocSAM来解决SAM的限制。

Efficient Attention

Transformer模型在计算机视觉领域取得了显著的进展。注意力模块的高计算复杂度引发了一系列研究[15, 38, 54, 60]。一种典型的方法是将每个 Token 的注意力区域从全注意力限制到局部/窗口注意力[17, 31, 38, 49]。这一策略引起了广泛关注，各种研究证明了这一点[7, 22, 24, 51, 56]。

最近，CSwin[8]引入了十字形窗口自注意力，以同时在两个方向上进行计算。Beyond Fixation[44]提出了DW-ViT来融合多尺度信息。在本文中，作者提出了Dwin-MSA，以对与目标相关的图像嵌入执行动态窗口注意力。

3 Method

作者提出了一个重新设计的SAM流程的FocSAM。在3.1节中，作者概述了SAM的流程和新流程。然后在3.2节和3.3节中，作者详细阐述了FocSAM聚焦精炼器的实现。最后，在第3.4节中讨论了训练损失。

Pipeline

SAM的流程。 在图2中，SAM [28] 包括一个图像编码器、一个提示编码器和解码器。图像编码器在交互前只预处理一次每个图像，尽管图像中的物体数量各不相同。相反，提示编码器和解码器积极参与每次交互，快速处理标注者的点击以预测分割结果。

图像编码器。 在SAM的预处理阶段，图像被调整大小并填充到，然后输入到基于ViT的图像编码器[9]中。这个编码器在四个等深度的阶段构建，每个阶段使用基于窗口的注意力进行有效计算[31]，在每个阶段的末尾应用全注意力。此后，简单的卷积层进一步降低维度，生成256维嵌入，对应于不重叠的图像块。

提示编码器。 在SAM的交互阶段，提示编码器[28]将标注者的提示转换为嵌入。这些提示包括第次交互中的个点击，每个点击带有坐标和一个指示正负的标签。在假阴性区域中的正向点击指示模型扩展该区域，而假阳性区域中的负向点击则建议移除。从每个目标的第二次交互开始，提示编码器还将之前预测的分割 Mask 转换为 Mask 嵌入。转换后的点击嵌入和 Mask 嵌入将被输入到SAM解码器中，如图2所示。

解码器。 在提示编码器之后，解码器接收图像嵌入， Mask 嵌入，点击嵌入和可学习的 Query 嵌入。 Query 嵌入的数量对应于解码器预期的输出 Mask 。在作者的工作中，作者使用单个 Query 嵌入。在解码过程中， ConCat 嵌入与 Mask 融合图像嵌入进行交叉注意力。它们在交叉注意力中交替充当 Query 和键/值的角色，而不涉及图像到图像的注意力。经过两个这样的交叉注意力块后，输出包括经过一些卷积上采样的点击融合图像嵌入和点击融合 Query 嵌入，点击嵌入被丢弃。它们的点积生成预测最终 Mask 的逻辑值。

FocSAM的流程。 在SAM流程的基础上，FocSAM的流程引入了焦点精炼器。精炼器每个目标使用一次。具体来说，在目标的第次交互中，精炼器接收图像嵌入，之前预测的 Mask 和之前的点击融合 Query 嵌入。然后，精炼器生成一个精炼的图像嵌入，其中包含与物体相关的嵌入。替换所有后续交互中此目标的原始嵌入。如图2（b）所示，这个焦点精炼器由一系列精炼块组成。这些块迭代地精炼图像和 Query 嵌入，共享相同的先前 Mask 。最终块的图像嵌入作为精炼器的输出。作者在下面的子节中详细介绍这些精炼块。

Refine Block

概述。在图3（a）中，普通的细化块和移位细化块在细化器中交替堆叠，使用共享 Mask 细化图像嵌入和点击融合的 Query 嵌入。它们共享大多数模块，主要区别在于Dwin和移位Dwin（图3（b)）。Dwin和移位Dwin都从 Mask 中识别围绕目标的边界框（图3（c)(d)）并在目标上细化嵌入。细化后的嵌入和相应复制的 Query 嵌入被送入MSA模块（图3（e)）。然后，作者详细介绍了Dwin和移位Dwin。

重新审视图像级放大。给定一个图像和一个边界框，图像级放大策略[46]被表述为，其中角坐标定义了边界框，是模型输入大小。将这种策略适应于嵌入通常涉及到RoIAlign[20]，它使用线性采样方法裁剪和调整嵌入的大小。然而，RoIAlign面临两个主要问题。首先，RoIAlign假设嵌入可以像图像一样线性插值，但由于缺乏相应的平滑性感知训练，这可能不适用于SAM的图像嵌入。其次，RoIAlign统一调整所有目标的大小，忽略了大小差异，这限制了较大目标的表示并为较小目标增加了冗余。

动态窗口。 作者不是使用RoIAlign，而是引入了动态窗口（Dwin）策略。给定窗口大小，一批个样本的图像嵌入可以窗口化为，其中。然后，选择与边界框相交的窗口（图3（c））。对于这些图像内的所有目标，作者可以同时选择与其各自边界框相交的所有窗口，而不管目标的大小。这导致了选定的嵌入窗口，其中是与框交互的窗口数。每个窗口在窗口内执行独立的计算，如自注意力，并用计算结果更新其自身的嵌入，冻结未选中的嵌入窗口。长距离 Patch 到 Patch 注意力。 作者进一步在Shift Dwin中采用移动策略[38, 39]（图3（d））。交替使用Dwin和Shift Dwin确保边界框内的所有 Patch 之间有足够的信息交换。此外，这些框通常限制了同一目标内嵌入的空间距离，意味着几个块和小窗口大小仍然允许足够的信息交换。

MSA模块。 MSA（图3（e））并行处理的每个窗口，具有复制的 Query 嵌入。令

表示传统的注意力[48]。MSA模块如下公式化。首先，与融合，_即_

然后，经过自注意力处理，得到

接下来，通过以下方式被 P-DyReLU 激活

最后，这个 MSA 模块输出

以及作为下一个块的输入。此外，通过平均求和的方式聚合每个窗口的。作者在下一小节详细描述 P-DyReLU。

Pixel-wise Dynamic ReLU

动态ReLU. DyReLU [6] 通过引入依赖于输入的激活参数来扩展传统的ReLU。对于输入向量，动态激活函数使用根据调整的参数。具体来说，传统的ReLU函数在DyReLU中推广为每个的元素的参数化分段线性函数。DyReLU根据调整系数和：

其中所有的系数都是超函数的输出。普通的ReLU是的特殊情况，具有和。

像素级DyReLU. 考虑到方程4，作者实现以融合方程2中的与方程3中的。这一实现受到SAM解码器使用图像和 Query 嵌入之间的点积来生成 Mask 预测的logits的启发[28]。这个过程有效地以像素级方式捕捉了每个图像嵌入与 Query 嵌入之间的非标准化相似性。作者采用这种相似性来增强与目标相关的嵌入并抑制不相关的嵌入，将表述为

其中复制以匹配图像嵌入，而执行空间平均池化。因此，系数与具有相同的形状。然后，作者在这些系数上应用通道MLP，以改变它们的比例和偏置，得到如下结果。

最后，方程式4中的P-DyReLU实现为

其中是逐元素乘积。

Training Loss

与先前的方法[23, 29, 36]一样，作者采用了在RITM [29]中提出的归一化Focal Loss（NFL）。此外，作者受到BRS [25]的启发，引入了点损失（PTL）作为辅助损失，其定义如下

其中是导致预测 Mask 的点击坐标，是指示点击是否为阳性的二进制标签。

4 Experiments

在第4.1节中，作者详细介绍了实验设置。第4.2节讨论了主要结果，比较了FocSAM在不同数据集上与先前方法的性能。在第4.3节中，作者统计评估了FocSAM在交互式分割中的稳定性，与SAM进行了比较。第4.4节探索了FocSAM模块的影响。最后，第4.5节展示了定性结果。

Experimental Setting

数据集。 与先前的方法[5, 23, 36, 37]一致，作者在COCO [33]和LVIS [16]上训练作者的模型，并评估所有方法在包括GrabCut [45]、Berkeley [26]、SBD [19]和DAVIS [43]在内的其他各种数据集上的零样本交互式分割能力。作者的评估还扩展到更具挑战性的数据集，包括MVTec [2]和COD10K [10]。有关数据集的更多详细信息，请参阅补充材料。

实现细节。 作者使用来自SAM [28]的预训练ViT-Huge作为主干网络，并配合提示编码器和解码器。对于所提出的焦点细化器，作者配置了总共个块，包括个普通细化块和个移位细化块。Dwin-MSA和P-DyReLU的嵌入维度均设置为与维的SAM图像嵌入对齐。Dwin-MSA的窗口大小设置为。细化步骤设置为，即焦点细化器在第二次点击后激活。更多详细信息可在补充材料中找到。

训练策略。 在训练FocSAM时，作者在损失计算之前采用Inter-Former的点击模拟策略[23]进行交互式模拟。在训练过程中，SAM的图像编码器和提示编码器被冻结。此外，作者使用图像编码器预先提取并存储COCO-LVIS图像嵌入以减少计算成本。作者将图像调整并填充以匹配SAM的输入大小。作者采用两阶段训练策略，首先以批量大小为的k迭代次数微调SAM解码器，然后在相同设置下以冻结的解码器训练FocSAM k迭代次数。这种策略解决了由于细化器的损失依赖于解码器而导致的训练不稳定问题。训练和评估在一台配备个NVIDIA RTX 3090 GPU和双Intel Xeon Silver CPU的服务器上进行。更多详细信息见补充材料。

评估。 在评估中，遵循SAM [28]，将图像调整并填充到，然后将解码器得到的分割结果调整回原始大小进行IoU计算。在测试中的点击模拟，作者按照先前的方法[5, 23, 36]，在错误预测区域的中心放置点击。每个点击的二进制标签由最大距离错误负区域和错误正区域边界决定。FocSAM在推理速度和分割性能方面进行评估。速度以CPU上的每点击秒数（SPC）量化，表示每次点击的平均推理时间。对于分割性能，作者使用点击次数（NoC）指标，即达到指定IoU所需的最小点击平均数。作者主要关注20次点击下的NoC@90，即达到 IoU所需的平均点击次数。如果需要超过20次点击，按照先前方法[5, 23, 36]进行评估一致性，点击次数上限为20。在消融研究中还采用了其他NoC指标。

Main Results

表1展示了FocSAM的主要结果，与先前方法进行了对比。实际上，尽管SAM在其实验[28]中是在SA-1B[28]上进行预训练，而不是用于这些方法的COCO+LVIS，但SAM已经与主流方法[5, 29, 36]进行了基准测试。SA-1B和COCO+LVIS都是为通用场景设计的，并且在范围上经常重叠，这有助于在SAM与这些方法之间进行有效比较。由于SA-1B包含了大量的由SAM生成的 Mask ，Focal SAM在COCO+LVIS上维持训练，以减轻SAM固有的偏见。

正如报告所述，FocSAM在六个评估数据集中的五个上取得了最先进的表现，尤其是在最大的SBD（6671个样本）和第二大的COD10K数据集（2026个样本）上。尽管在MVTec数据集上表现略逊一筹，但FocSAM仍然保持了所有数据集上的最佳平均NoC，与之前的最先进水平SimpleClick-ViT-H[36]相当接近。

然而，FocSAM最突出的特点是它的时间效率，SPC为0.39，远快于SimpleClick-ViT-H的6.99 SPC。这归功于FocSAM使用了SAM的处理流程，预提取图像嵌入以实现高效交互，而SimpleClick在每次交互时都要进行完整的模型推理。另一方面，尽管SAM的推理时间稍短，但与早期的FocalClick[5]等方法相比，其分割性能较低。FocSAM在交互式分割中将SAM的性能提升到顶级水平，同时只增加了约10%的计算成本。下面的子节将进一步验证FocSAM是否真正解决了SAM中的不稳定问题。

Stability Analysis

实验设置。 为了评估稳定性，作者在三个大规模数据集上进行了统计分析，即SBD、MVTec和COD10K。与对NoC指标的评估类似，每次点击都放置在错误预测区域的中心。每个样本的模拟点击数量从20增加到100。对于每个样本，从第二次交互点击开始，作者计算连续点击之间的IoU差异（IoU），并过滤掉大于-1%的IoU。这确保只考虑分割质量显著下降的情况。剩余的IoU随后被可视化。

结果。 如图4所示，与SAM相比，FocSAM在所有数据集上展现出明显的稳定性优势。FocSAM的IoU分布向右偏移，表明随着后续点击，分割质量下降的样本较少。尽管SAM偶尔会取得良好结果，但其固有的不稳定性通常需要额外的标注者交互来纠正错误。因此，在稳定性分析中，FocSAM在真实世界的交互效率方面代表了比SAM的稳定性提升。

Ablation Study

实验设置。 在消融研究中，作者评估了Dwin-MSA和P-DyReLU对FocSAM性能的个别影响。由于这些模块之间的相互依赖性，作者对模块进行了轻微修改。对于仅Dwin-MSA的情况，作者移除了所有的P-DyReLU模块，并将Dwin-MSA中的P-DyReLU激活替换为标准的ReLU。对于仅P-DyReLU的情况，作者移除了动态窗口以保留所有的图像嵌入，并移除了Dwin的注意力计算。作者使用20次点击内的NoC@90，并在100次点击内扩展到NoC@95进行更深入的分析，在这些三个最大数据集上评估了这些变体，包括SBD、MVTec和COD10K。这个NoC@95指标量化了每个模块在更具挑战性的样本上的个别贡献。所有的消融模型都按照主实验的相同协议进行训练。

结果。 表2显示，Dwin-MSA和P-DyReLU单独对FocSAM性能的贡献相似，表明它们提供了可比较的交互信息。Dwin-MSA主要关注最初预测的 Mask 以定位主要目标区域，类似于SAM中的边界框提示，而P-DyReLU利用初始点击来进行主要目标的勾勒。它们的交互信息是互补的。因此，它们的组合导致了整体性能的增强，特别是在100次点击内的NoC@95中最为明显。这个指标强调了对于挑战性样本实现95% IoU所需的点击次数的增加。Dwin-MSA和P-DyReLU的集成进一步稳定了FocSAM在挑战性样本上的性能。更多的消融研究在附录材料中提供。

Qualitative Results

在图5中，作者提供了一个具有挑战性的示例，对比了FocSAM和SAM，并可视化了四个不同点击位置的分段结果。这种可视化清楚地展示了FocSAM相对于SAM的增强稳定性。作者的定性分析确认，在这样一个具有挑战性的示例下，FocSAM保持了性能的一致性，与SAM相比，提供了更优的分段质量。额外的定性结果可在补充材料中找到。