【论文阅读】Segment Anything

论文链接： https://arxiv.org/abs/2304.02643

Code： https://github.com/facebookresearch/segment-anything

来源： 2023 IEEE/CVF International Conference on Computer Vision (ICCV)

1. 引言

基础模型在NLP中的成功： 大规模预训练的语言模型（如GPT系列）利用大量网络文本数据，展示了强大的zero-shot和few-shot学习能力，即在没有专门针对某个任务的训练的情况下，也能完成新任务。这些模型通过提示工程（prompt engineering），即设计特定的文本提示，引导模型产生符合需求的回答。模型规模和训练数据越大，其性能越好，有时甚至可以与专门调优的模型相媲美或优越。

扩展到计算机视觉： 在视觉领域，也开始尝试类似的方法。例如，CLIP和ALIGN训练了图文编码器，能够把图像和文本对应起来，通过给模型设计的文本提示，实现对新视觉概念的zero-shot识别。这些模型还能与其他模块结合，用于图像生成等任务。

尽管取得了进展，但计算机视觉的问题比自然语言更复杂，有些问题缺乏大规模标注数据。因此，本文的目标是：

• 打造一种基础模型，专用于图像分割
• 让它具有promptable（可用提示引导）和zero-shot泛化的能力
• 让模型能在没有针对特定任务的训练数据情况下，通过设计提示解决新任务

解决方案：

• 定义promptable segmentation的任务：给出任何提示（空间位置或文本），模型输出对应的有效分割掩码。这相当于让模型“听懂”各种提示，进行对象分割。
• 训练过程中，搭建“数据引擎”——利用已有模型协助收集和标注数据，不断反哺模型，扩大数据规模。

总结来说，通过借鉴语言模型的思路，旨在通过设计一个支持多种提示的图像分割模型，实现zero-shot和few-shot的泛化能力，并构建了相关的任务、架构和数据链路，以实现这一目标。

2. Segment Anything任务

灵感来源：模型的训练方式借鉴了NLP中的预训练方法，即通过预测下一词（token）来学习丰富的语义和结构信息，然后利用“提示（prompt）”方法在多种任务中实现零样本（zero-shot）适应能力。

任务定义：在视觉分割中，提示可以是前景/背景点、粗略的边界框或掩码、自由文本，或者任何指示需要分割内容的信息。目标是给定任何提示，模型都能输出一个“有效”的分割掩码。这里的“有效”意味着，即使提示模糊或可能对应多个对象（例如区分“衬衫”或“人”），输出也应合理对应至少一个对象。

预训练策略：这种提示可分割任务启发了一种自然的预训练算法：让模型在训练中模拟一系列提示（点、框、掩码）并将预测的掩码与标注的真实掩码进行比较。这种方法不同于传统的交互式分割（只在用户大量输入后预测掩码），其目标是让模型始终能对任何提示生成合理的“有效”掩码，即便提示模糊或多义。

零样本迁移能力：预训练使模型在推理（inference）时可以根据不同任务的提示，自动适应不同的分割需求。例如，利用已训练的边界框检测器提供定位信息，将其作为提示识别具体某类对象的实例。

区别于多任务模型：传统的多任务系统在训练和测试时执行固定任务（如语义分割、实例分割、全景分割等），而本文所描述的“提示可分割”模型的特点是训练后可以通过提示在不同任务间灵活转换，实现任务通用化，这极大增强了模型的适应性和扩展性。

3. Segment Anything模型

SAM基于Transformer视觉模型（如ViT）构建，经过特殊设计以实现实时性能。整体架构包括：图像编码器，提示编码器，和掩码解码器。图像编码器在每张图像上运行一次，然后生成高质量的图像嵌入，供后续步骤使用。提示编码器接受各种输入提示（点、框、文本、掩码）并将它们转换为嵌入表示。掩码解码器结合图像嵌入和提示嵌入，快速生成目标掩码。

• 图像编码器（Image encoder）：

  • 采用预训练的MAE（Masked Autoencoder）结合Vision Transformer（ViT），经过微调以适应高分辨率输入。
  • 作用：提取图像的深层次特征表示，只需一次运行，便于后续交互操作。

• 提示编码器（Prompt encoder） 处理两类提示：

  • 稀疏提示（Points, boxes, text）：
    • 点和框通过位置编码（positional encodings）加上学习的嵌入表示。
    • 文本用现成的CLIP模型编码。
  • 密集提示（Masks）：
    • 通过卷积进行嵌入，并与图像嵌入逐元素相加，形成丰富的提示信息。

