【论文笔记】SAM 2——革新视频与图像分割的交互式基础模型

原文链接：https://arxiv.org/abs/2408.00714

原文摘要：我们提出了 Segment Anything Model 2（SAM 2），这是一个用于解决图像和视频中提示驱动视觉分割的基础模型。我们构建了一个数据引擎，通过用户交互改进模型和数据，收集了迄今为止最大的视频分割数据集。我们的模型是一个简单的带有流式记忆的 Transformer 架构，用于实时视频处理。在我们的数据上训练的 SAM 2 在广泛的任务中表现出色。在视频分割中，我们观察到其准确性更高，且比以往方法少用3倍的交互。在图像分割中，我们的模型比 Segment Anything Model（SAM）更准确，速度提高了6倍。我们相信我们的数据、模型和见解将成为视频分割及相关感知任务的一个重要里程碑。我们将在开放许可下发布我们的主要模型、数据集以及用于模型训练和我们演示的代码。

一、研究背景

• 视频分割的挑战：
  • （1）运动、变形、遮挡、光照变化和其他因素，实体的外观可能会发生显著变化。
  • （2）相机运动、模糊和分辨率较低，视频通常比图像质量更低。
  • （3）高效处理大量帧是一个关键挑战。
• 现有交互式视频目标分割（iVOS）局限性：SAM 在视频帧上可能表现不佳，且没有机制可以交互式地修正模型的错误，除非在每一帧中重新使用 SAM 进行注释并重新启动跟踪。
• 视频分割数据集缺乏足够的覆盖范围，通常围绕特定目标类别（如人、车辆和动物）展开，无法实现“分割视频中任何内容”的能力。
• 实际应用需求推动：AR/VR、机器人、自动驾驶和视频编辑等领域，需要超越图像级分割的时间定位能力，要求模型能处理更复杂的视频数据，实现精准的时空分割。

二、核心任务：提示驱动的视觉分割（PVS）

PVS任务允许在视频的任何帧上向模型提供提示。提示可以是正向/负向的点击、框或掩码，用于定义要分割的对象或优化模型预测的对象。为了提供交互式体验，在接收到特定帧上的提示后，模型应立即响应。
在接收到初始提示（无论是在同一帧还是不同帧）后，模型应传播这些提示以获取整个视频中对象的掩码片段（masklet），在每个视频帧上定位目标的分割掩码。可以在任何帧上向模型提供额外的提示，以在整个视频中优化分割。

三、主要贡献

• SAM 2模型的创新：SAM 2是首个将提示分割任务扩展至视频领域的基础模型，能处理图像和视频，具备实时处理能力，适用于交互式分割和掩码传播。
• 数据引擎与大规模数据集：构建的数据引擎可收集大型多样化视频分割数据集SA-V，包含50.9K视频和642.6K掩码片段，为视频分割研究提供丰富资源。
• 性能提升与应用拓展：SAM 2在视频分割中准确性更高，交互次数减少3倍，图像分割速度比SAM快6倍且更准确，推动了视觉感知领域发展。

四、模型架构

对于给定的帧，分割预测是基于当前提示和/或先前观察到的记忆。视频以流式方式处理，图像编码器一次处理一个帧，并与来自先前帧的目标对象的记忆进行交叉注意。掩码解码器，它也可以选择性地接收输入提示，预测该帧的分割掩码。最后，记忆编码器转换预测和图像编码器嵌入（在图中未显示），以供未来帧使用。

1. 输入与输出

• 输入：图像或视频帧、提示（点击（正向/负向）、边界框、掩码）、先前记忆（Memory）
• 输出：图像/视频分割掩码、遮挡预测（通过一个额外的MLP头部实现，它为每个掩码生成一个分数，指示对象在当前帧中可见的可能性）、对象指针（基于掩码解码器的输出令牌）

2. 图像编码器

• 采用流式方法，随着视频帧的可用而逐个处理。
• 仅在整个交互中运行一次，其作用是提供代表每个帧的未条件化标记（特征嵌入）。

3. 记忆注意力机制

• 根据过去的帧特征和预测以及任何新提示来条件化当前帧特征。
• 堆叠L个Transformer块，第一个块以当前帧的图像编码为输入。每个块执行自注意力，然后是跨注意力到（提示/未提示）帧的记忆和对象指针，存储在记忆库中，然后是一个MLP。

4. 提示编码器

与SAM相同，可以通过点击（正向或负向）、边界框或遮罩来提示，以定义给定帧中对象的范围。稀疏提示通过位置编码与每种提示类型学习到的嵌入相加来表示，而遮罩则通过卷积嵌入并与帧嵌入相加。

5. 解码器

堆叠了“双向”transformer块来更新提示和帧嵌入：

• 对于模糊提示（例如，一个单击），可能有多个兼容的目标遮罩，预测多个遮罩。在视频中，模糊性可能延伸到视频帧之间，模型在每一帧上预测多个遮罩。如果没有后续提示解决模糊性，模型只传播当前帧预测IoU最高的遮罩。（同SAM）
• 对于可能不存在有效对象（例如，由于遮挡）的帧。增加了一个额外的头来预测当前帧上是否存在感兴趣的对象。（不同于SAM）
• 使用来自层次图像编码器的跳跃连接（绕过记忆注意力）来合并高分辨率信息进行遮罩解码。（不同于SAM）

