【分割大模型】SAM2（Segment Anything2）新的分割一切大模型（原理+安装+代码）

在这里插入图片描述

文章目录

前言
- 1.特点
- 2.结构
- 3.应用场景
一、原理
- 1.1 引言
- 1.2 任务：可提示的视觉分割
- 1.3 模型
- 1.4 数据引擎与 SA-V数据集
二、安装与使用

项目地址：https://github.com/facebookresearch/segment-anything-2

前言

在这里插入图片描述

Segment Anything Model 2（SAM 2）是由Meta公司发布的一个先进的图像和视频分割模型。它是Segment Anything Model（SAM）的升级版本，SAM是Meta的FAIR实验室发布的一款用于图像分割的基础模型，能够在给定提示的情况下生成高质量的对象mask。

1.特点

准确性提升：SAM 2相比原始的SAM模型在分割精度上有所提高。
速度加快：SAM 2的处理速度提高了大约六倍，可以更快地生成分割mask。
支持视频分割：除了图像分割之外，SAM 2还支持视频中的对象分割。
实时处理：非常适合于需要快速响应的应用场景，如增强现实（AR）和虚拟现实（VR）应用。
Zero-Shot泛化：良好的zero-shot迁移能力，可以在未见过的数据上工作而不需要额外的训练。
可提示的模型架构：SAM 2继承了SAM的特性，可以根据不同的提示（如点、框、甚至是文本）来生成分割结果。

2.结构

编码器-解码器架构：SAM 2很可能会继续使用编码器-解码器架构，其中编码器负责提取特征，解码器则用于生成分割掩模。
高效网络设计：为提高处理速度，SAM 2可能采用了优化过的网络结构或计算效率更高的组件。
适应性强的分割头：模型可能包含了一个高度灵活的分割头，能够根据不同的提示生成相应的掩模。
多模态输入支持：除了传统的图像输入外，SAM 2还可能支持视频帧序列作为输入，以实现视频分割。
训练数据集：SAM 2的训练数据集可能包含了大量多样化的图像和视频样本，以确保模型的泛化能力和鲁棒性。

3.应用场景

增强现实（AR）和虚拟现实（VR）：SAM 2可以用于实时分割用户周围的环境，从而增强用户体验。
自动驾驶：在自动驾驶系统中，SAM 2可以帮助车辆识别和理解道路场景中的不同元素。
医学影像分析：SAM 2可以用来自动分割医学影像中的器官或病变区域。

提示：以下是本篇文章正文内容，下面案例可供参考

一、原理

1.1 引言

图像只是真实世界的静态快照，在其中视觉片段可以表现出复杂的运动。AR/VR、机器人技术、自动驾驶汽车和视频编辑中的许多重要应用都需要在图像级分割之外的时间定位。因此，一个通用的视觉分割系统应该同时适用于图像和视频。

图像之外的挑战：由于运动、变形、遮挡、照明变化和其他因素，实体可以在外观上发生显著的变化。由于摄像机的运动、模糊度和较低的分辨率，视频的质量通常低于图像。

SAM 2配备了一个记忆注意模块，关注目标对象之前的记忆，存储关于对象和之前交互的信息，这允许它在整个视频中生成mask预测，并且根据之前观察到的帧中存储的对象记忆上下文有效地纠正这些预测。当应用于图像时，内存为空。

SAM 2使用一个数据引擎(5)，在循环中使用SAM2来交互式地注释新的和具有挑战性的数据，包括分割任何具有有效边界的对象，包括部分和子部分。最终的Segment Anything Video（SA-V）数据集包含50.9K个视频的35.5M掩码。SA-V对小物体和部分被遮挡和重新出现的挑战。

1.2 任务：可提示的视觉分割

PVS任务允许在视频的任何一帧上向模型提供提示。提示可以是正/负单击、边界框或mask ，既可以定义要分割的对象，也可以细化模型预测的对象。在接收到初始（一个或多个）提示（在同一帧或不同帧上）后，模型应该传播这些提示，以获得整个视频中的对象mask，其中包含目标对象在每个视频帧上的分割掩码。

1.3 模型

对于图像输入，该模型的行为与SAM类似。一个可提示的轻量级 mask decoder 接受当前帧上的 frame embedding 和 prompts（如果有），并输出该帧的mask。提示可以在帧上迭代添加，以细化掩码。 SAM 2解码器使用的frame embedding 并不是直接来自图像编码器，而是基于对过去预测和提示帧的记忆。相对于当前帧，提示帧也有可能来自“来自未来” 。帧的记忆是由内存编码器根据当前的预测创建的，并放置在内存库中，以便在后续的帧中使用。记忆注意操作从图像编码器中获取每一帧的特征，并将其作为memory bank的条件，以产生embedding，然后传递到mask decoder。

图像编码器 image encoder 。对于任意视频，采用 streaming 流的方法，在视频帧可用时使用它们。图像编码器对整个交互只运行一次，提供每一帧的无条件token（特征）。 具体是使用MAE预训练的Hiera（Ryali等人，2023）图像编码器。它是多层的，允许解码过程中使用多尺度特征。

记忆注意力 Memory attention 。作用是根据过去的frame 特征来condition 当前帧的特征，并新的提示来做预测。具体的，堆叠L个Transformer，第一个以当前帧的图像编码为输入。每个块先执行自注意，然后是对存储在Memory bank中的memory of (prompted/unprompted) frames 和 object pointers (见下文) 的交叉注意，然后是一个MLP。

提示编码器和掩码解码器 Prompt encoder and mask decoder。Prompt encoder与SAM的相同，可以通过点击（正或负）、边界框或掩码来提示，以定义对象在给定帧中的范围。稀疏提示由每个提示类型的位置编码与可学习embedding求和表示，而mask使用卷积并与 frame embedding相加。