• 掩码解码器（Mask decoder）：

  • 负责将图像嵌入和提示嵌入结合起来，快速预测目标区域的掩码。
  • 采用改良的Transformer解码器：
    • 包含提示自注意力（prompt self-attention）
    • 跨注意力（cross-attention）——用于融合提示与图像特征
    • 最后通过多层感知机（MLP）和动态线性分类器，将输出映射到每个像素的前景概率。

• 解决歧义（Resolving ambiguity）：

  • 对于模糊或多义提示（如“shirt”可能指衣服或人），模型会预测多个掩码（一般为3个），代表不同的可能对象。
  • 训练时，仅对最优掩码应用损失（如最小损失），为模型提供处理歧义的能力。
  • 模型会为每个掩码预测置信度（估计IoU），以便排序。

• 效率（Efficiency）：

  • 设计目标是快速响应，给定预先计算的图像嵌入后，提示和掩码预测在浏览器（CPU上）大约50毫秒内完成，支持交互式应用。

• 训练（Losses and training）：

  • 使用Focal loss和Dice Loss的线性组合，监督掩码预测。
  • 模拟交互操作：在训练过程中，随机采样11轮提示，使模型学习在多轮交互中的表现。

4. Segment Anything数据引擎

构建SA-1B数据集的三阶段的具体操作和技术细节：

• 手动标注阶段（Assisted-manual stage）：
  • 类似于传统的交互式分割方法，专业标注员使用由SAM驱动的浏览器工具，通过点击前景/背景点来标注掩码。
  • 标注员还可以使用“画笔”和“橡皮擦”工具进行细化。
  • 该过程实现了实时交互（利用预先计算的图像嵌入），提升效率。
  • 标注忽略了语义限制，标注员可以自由标注“物体”和“材料”类别，但未收集名称或描述。
  • 初始阶段，SAM基于公开数据集训练，随着数据量增加，模型在逐轮重新训练中不断优化，平均每个掩码标注时间从34秒缩短到14秒，比COCO标注快得多。
• 半自动阶段（Semi-automatic stage）：
  • 目标是增加掩码多样性，提升模型的“anything segmentation”能力。
  • 先用训练好的边界框检测器自动识别“有信心的掩码”。
  • 标注员在这些预填充掩码的基础上，补充未标注的对象。
  • 通过反复训练SAM，模型性能逐步提升。
  • 在这一阶段收集到约5.9百万掩码，总计约10.2百万掩码，平均每图标注72个掩码。
• 全自动阶段（Fully automatic stage）：
  • 利用经过大量多样化掩码训练、具备“模糊理解”能力（ambiguity-aware）的模型完全自动生成掩码。
  • 以一个32×32的规则点阵作为提示，模型为每个点预测多个可能对应的掩码（包括子部分、部分和整体物体）。
  • 利用IoU估算模块筛选可信且稳定的掩码（通过概率阈值和一致性检测）。
  • 采用非极大值抑制（NMS）过滤重复掩码。
  • 进一步利用重叠的放大图像裁剪，提升小目标掩码的质量。
  • 最终，自动处理了全部11M图像，产生了总计约1.1亿个高质量掩码（SA-1B dataset）。

5. Segment Anything数据集

数据集概述：SA-1B包含11百万（11M）张高质量、多样化、已授权且保护隐私的图片，以及1.1十亿（1.1B）个自动生成的分割掩码。这些掩码由我们的“数据引擎”自动收集产生。该数据集旨在推动计算机视觉基础模型的发展。

• 图片方面：
  • 这些图片由合作摄影师提供，具有高分辨率（平均约3300×4950像素），但大数据量带来存储和可访问性挑战。
  • 为解决此问题，我们同时发布了缩小版（最短边设置为1500像素）的图片，虽然减小了分辨率，但仍然远比许多现有数据集（如COCO，约480×640像素）更加高清。
  • 出于隐私和安全考虑，图片中的面孔和车牌已被模糊处理。
• 掩码方面：
  • 绝大多数掩码（99.1%）由自动化模型生成。
  • 这些掩码的质量是关键，我们通过随机抽取500张图片（约5万个掩码），让专业标注人员对自动生成的掩码进行细节修正，从而得到“专业校正”版本与自动掩码比较。
  • 评估显示，94%的自动掩码与校正掩码的交并比（IoU）超过90%，而97%的超过75%%。这一水平与不同标注者之间的共识（85-91%）相当，表明掩码质量较高。

