本文首发于 CVHub，白名单账号转载请自觉植入本公众号名片并注明来源，非白名单账号请先申请权限，违者必究。
Title: Segment Anything in High Quality
PDF: https://arxiv.org/pdf/2306.01567v1.pdf
Code: https://github.com/SysCV/SAM-HQ
导读
SAM 拥有强大的零样本能力和灵活的提示功能，尽管它已经通过11亿个掩膜进行了训练，但在很多情况下，特别是处理结构复杂的对象时，其掩膜预测质量还是有所欠缺。
为此，本文提出了HQ-SAM，为SAM赋予了准确分割任何对象的能力，同时保持了SAM的原始设计、效率和零样本泛化能力。作者设计重用并保留了SAM的预训练模型权重，同时仅引入了最小的额外参数和计算。与此同时还设计了一个可学习的高质量输出令牌，该令牌被注入到SAM的掩膜解码器中，负责预测高质量的遮罩。该方法不仅将其应用于掩膜解码器特征，而且首先将其与初期和最终的ViT特征融合，以改善细节。为了训练引入的可学习参数，本文构建了一个包含44K细粒度掩膜的数据集。HQ-SAM仅在这个引入的44k 掩膜数据集上进行训练，这在8个GPU上只需4小时。
最终，本文在9个不同的分割数据集中展示了HQ-SAM的有效性，这些数据集覆盖了不同的下游任务，其中7个任务在零样本迁移中进行了评估。
引言
对多样化物体进行精确的分割是一系列场景理解应用的基础，包括图像/视频编辑、机器人感知，以及AR/VR等。"Segment Anything Model"（SAM）模型被设计出来，作为通用图像分割的基础视觉模型，它经过了十亿级别的掩膜标签的训练。SAM模型通过接受一个包含点、边界框或者粗糙掩膜的提示，可以在各种场景下分割出一系列物体、部分和视觉结构。尽管SAM模型取得了令人印象深刻的性能，但其分割结果在许多情况下仍然无法满足需求，特别是对于自动标注和图像/视频编辑任务，这些任务对于图像掩膜的精度有着极高的要求。
因此，作者提出了一个新的模型HQ-SAM，能够在保持原始SAM模型的零样本能力和灵活性的同时，预测出极高精度的分割掩膜。为了维护效率和零样本性能，作者在SAM模型上进行了微小的改动，只增加了少于0.5%的参数，以提升其高质量分割的能力。他们设计了一个可学习的HQ-Output令牌，该令牌被输入到SAM的掩膜解码器中，并且被训练以预测高质量的分割掩膜。此外，HQ-Output令牌在一个优化后的特征集上运行，以实现精确的掩膜细节。
为了学习精确的分割，需要一个包含精确掩膜注释的数据集。因此，作者构建了一个新的数据集，名为HQSeg-44K，该数据集包含了44K极其细粒度的图像掩膜注释，覆盖了超过1000种不同的语义类别。由于数据集规模较小和他们的最小化集成架构，HQ-SAM在8个RTX 3090 GPU上只需4小时就能完成训练。
为了验证HQ-SAM的有效性，作者进行了大量的定量和定性实验分析。他们将HQ-SAM与SAM在9个不同的分割数据集中进行比较，这些数据集涵盖了不同的下游任务，其中7个任务采用了零样本转移协议。严格的评估表明，与SAM相比，本文提出的HQ-SAM能够生成更高质量的掩膜，同时保持零样本的能力。
方法
HQ-SAM为了实现高质量的掩膜预测，将HQ-Output Token（高质量输出标记）和全局-局部特征融合引入到SAM中。为了保持SAM的零射能力，轻量级的HQ-Output Token复用了SAM的掩膜解码器，并生成了新的MLP（多层感知器）层来执行与融合后的HQ-Features（高质量特征）的逐点乘积。在训练期间，将预训练的SAM的模型参数固定，只有HQ-SAM中的少数可学习参数可以进行训练。
为了提升原SAM模型在零样本分割任务上的性能，同时保留其零样本的特性。HQ-SAM对SAM模型进行了两处关键的改动。
首先，作者在SAM模型的基础上引入了一个新的输出令牌（High-Quality Output Token）和全局-局部特征融合。HQ-Output Token可以更好地指导高质量的掩模生成，而全局-局部特征融合则可以将来自不同阶段的特征提取和融合，从而富化掩模特征的全局语义上下文和局部边界细节。
HQ-Output Token的引入对SAM模型的掩模预测能力进行了提升。在原SAM模型的设计中，掩模解码器使用一个输出令牌（类似于DETR中的对象查询）进行掩模预测。在HQ-SAM中，作者引入了一个新的可学习的HQ-Output Token，并加入了一个新的掩模预测层来进行高质量的掩模预测。
其次，全局-局部特征融合通过提取和融合SAM模型不同阶段的特征来提升掩模质量。具体来说，作者将SAM模型的ViT编码器的早期层次特征、ViT编码器的最后一层全局特征以及SAM模型掩模解码器的掩模特征进行了融合，生成了新的高质量特征（HQ-Features）。
基于ViT-L的SAM和HQ-SAM的训练和推理比较。HQ-SAM给SAM带来了微不足道的额外计算负担，模型参数的增加少于0.5%，并且达到了其原始速度的96%。SAM-L在128个A100 GPU上进行了180k次迭代的训练。基于SAM-L，只需要在8个RTX3090 GPU上训练HQ-SAM 4小时。
HQ-SAM的训练和推理过程是数据和计算高效的。在训练阶段，作者固定了预训练的SAM模型的参数，只对HQ-SAM中的新引入的可学习参数进行训练。在推理阶段，作者遵循了SAM的推理流程，但使用了HQ-Output token的掩模预测作为高质量的掩模预测。
总体来说，HQ-SAM相比原SAM模型，在提升分割质量的同时，训练过程更加高效，只需要4小时就可以在8块RTX3090 GPUs上完成训练。HQ-SAM也是非常轻量级，增加的模型参数、GPU内存使用和每图像推理时间都可以忽略不计。
实验
SAM与我们的HQ-SAM预测的掩膜比较，输入提示是相同的红色框或者在对象上的几个点。HQ-SAM产生了更为详细且具有非常准确边界的结果。在最右边的列中，SAM错误地解读了风筝线的细长结构，并且在输入框提示下产生了大量带有断裂孔洞的错误。
对四个极端细粒度的分割数据集进行的HQ-Output Token的消融实验。本文采用从它们的GT（Ground Truth，真值）mask转换过来的框作为框提示输入。默认情况下，通过计算全GT掩膜损失来训练HQ Output-Token的预测掩膜。
关于HQ-Features来源的消融实验。早期层（Early-layer）表示ViT编码器的第一个全局注意力块之后的特征，而最终层（final-layer）表示最后一个ViT块的输出。四个HQ数据集分别是DIS（验证集），ThinObject-5K（测试集），COIFT 和HR-SOD 。
模型微调或额外后处理的比较。对于COCO数据集，作者使用在COCO数据集上训练的最新水平的目标检测器FocalNet-DINO作为边界框提示生成器。
上图展示了COIFT和HRSOD在零样本协议下的召回率比较，使用了从宽松到严格的BIoU阈值。结果显示，当阈值从0.5变化到0.9时，SAM和HQ-SAM之间的性能差距显著增大。这表明HQ-SAM在预测非常准确的分割掩码方面具有优势，即HQ-SAM能够更准确地进行目标分割，尤其在对于严格的阈值要求下表现更好。
在UVO数据集上进行零样本开放世界实例分割的结果比较。为了生成边界提示，作者使用在COCO数据集上训练的FocalNet-DINO模型。其中，符号表示使用更严格的阈值来定义边界区域。
在高质量BIG基准测试集上进行零样本分割结果的比较。为了生成输入提示，作者使用了PSPNet来生成粗糙的掩码提示。通过比较不同类型的输入提示，对零样本分割结果进行了评估。
在COCO和LVISv1数据集上进行零样本实例分割结果的比较。对于COCO数据集，作者使用在COCO上训练的FocalNet-DINO模型进行检测，而对于LVIS数据集，则采用在LVIS数据集上训练的ViTDet-H作为它们的边界提示生成器。在SAM模型中，作者使用了ViT-L作为骨干网络，并使用了边界提示。作者在保持原始SAM的零样本分割能力的同时，改善了边界区域的掩码质量。
上图展示了SAM和HQ-SAM在零样本迁移设置下，给定相同的红色框或点提示时的视觉结果比较。从结果可以看出，HQ-SAM产生了明显更多保留细节的结果，并且还修复了掩码中的错误孔洞。相比之下，HQ-SAM在零样本迁移任务中能够更好地保留目标细节，并处理掩码中的错误。
上图展示了在COIFT（零样本）和DIS验证集上，使用不同数量的输入点进行交互式分割的结果比较。结果显示，HQ-SAM在各种点的数量下始终优于SAM，并且在提示模糊度较小的情况下，相对改进更为明显。这表明HQ-SAM在交互式分割任务中对于不同数量的输入点都具有更好的性能，尤其在输入点较少且提示不明确的情况下，HQ-SAM的改进效果更加明显。
上表展示了在HQ-YTVIS基准测试集上进行零样本视频实例分割的比较结果。在该比较中，作者使用了在YTVIS数据集上预训练的基于Swin-L的Mask2Fromer模型作为边界框提示输入，并且重复使用了其对象关联预测。通过这样的设计，作者对于零样本视频实例分割方法进行了评估和比较。
结论
本文提出了HQ-SAM，这是第一个通过对原始SAM引入可忽略开销而实现高质量零样本分割的模型，探讨了如何在数据高效和计算经济的方式下利用和扩展类似SAM的基础分割模型。作者在HQ-SAM中引入了一个轻量级的高质量输出标记，用于替换原始SAM的输出标记，以实现高质量的掩码预测。在仅使用了44K个高度准确的掩码进行训练后，HQ-SAM显著提升了SAM的掩码预测质量，而SAM本身是在11亿个掩码上进行训练的。作者对包括图像和视频任务在内的7个分割基准进行了零样本转移评估，涵盖了各种对象和场景。
作者： 派派星
文章来源：CVHub
推荐阅读
- YOLOv8太卷啦 | YOLOv8官方仓库正式支持RT-DETR训练、测试以及推理
- 软硬协同优化 (1)：高层综合编译技术
- 编译器优化那些事儿（10）：区域分析
- 编译器优化那些事儿（9）：Machine Outliner
- 编译器优化那些事儿（8）：指令调度概述
更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