6. 记忆编码器

使用卷积模块对输出掩码进行下采样，并将其元素级与图像编码器的未条件帧嵌入相加，然后通过轻量级卷积层融合信息。

7. 记忆库

• 维护最多N个最近帧的记忆FIFO队列和存储来自提示的最多M个提示帧的信息。
• 存储了一组对象指针作为轻量级向量，以获取要分割对象的高级语义信息，基于每帧掩码解码器输出令牌。
• 将时间位置信息嵌入N个最近帧的记忆，允许模型表示短期对象运动，但不嵌入提示帧，（因为提示帧的训练信号更稀疏，并且更难以泛化到推理设置中。）

训练策略：模型在图像和视频数据上联合训练，通过模拟交互式提示，使用8帧序列和纠正点击，提升模型性能。

五、数据引擎工作流程

• 阶段1：用SAM逐帧标注；
• 阶段2：引入SAM 2 Mask，结合SAM和其他工具生成掩码；
• 阶段3：使用完整SAM 2，提升标注效率和质量。
• 验证步：一组独立的标注者负责验证每个标注掩码片段的质量，将其标记为“满意”或“不满意”。不满意的掩码片段被送回标注流程进行细化。任何跟踪定义不清的对象的掩码片段被完全拒绝。

六、实验结果

1. 提示性视频分割

SAM 2在离线和在线评估设置中都优于SAM+XMem+++和SAM+Cutie。

2. 半监督视频对象分割

3. 图像分割

SAM 2在23个数据集上平均1-click mIoU达到58.9，高于SAM的58.1。

七、消融研究

1. 容量消融

• 输入尺寸：提高输入分辨率和帧数可以显著提升图像和视频任务的性能。（最终选择了1024的输入分辨率和8帧的默认设置。）
• 内存大小：增加内存数量有助于性能提升（最终选择了6个过去的帧作为默认值，减少内存通道对性能影响不大，同时显著降低了存储需求。）
• 模型大小：增加图像编码器和记忆-注意力的容量可以提高性能（默认使用B+图像编码器，它为速度和准确性提供了合理的平衡。）
  

2. 记忆构架消融

将记忆特征直接存储在记忆库中比先通过GRU处理更简单且高效，尽管在LVOSv2基准测试中GRU略有提升。此外，发现在其他帧中交叉关注掩模解码器输出的对象指针向量可以显著提高SA-V验证集和LVOSv2基准测试的性能。

八、小结

SAM 2将Segment Anything自然地进化到视频领域，基于三个关键方面：

• (i) 将提示分割任务扩展到视频，
• (ii) 将SAM架构应用于视频时使用记忆
• (iii) 多样化的SA-V数据集用于训练和基准测试视频分割。

---

补充本文用到的几个评价指标：

1. J&F（Jaccard & F-measure）指标

视频分割任务中的标准评估指标，用于衡量预测掩码与真实掩码之间的重叠程度。J&F值越高，表示模型的分割性能越好。

计算公式

• Jaccard Index（IoU）：
  衡量预测掩码与真实掩码的重叠程度。
  $$\text{IoU}=\frac{\text{TP}}{\text{TP}+\text{FP}+\text{FN}}$$
  其中：

  • TP（True Positive）：预测为正且实际为正的像素数；
  • FP（False Positive）：预测为正但实际为负的像素数；
  • FN（False Negative）：预测为负但实际为正的像素数。

• F-measure（F1分数）：
  衡量预测掩码的精确度和召回率的调和平均值。
  $$\text{F-measure}=2\times \frac{\text{Precision}\times \text{Recall}}{\text{Precision}+\text{Recall}}$$

2. mIoU（Mean Intersection over Union）

图像分割任务中常用的评估指标，用于衡量预测掩码与真实掩码之间的平均交并比。mIoU值越高，表示模型的分割准确性越好。

计算公式
对于每个类别 $c$，计算其交并比（IoU）：
$${\text{IoU}}_c=\frac{{\text{TP}}_c}{{\text{TP}}_c+{\text{FP}}_c+{\text{FN}}_c}$$
其中：

• ${\text{TP}}_c$：类别 $c$ 的预测为正且实际为正的像素数；
• ${\text{FP}}_c$：类别 $c$ 的预测为正但实际为负的像素数；
• ${\text{FN}}_c$：类别 $c$ 的预测为负但实际为正的像素数。

然后计算所有类别的平均交并比：
$$\text{mIoU}=\frac{1}{C}\sum _{c=1}^C{\text{IoU}}_c$$
其中 $C$ 是类别总数。

3. FPS（Frames Per Second）

衡量模型速度的指标，表示模型每秒可以处理的帧数。FPS值越高，表示模型处理视频的速度越快。

计算公式
$$\text{FPS}=\frac{1}{\text{平均处理时间（秒）}}$$
其中，平均处理时间是指模型处理一帧图像或视频所需的平均时间。