6. 零样本迁移实验

6.1 零样本单点有效掩模评估

任务定义：用户只提供一个点作为提示，旨在分割对应的目标对象。这个任务本身具有不确定性，因为一个点可以对应多个对象。此外，大多数数据集中的标签没有列出所有可能的掩码，所以自动评价指标（如IoU）可能不够精确，因此作者采用了人类评估。人类评估的标准从1（无意义）到10（像素级完美）打分。

• 评估方式：
  • 采样点：默认从目标掩码的“中心”位置采样，这样符合交互式分割的标准评测方法。
  • 比较模型：除了SAM，还与强基线RITM进行了比较。SAM会生成多个候选掩码，但只用最有信心的那个掩码进行评估。
• 数据集：
  • 使用一个包含23个多样化图像分布数据集的全套评估，包括常规指标（mIoU）和人类评分。
  • 只用部分数据集进行人类评分，以免评估资源过大。
• 结果分析：
  • 自动评估（mIoU）：SAM在16个数据集上优于其他方法，最大提升到约47点IoU。在“oracle”条件下（选择最符合地面真实的掩码）SAM在所有数据集上都优于RITM，说明SAM的多掩码输出具有潜力，但存在歧义问题。
  • 人类评分：无论在哪个数据集，评估者都认为SAM的掩码质量明显优于RITM，平均分在7到9之间，表明其生成的掩码大部分都可以认为是“合理、细节良好”。即使在自动指标表现较弱的数据集如DRAM和IBD中，人类评估仍给出较高评分。
  • 变体比较：增加提示点数量（从1到9）会降低模型之间的差距，因为任务变简单。随机采样点时，SAM表现稳健，甚至优于基线方法。

6.2 零样本边缘检测

将SAM应用到零样本边缘检测任务中的方法和结果，具体如下：

• 方法（Approach）：
  • 任务目标：在经典的低级视觉任务——边缘检测上评估SAM的能力。
  • 数据集：使用BSDS500数据集，该数据集包含手工标注的边缘真值图。
  • 生成边缘图的过程：
    • 首先，将一个16×16的规则网格（共256个点）作为输入提示点喂给SAM，每个点生成3个掩码，共计768个预测掩码。
    • 移除冗余掩码：使用非极大值抑制（NMS）技术。
    • 计算边缘图：对未阈值化的掩码概率图使用Sobel滤波算子产生边缘线，然后经过轻量级的后处理步骤（包括边缘NMS）得到最终的边缘图。
• 结果（Results）：
  • 可视化分析：尽管SAM没有专门训练进行边缘检测，却能够生成合理的边缘图。与真值边缘相比，SAM输出的边缘更丰富，包含许多未标注的细节边缘。
  • 定量指标：主要指标是“在50%精度下的召回率（R50）”，即在不同的边缘阈值中达成50%假阳性比例时的召回率。结果显示，SAM在R50上表现较高，但在精确度方面有所欠缺，表现出一定的偏差。
  • 性能方面，SAM略落后于专门训练的最先进边缘检测方法（如HED），但优于早期的零样本传输方法，显示其在无任务特定训练下也能产生具有参考价值的边缘图。

6.3 零样本目标检测

将SAM应用于中层视觉任务——目标提名（object proposal）生成，以及与基线方法的性能对比，具体内容如下：

• 方法（Approach）：
  • 任务目标：评估SAM在目标提名任务中的表现。这一任务在目标检测研究中非常重要，通常作为早期系统中中间步骤，旨在产生候选目标区域。
  • 生成方法：
    • 使用经过轻微修改的自动掩码生成流程。
    • 将生成的掩码作为候选目标提名（proposals）输出。
  • 评价指标：
    • 使用标准的平均召回率（AR）指标，计算在LVIS v1数据集上的表现。
    • 选择LVIS数据集，是因为其类别丰富，具有挑战性。
  • 基线系统：
    • 使用由ViTDet（携带级联Mask R-CNN的ViT-H模型）实现的“强大基线”。
• 结果（Results）：
  • 基于ViTDet-H的检测器作为提名器（即DMP方法）表现最佳，能充分利用检测结果的优势。
  • SAM：
    • 在中等和大型目标、稀有类别目标以及常见类别目标上，性能优于ViTDet-H。
    • 仅在小目标和频繁类别目标上略逊于ViTDet-H。原因是ViTDet-H在LVIS数据集训练，能学习到LVIS特有的类别偏差，而SAM作为零样本模型，没有在LVIS专门训练，因此在这些类别表现稍差。
      

