SAM 2: Segment Anything in Images and Videos
我们推出了分割一切模型2(SAM 2),这是一个旨在解决图像和视频中可提示的视觉分割问题的基础模型。我们构建了一个数据引擎,通过用户交互来改进模型和数据,从而收集到了迄今为止最大的视频分割数据集。我们的模型采用了简单的Transformer架构,并配备了流式内存,可实现实时视频处理。在我们的数据上训练的SAM 2在广泛的任务中表现出强大的性能。在视频分割方面,与先前的方法相比,我们观察到其精度更高,且所需的交互次数减少了三倍。在图像分割方面,我们的模型比分割一切模型(SAM)精度更高,速度快六倍。我们相信,我们的数据、模型以及见解将成为视频分割及相关感知任务的一个重要里程碑。我们将发布我们的主模型、数据集,以及用于模型训练的代码和演示。
1 Introduction
“分割一切”(SA)引入了一个用于图像中可提示分割的基础模型(基里洛夫等人,2023)。然而,图像仅仅是现实世界的一个静态快照,在现实世界中,视觉片段可能会表现出复杂的运动,并且随着多媒体内容的快速增长,现在很大一部分内容都记录了时间维度的信息,尤其是在视频数据中。增强现实/虚拟现实、机器人技术、自动驾驶汽车以及视频编辑等领域的许多重要应用,除了图像级别的分割之外,还需要进行时间定位。我们认为,一个通用的视觉分割系统应该既适用于图像,也适用于视频。
视频分割旨在确定实体在时空上的范围,这带来了超越图像分割的独特挑战。由于运动、变形、遮挡、光照变化以及其他因素,实体的外观可能会发生显著变化。由于相机运动、模糊以及较低的分辨率,视频的质量通常比图像要低。此外,高效处理大量帧也是一个关键挑战。虽然“分割一切”(SA)成功解决了图像分割的问题,但现有的视频分割模型和数据集在实现“分割视频中的一切”这一能力方面仍存在不足。
我们推出了分割一切模型2(SAM 2),这是一个适用于视频和图像分割的统一模型(我们将图像视为单帧视频)。我们的工作包括一项任务、一个模型和一个数据集(见图1)。
我们专注于可提示视觉分割(PVS)任务,该任务将图像分割推广到视频领域。这项任务将视频任意一帧上的点、框或掩码作为输入,以定义一个感兴趣的分割区域,进而预测其时空掩码(即“掩码小块”)。一旦预测出掩码小块,可通过在其他帧中提供提示来对其进行迭代优化。
我们的模型(第4节)可生成单张图像以及跨视频帧的感兴趣对象的分割掩码。SAM 2配备了一种存储器,该存储器可存储有关对象和先前交互的信息,这使得它能够在整个视频中生成掩码小块预测,并且还能基于先前观察到的帧中对象的存储记忆上下文有效地校正这些预测。我们的流式架构是将SAM自然推广到视频领域的一种方式,它一次处理一个视频帧,并配备了一个记忆注意力模块,用于关注目标对象的先前记忆。当应用于图像时,存储器为空,并且该模型的表现与SAM类似。
我们使用一个数据引擎(第5节),通过让我们的模型与注释员协同工作来生成训练数据,以交互式地注释新的、具有挑战性的数据。与大多数现有的视频分割数据集不同,我们的数据引擎并不局限于特定类别的对象,而是旨在为分割任何具有有效边界的对象(包括对象的部分和子部分)提供训练数据。与现有的模型辅助方法相比,我们让SAM 2参与其中的数据引擎在质量相当的情况下速度快8.4倍。我们最终的 “分割一切视频”(SA-V)数据集(第5.2节)由50900个视频中的3550万个掩码组成,其掩码数量比任何现有的视频分割数据集多53倍。SA-V数据集具有挑战性,其中包含在整个视频中被遮挡后又重新出现的小对象和对象的部分。我们的SA-V数据集在地理来源上具有多样性,并且对SAM 2的公平性评估表明,在基于感知的性别进行视频分割时性能差异极小,并且在我们评估的三个感知年龄组之间差异也很小。
我们的实验(第6节)表明,SAM 2在视频分割体验上实现了质的飞跃。SAM 2能够在交互次数比先前方法少三倍的情况下,实现更高的分割精度。此外,在多种评估设置下,SAM 2在已有的视频目标分割基准测试中优于先前的研究成果。并且,在图像分割基准测试中,SAM 2的性能比SAM更好,速度快六倍。通过众多的零样本基准测试(其中17个用于视频分割,37个用于单张图像分割)可以看出,SAM 2在各种视频和图像数据分布上都十分有效。
我们将在宽松的开源许可下发布我们的研究成果,包括SA-V数据集(知识共享署名4.0许可协议)、SAM 2模型的检查点,训练代码(遵循Apache 2.0许可协议),以及交互式在线演示的代码(遵循Apache 2.0许可协议)。
2 Related work
Image segmentation.“分割一切”(基里洛夫等人,2023)引入了一种可提示的图像分割任务,其目标是在给定输入提示(如指代感兴趣对象的边界框或点)的情况下,输出有效的分割掩码。在SA-1B数据集上训练的“分割一切”模型(SAM)支持零样本分割,这使得它能够广泛应用于各种场景。最近的研究对SAM进行了扩展,例如,引入高质量输出标记以在细粒度掩码上进行训练(柯等人,2024),或者提高SAM的效率(熊等人,2023;张等人,2023a;赵等人,2023)。从更广泛的角度来看,SAM被应用于众多领域,包括医学成像(马等人,2024;邓等人,2023;马祖罗夫斯基等人,2023;吴等人,2023a)、遥感(陈等人,2024;任等人,2024)、运动分割(谢等人,2024)以及伪装物体检测(唐等人,2023)。
Interactive Video Object Segmentation (iVOS).交互式视频目标分割已成为一项关键任务,旨在在用户引导下(引导形式通常为涂鸦、点击或边界框)高效地获取视频中的目标分割结果(掩码小块)。一些早期方法(王等人,2005;白和萨皮罗,2007;范等人,2015)采用基于图的优化方法来指导分割注释过程。较新的方法(赫奥等人,2020;程等人,2021b;德拉托拉斯等人,2024)通常采用模块化设计,将用户输入转换为单帧上的掩码表示,然后将其传播到其他帧。
基于点击的输入在交互式视频分割中更易于收集(霍马云法尔等人,2021)。近期的研究工作结合了图像上的“分割一切”模型(SAM),并使用了基于掩码的视频跟踪器(程等人,2023b;杨等人,2023;程等人,2023c)或基于点的跟踪器(拉伊奇等人,2023)。然而,这些方法存在局限性:跟踪器可能并非对所有目标都有效,“分割一切”模型(SAM)在视频帧上的表现可能不佳,并且除了在每一帧中使用“分割一切”模型重新注释并从该帧重新开始跟踪之外,没有其他机制来交互式地修正模型的错误。
我们的工作与这些研究有着相似的目标,即交互式地分割视频中的目标。为实现这一目标,我们构建了一个强大的统一模型,该模型可直接接收用于交互式视频分割的提示,同时还构建了一个规模庞大且多样化的数据集。
Video Object Segmentation (VOS).视频目标分割(VOS)任务以第一帧中的目标掩码作为输入,并且必须在整个视频中准确跟踪该目标(蓬特 - 图塞特等人,2017)。这项任务被称为 “半监督视频目标分割”,因为输入的掩码可以被视为仅在第一帧中可用的目标监督信号。由于其在视频编辑或机器人技术等应用中的相关性,该任务受到了广泛关注。
早期基于深度学习的方法通常在视频的第一帧上进行在线微调,或者在所有帧上进行微调,以使模型适应目标对象。通过离线训练的模型实现了更快的推理,这些模型要么仅以第一帧为条件,要么还整合前一帧的信息。 这种多条件方法已通过循环神经网络(RNN)和Transformer扩展到所有帧。
半监督视频目标分割可以被视为我们的可提示视觉分割(PVS)任务的一个特殊情况,即在视频的第一帧中仅使用掩码提示。值得注意的是,在视频目标分割任务中,为第一帧标注所需的高质量目标掩码在实际操作中既具有挑战性又耗时。
Video segmentation datasets.已经有许多数据集被提出用于支持视频目标分割(VOS)任务。早期的VOS数据集,比如DAVIS数据集,包含了高质量的标注,但其规模限制了基于深度学习的方法的应用。YouTube-VOS数据集是第一个用于VOS任务的大规模数据集。随着算法的不断改进以及基准测试性能开始趋于饱和,研究人员开始关注通过特别聚焦于遮挡情况、长视频、极端变换、目标多样性或场景多样性来增加VOS任务的难度。
我们发现,当前的视频分割数据集缺乏足够的覆盖范围,无法实现 “分割视频中的任何对象” 的能力。这些数据集的标注通常只涵盖完整的对象(而非对象的部分),并且数据集往往围绕着特定的对象类别,例如人、车辆和动物。与这些数据集相比,我们发布的SA-V数据集不仅关注完整的对象,还广泛涵盖了对象的各个部分,并且包含的掩码数量比其他数据集多一个数量级以上。
3 Task: promptable visual segmentation
我们的可提示视觉分割(PVS)任务允许在视频的任意一帧上向模型提供提示。这些提示可以是正/负点击、边界框或者掩码,既可以用来定义要分割的对象,也可以用来优化模型预测的结果。为了提供交互式体验,当在特定帧上接收到提示时,模型应立即给出该帧上对象的有效分割掩码。在接收到初始提示(可以在同一帧上,也可以在不同帧上)后,模型应传播这些提示,以获取整个视频中对象的掩码小块,从而在视频的每一帧上定位目标的分割掩码。可以在视频的任意一帧上向模型提供额外的提示,以便在整个视频中优化分割结果(示例见图2)。有关该任务的详细信息,请参阅第B节。
我们将SAM 2(第4节)应用于PVS任务以构建我们的SA-V数据集(第5节)。我们在多个帧上模拟交互式视频分割场景,以在传统的半监督VOS(视频目标分割)设置中评估该模型(第6节),其中注释仅限于第一帧,并且还在SA基准测试中进行图像分割。
使用SAM 2进行交互式分割。步骤1(选择):我们在第1帧提示SAM 2以获得目标对象(舌头)的分割。绿色/红色点分别表示正/负提示。SAM 2自动将分割传播到后续帧(蓝色箭头),形成一个掩码片段。如果SAM 2在第2帧之后丢失目标对象,我们可以通过在新帧中提供额外的提示(红色箭头)来纠正掩码片段。步骤2(细化):在第3帧中单击一次就足以恢复目标对象并将其传播以获得正确的掩码片段。如果采用解耦的SAM+视频跟踪器方法,则需要在第3帧中多次点击(如同在第1帧中一样),才能正确重新标注目标对象,因为分割是从头开始的。凭借SAM 2的记忆功能,单击一次就可以恢复舌头。
4 Model
SAM 2(图3)可以被视为SAM在视频(和图像)领域的扩展,它在各个帧上接受点、框和掩码提示,以在时空上定义要分割的对象的范围。在空间上,该模型的行为与SAM类似。一个可提示的轻量级掩码解码器接收图像嵌入和提示(如果有),并为该帧输出一个分割掩码。可以通过在帧上迭代添加提示来细化掩码。
SAM 2解码器所使用的帧嵌入并非直接来自图像编码器,而是基于过去的预测和被提示帧的记忆进行条件化。相对于当前帧,被提示帧也可以“来自未来”。帧的记忆是由记忆编码器根据当前预测创建的,并存储在记忆库中,供后续帧使用。记忆注意力操作会将图像编码器的每帧嵌入基于记忆库进行条件化,然后掩码解码器将其输入以形成预测。
我们在下面描述各个组件和训练过程,并在附录D中提供更多详细信息。
image encoder:为了实时处理任意长度的视频,我们采用流式处理方法,随着视频帧的可用性逐帧处理。图像编码器在整个交互过程中只运行一次,其作用是为每一帧提供未经过条件化的标记(特征嵌入)。我们使用了一个经过掩码自编码器(MAE,He等人,2022年)预训练的Hiera(Ryali等人,2023年;Bolya等人,2023年)图像编码器,该编码器是分层的,这使得我们在解码过程中能够使用多尺度特征。
Memory attention.记忆注意力的作用是将当前帧的特征基于过去的帧的特征和预测以及任何新的提示进行条件化。我们堆叠了L个Transformer块,第一个块以当前帧的图像编码作为输入。每个块执行自注意力操作,随后对存储在记忆库(见下文)中的(被提示/未被提示)帧的记忆和目标指针(见下文)进行交叉注意力操作,最后是一个MLP。我们使用普通的注意力操作来进行自注意力和交叉注意力,这使我们能够受益于最近在高效注意力内核方面的进展(Dao,2023)。
Prompt encoder and mask decoder.我们的提示编码器与SAM的完全相同,可以通过点击(正向或负向)、框或掩码来定义给定帧中对象的范围。稀疏提示通过位置编码与每种提示类型的学习嵌入相加来表示,而掩码则通过卷积嵌入并与帧嵌入相加。
我们的解码器设计在很大程度上遵循SAM。我们堆叠了“双向”Transformer块,用于更新提示和帧嵌入。与SAM一样,对于模糊提示(例如,单次点击),可能存在多个兼容的目标掩码,我们预测多个掩码。这种设计对于确保模型输出有效的掩码至关重要。在视频中,模糊性可能会延伸到多个视频帧,模型会在每一帧上预测多个掩码。如果后续提示没有解决这种模糊性,模型只会传播当前帧预测的具有最高IoU(交并比)的掩码。
与SAM不同,SAM在给定正提示的情况下总是存在一个有效的对象可供分割,而在PVS任务中,某些帧上可能不存在有效的对象(例如,由于遮挡)。为了支持这种新的输出模式,我们增加了一个额外的头部,用于预测当前帧上是否存在感兴趣的对象。另一个新特点是,从我们的分层图像编码器(绕过记忆注意力)添加了跳跃连接,以引入高分辨率嵌入用于掩码解码(见第D节)。
Memory encoder.记忆编码器通过使用卷积模块对输出掩码进行下采样,并将其与图像编码器提供的未经过条件化的帧嵌入(图3中未显示)逐元素相加,随后通过轻量级卷积层融合信息,从而生成记忆。(记忆是如何生成的?这里给出了答案)
Memory bank.记忆库通过维护一个包含最多N个最近帧的先进先出(FIFO)记忆队列来保留关于视频中目标对象的过往预测信息,并将来自提示的信息存储在一个最多包含M个已提示帧的先进先出队列中。例如,在视频目标分割(VOS)任务中,初始掩码是唯一的提示,记忆库会持续保留第一帧的记忆以及最多N个最近(无提示)帧的记忆。这两组记忆都以空间特征图的形式存储。
除了空间记忆之外,我们还会基于每一帧的掩码解码器输出标记,将对象指针列表存储为轻量级向量,用于表示待分割对象的高级语义信息。我们的记忆注意力会交叉关注空间记忆特征和这些对象指针。
我们将时间位置信息嵌入到最近N帧的记忆中,使模型能够表示短期目标运动,但不嵌入到被提示帧的记忆中,因为被提示帧的训练信号更稀疏,并且在推理设置中,被提示帧可能来自与训练期间看到的非常不同的时间范围,因此更难以泛化。
Trainging:该模型在图像和视频数据上联合训练。类似于之前的研究(Kirillov等人,2023年;Sofiiuk等人,2022年),我们模拟对模型进行交互式提示。我们采样包含8帧的序列,并随机选择最多2帧进行提示,并且在训练过程中使用真实掩码片段和模型预测以概率方式接收校正点击,这些点击是根据真实掩码片段和模型预测采样得到的。训练任务是顺序地(并且是“交互式地”)预测真实掩码片段。模型的初始提示可以是真实掩码(概率为0.5),从真实掩码中采样的正点击(概率为0.25),或者是边界框输入(概率为0.25)。更多详细信息见第D节。