在这里插入图片描述

堆叠“two-way” transformer blocks 来更新 prompt 和 frame embeddings。对于可能存在多个mask的模糊提示（即单次点击），预测多个结果。在视频中，模糊性可以扩展到各个视频帧，该模型预测了每一帧上的多个掩模。如果没有后续提示解决歧义，该模型只传播当前帧预测最高IoU的mask。

PVS任务中，某些帧上可能不存在有效的对象（例如遮挡）为此添加了一个额外的head，来预测感兴趣的对象是否存在于当前帧上。与SAM的另一个不同之处是，使用来自image encoder 的 skip connections（绕过 memory attention）来合并高分辨率信息以进行 mask decoding。

内存编码器 Memory encoder.使用卷积对输出mask降采样，并与图像编码器的unconditioned frame embedding逐元素求和来生成记忆。

Memory bank。维护多达N个最近帧的FIFO记忆队列来保留视频中目标对象过去的预测信息 ，将提示信息存储在 多达M个提示帧 的FIFO队列中。例如，在VOS任务中，初始mask是唯一的提示，Memory bank始终保留第一帧的记忆以及多达N个最近（unprompted）帧的记忆。这两组记忆都被存储为空间特征图。

除了空间内存之外，还基于每一帧的mask decoder的输出token，将一个object pointers 的轻量级向量，用于分割对象的高级语义信息。我们的记忆attention交叉关注空间记忆特征和这些object pointers。

我们嵌入temporal的位置信息到N个最近帧的记忆中，来允许模型表示短期的对象运动，而但不是提示帧。（因为来自提示帧的训练信号更稀疏，并且更难以推广到推理设置，其中提示帧可能来自一个与在训练中看到的非常不同的时间范围）

训练。图像和视频联合训练。follow（Kirillov等人，2023年；Sofiiuk等人，2022年），模拟了模型的交互式提示。我们对8帧的序列进行采样，并随机选择最多2帧来提示和概率接收纠正点击，并在训练期间使用真实mask和模型预测进行采样。训练任务是顺序（和“交互”）预测真实mask。模型的初始提示可以是概率为0.5的地面真掩模，从概率为0.25的地面真掩模中采样的正点击，或者概率为0.25的边界框输入。

1.4 数据引擎与 SA-V数据集

第一阶段：逐帧SAM。使用SAM以每帧6秒（FPS）的速度注释视频的每帧中目标对象的掩码，此外，还使用“brush”等像素精确的手动编辑工具。每帧的平均注释时间为37.8秒，但为每帧产生了高质量的空间注释。此阶段1.4K的视频中收集了16K个mask。阶段2： SAM + SAM 2 mask。第二个阶段将SAM 2添加到循环中，其只接受掩码作为提示。将第一帧的空间掩模，然后使用SAM 2将带注释的mask在时间上传播到其他帧，以获得完整的时空掩模。在任何后续的视频帧中，注释者都可以通过使用SAM、“笔刷”和/或“brush”从头开始注释SAM 2mask所做的预测，并使用SAM 重新传播，重复这个过程，直到mask正确。该阶段收集了63.5K的msk。注释时间下降到7.4秒/帧。阶段3: SAM 2。 此阶段SAM 2接受各种类型的提示，包括点和mask。其受益于跨时间维度的物体记忆来生成mask预测。每帧的注释时间下降到4.5秒，收集了19个7.0K的掩模。

SA-V数据集包含50.9K的视频和64个2.6K的mask，见下图：

在这里插入图片描述

二、安装与使用

1.先安装 SAM 2。代码需要python>=3.10，以及torch>=2.3.1和。请按照此处的torchvision>=0.18.1说明安装 PyTorch 和 TorchVision 依赖项。您可以使用以下方式在 GPU 机器上安装 SAM 2：

git clone https://github.com/facebookresearch/segment-anything-2.git

cd segment-anything-2; pip install -e .

2.下载离线模型（在项目主页下载，推荐large版本）：

在这里插入图片描述

3.图像预测（需要输入对应的img/video路径，以及prompt）：

import torch
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

checkpoint = "./checkpoints/sam2_hiera_large.pt"
model_cfg = "sam2_hiera_l.yaml"
predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))

with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    predictor.set_image(<your_image>)
    masks, _, _ = predictor.predict(<input_prompts>)

视频预测：（需要输入对应的img/video路径，以及prompt）

import torch
from sam2.build_sam import build_sam2_video_predictor

checkpoint = "./checkpoints/sam2_hiera_large.pt"
model_cfg = "sam2_hiera_l.yaml"
predictor = build_sam2_video_predictor(model_cfg, checkpoint)

with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    state = predictor.init_state(<your_video>)

    # add new prompts and instantly get the output on the same frame
    frame_idx, object_ids, masks = predictor.add_new_points(state, <your_prompts>):

    # propagate the prompts to get masklets throughout the video
    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
        ...

4.还可以使用ultralytics包的封装来调用（选用）

安装必要的包。安装ultralytics并确保其版本>=8.2.70，torch版本也需>=2.0或使用最新版本

pip install -U ultralytics

下载分割模型。下面网址中提供了4个模型，大家可以根据自己需要下载

https://docs.ultralytics.com/models/sam-2/#how-can-i-use-sam-2-for-real-time-video-segmentation

全局目标分割：

from ultralytics import ASSETS, SAM

# Load a model
model = SAM("sam2_s.pt")

# Display model information (optional)
model.info()

# Segment image or video
results = model('d.jpg') # 图片推理
#results = model('d.jpg') # 视频推理

# Display results
for result in results:
    result.show()