6.4 零样本实例分割

将SAM集成到更高层次的视觉系统中进行实例分割，以及对其性能和质量的评估，主要内容如下：

• 方法（Approach）：

  • 先运行一个目标检测器（这里使用之前提到的ViTDet）以检测出目标边界框。
  • 将检测得到的边界框作为提示（prompt）输入到SAM中。
  • 框架示例：通过“串联”检测器和SAM，实现实例分割。

• 结果（Results）：在COCO和LVIS两个数据集上，比较了SAM和ViTDet预测的掩码

  • Mask AP（平均精度）指标显示两者存在差距，ViTDet性能略优，但差距不大。
  • 通过可视化分析，发现SAM的掩码边界更清晰、更具“锐度”。
  • SAM在人类评价中持续优于ViTDet。

  

  • 在COCO数据集上，由于掩码AP差距较大且真实标注质量较低（由人类研究支持），ViTDet可能学到了COCO掩码的特定偏差，即针对训练数据调优的“偏见”。SAM作为零样本方法，没有这些偏差，使得它在掩码质量上表现更佳（尤其在人类主观质量评估方面）。
    

6.5 零样本文本分割

SAM在处理基于自由文本的对象分割任务中的尝试，从方法到结果如下：

• 方法（Approach）：

  • 任务：更高层次的目标——用自然语言文本描述进行对象分割。
  • 样本准备：
    • 使用手动收集的掩码（面积大于100²像素）对应的图像。
    • 提取这些图像对应的CLIP（Contrastive Language-Image Pretraining）图像嵌入（image embeddings）。
  • 训练过程：
    • 在训练时，将从图像中提取的CLIP图像嵌入作为“prompt”输入给SAM（首次交互）。
    • 关键点：因为CLIP的图像嵌入是与文本嵌入对齐训练的（即图像和文本在潜在空间中相互对应），
  • 在推理阶段，将自然语言文本输入CLIP的文本编码器，得到对应的文本嵌入，再作为prompt输入给SAM。

• 结果（Results）：

  • SAM可以根据简单的文本如“a wheel”、“beaver tooth grille”进行对象分割。
  • 当只用文本提示时系统偶尔会选错目标。
  • 但加入额外的点（点提示）一般能改善预测表现，类似于其他研究中的“点提示”方法。
    

6.6 消融实验

对其训练流程和数据规模进行的消融实验（ablation studies），重点在于不同训练策略和数据量对SAM性能的影响。

• 单点提示协议的消融研究：
  • 实验中采用单点（center point）提示，意味着每次只用一个点引导模型进行分割。
  • 单点可能存在歧义，且该歧义可能未反映在 ground truth（标注的真实掩码）中，因为ground truth通常只有一个唯一的掩码。
• 偏差和“oracle”的引入：
  • 由于SAM在“零样本（zero-shot）”的转移场景中，其最优预测（top-ranked mask）可能会偏离标注指南中定义的ground truth掩码。
  • 因此，除了模型的top预测外，还报告“oracle”结果，即：在所有模型输出中选择最接近ground truth的掩码（最佳掩码），用于评估模型的潜力。
• 训练数据的阶段性（cumulative data from不同“阶段”）：
  • 展示在逐步加入不同阶段的训练数据后，模型的性能（以mIoU衡量）逐步提升。
  • 每个阶段都带来性能的提升，说明数据的丰富性有助于模型表现。
• 自动掩码与手工/半自动掩码的比例：
  • 训练数据中，自动生成的掩码远多于手动或半自动标注的掩码。
  • 为了平衡数据，用手工和半自动掩码进行10倍的过采样（oversampling），取得最优效果。
  • 但这使训练变得复杂。
• 只用自动掩码的简化策略：
  • 为简化训练流程，作者也试验了只用自动掩码的方案。
  • 结果显示，性能（mIoU）只略低于使用全部数据（差异约0.5），证明使用自动掩码训练效果还不错。
• 数据量影响：
  • 完整的SA-1B数据集包含约11M图像，为了实验，均匀抽样得到不同规模的子集：1M和0.1M图像。
  • 在仅使用0.1M（约10%的全部数据）时，模型性能明显下降。
  • 使用1M图像（大约10%数据规模）时，模型性能（mIoU）接近完整数据集的表现。

A. Segment Anything模型细节

SAM（Segment Anything Model） 中 图像编码器（Image encoder） 和 提示编码器（Prompt encoder） 以及轻量级的 掩码解码器（Mask decoder） 的设计与功能，具体如下：

• 图像编码器（Image encoder）：

  • 作用：将输入图像转换为C×H×W的图像嵌入（image embedding）。这是一种紧凑的表征，用于后续的提示融合和掩码预测。
  • 实现方案：采用预训练的MAE（Masked Autoencoder）的Vision Transformer（ViT），特别是ViT-H/16模型，进行了最小改动，支持高分辨率输入（1024×1024像素）。
  • 处理流程：
    • 输入图像经过缩放和补边，使分辨率为1024×1024，并生成大小为64×64的特征映射（因为采用了一次64×下采样）。
    • 提取的特征通过1×1卷积将通道数减少到256，随后通过3×3卷积保持256通道，每个卷积层后接层归一化（layer normalization）。
  • 重要性：该图像嵌入在整个流程中只需计算一次，避免重复计算以支持实时交互。

• 提示编码器（Prompt encoder）：

  • 作用：将不同类型的提示（稀疏提示：点和框；密集提示：掩码；文本提示）编码成向量，以便与图像嵌入结合。
  • 具体实现：
    • 点：用位置编码（positional encoding）叠加到点的坐标上，再加上表示“前景”或“背景”的可学习嵌入。
    • 框：用两个边界的角点位置编码，加上对应的“左上角”或“右下角”嵌入。
    • 文本：用CLIP的文本编码器处理。
    • 密集掩码（mask prompt）：在分辨率比输入图像低4倍的空间上输入，然后通过两个步幅为2的卷积逐步缩小4倍，最后用1×1卷积映射到256维。若无掩码提示，则加入表示“无掩码”的嵌入。

• 轻量级掩码解码器（Mask decoder）：

  • 作用：结合图像嵌入和提示嵌入（tokens）预测目标掩码。
  • 结构：
    • 使用类似Transformer的解码器：包含自注意（self-attention）、跨注意（cross-attention）和MLP（多层感知机）模块。
    • 在解码过程中，将一个“输出”token（类似于分类任务中的[cls] token）加入提示集合。
    • 每一层解码器包括：
      1. 自注意：在tokens间进行
      2. 跨注意：用tokens（查询）对图像嵌入（作为key和value）进行
      3. MLP：更新tokens信息
      4. 反向跨注意：用图像嵌入（查询）对tokens（作为key/value）进行更新
  • 多层解码后，将更新的图像嵌入进行放大（Upsample 4×），并通过MLP和点乘操作，预测出最终的掩码。

• Transformer的嵌入维度：

  • 所有Transformer层（包括自注意和交叉注意层）中的嵌入空间使用的向量维度为256。这意味着每个token（如提示向量或图像特征）在空间中表示为256维的向量。

• MLP块的规模与应用范围：

  • Transformer中的MLP（多层感知机）块具有较大的内部维度，为2048。这意味着每个MLP层的隐藏层有2048个神经元，增加模型的表达能力。
  • 但为了计算效率，这个大规模的MLP仅应用于prompt tokens（提示token），而提示token一般相对较少（通常少于20个），因此在实际中处理量有限。

• 交叉注意层中的通道维度缩减：

  • 在交叉注意（cross-attention）层中，处理的是64×64的图像嵌入（即空间尺寸，空间区域为4096个位置）。
  • 为了提高计算效率，将查询（queries）、键（keys）和值（values）中的通道数减半，从256减到128。这是因为在这种注意机制中，将通道数减半可以显著减少矩阵乘法的计算量，同时对性能影响较小。

• 注意力头数：

  • 所有注意力层（self和cross）都使用8个注意力头（heads），这允许模型在不同子空间中并行关注不同的特征，从而提升建模能力。

• 反卷积（transposed convolution）用于上采样：

  • 用于将输出的图像嵌入放大（Upsample）以恢复到更高的空间分辨率。
  • 具体操作：使用2×2的转置卷积核（kernel size = 2×2），步长为2（stride = 2），以实现放大2倍的空间尺寸（比如64×64变为128×128）。
  • 输出通道数：设为两层，分别为64和32通道，这决定了上采样后图像的深度（特征维度）。

• 激活函数：使用GELU（高斯误差线性单元）激活，确保非线性变换。

• 转置卷积层之间还加入了层归一化（layer normalization），以稳定训练过程并改善模型性能。