SAM2（Segment Anything Model 2）新一代分割一切大模型实战总结

AI浩

发布于 2024-10-22 13:18:06

3.7K00

代码可运行

文章被收录于专栏：AI智韵AI智韵

运行总次数：0

代码可运行

Segment Anything Model 2（SAM 2）作为Meta公司发布的Segment Anything Model（SAM）的升级版本，在图像和视频分割领域展现出了显著的优点和特性。

论文连接：https://arxiv.org/pdf/2408.00714

Demo: https://sam2.metademolab.com Code: https://github.com/facebookresearch/segment-anything-2 Website: https://ai.meta.com/sam2

优点

分割精度提升：
- 相比原始的SAM模型，SAM 2在分割精度上有所提高，能够更准确地识别并分割图像和视频中的对象。
处理速度加快：
- SAM 2的处理速度提高了大约六倍，能够更快地生成分割掩模，适用于需要快速响应的应用场景。
支持视频分割：
- 除了图像分割之外，SAM 2还支持视频中的对象分割，为视频处理和分析提供了强大的工具。
实时处理能力：
- SAM 2非常适合于需要实时处理的应用场景，如增强现实（AR）和虚拟现实（VR）应用，能够实时分割用户周围的环境。
Zero-Shot泛化能力：
- SAM 2具有良好的zero-shot迁移能力，可以在未见过的数据上工作而不需要额外的训练，这使得它能够在多种不同的视觉域中应用。
广泛的训练数据：
- SAM 2的训练数据集包含了1100万张图像和110亿个掩码，这种广泛的训练数据使得SAM 2成为训练新图像分割任务的强大起点。

特性

可提示的模型架构：
- SAM 2可以根据不同的提示（如点、框、甚至是文本）来生成分割结果，这种灵活性使得用户可以根据需要轻松定义目标对象。
编码器-解码器架构：
- SAM 2很可能继续使用了编码器-解码器架构，其中编码器负责提取特征，解码器则用于生成分割掩模。
高效网络设计：
- 为提高处理速度，SAM 2可能采用了优化过的网络结构或计算效率更高的组件。
适应性强的分割头：
- 模型包含了一个高度灵活的分割头，能够根据不同的提示生成相应的掩模，从而适应不同的分割任务。
多模态输入支持：
- 除了传统的图像输入外，SAM 2还支持视频帧序列作为输入，以实现视频分割。
记忆注意模块：
- SAM 2配备了一个记忆注意模块，关注目标对象之前的记忆，存储关于对象和之前交互的信息，这允许它在整个视频中生成mask预测，并根据之前观察到的帧中存储的对象记忆上下文有效地纠正这些预测。
PVS任务支持：
- 在视频的任何一帧上向模型提供提示（如正/负单击、边界框或mask），模型可以传播这些提示以获得整个视频中的对象mask。
开源与社区支持：
- Meta公司以宽松的Apache 2.0许可证共享了SAM 2的代码和模型权重，促进了社区的研究和应用。

SAM 2以其高精度、高速度、广泛的支持能力以及强大的zero-shot泛化能力，在图像和视频分割领域展现出了显著的优势和广泛的应用前景。

实战

实现SAM推理有两种方法，一种是直接使用官方的SAM2模型，另一种使用Ultralytics 。

基于官方模型的SAM2实战

GitHub链接：

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

使用前需要先安装 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 .

然后，下载模型： sam2_hiera_tiny.pt：

https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt

sam2_hiera_small.pt：

https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt

sam2_hiera_base_plus.pt：

https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt

sam2_hiera_large.pt：

https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt

测试图片

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>)

测试视频

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_or_box(state, <your_prompts>):

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

方法二：基于ultralytics包的封装来调用

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

pip install -U ultralytics

下载分割模型。链接如下：

SAM 2 tiny链接：

https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_t.pt

SAM 2 small：

https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_s.pt

SAM 2 base

https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_b.pt

SAM 2 large

https://github.com/ultralytics/assets/releases/download/v8.2.0/sam2_l.pt

全局分割

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('car.jpg') # 图片推理
# Display results
for result in results:    
 result.show()

指定点或者矩形区域推理图片

from ultralytics import SAM

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

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

# Segment with bounding box prompt
results = model("path/to/image.jpg", bboxes=[100, 100, 200, 200])

# Segment with point prompt
results = model("path/to/image.jpg", points=[150, 150], labels=[1])

for result in results:    
 result.show()

视频推理

from ultralytics import SAM

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

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

# Run inference
model("path/to/video.mp4")

SAM 2 与 YOLOv8 的比较

在这里，我们将 Meta 的最小 SAM 2 模型 SAM2-t 与 Ultralytics 的最小分割模型 YOLOv8n-seg 进行比较：

这种比较显示了模型大小和速度之间的数量级差异。虽然 SAM 具有自动分割的独特能力，但它并不是 YOLOv8 分割模型的直接竞争对手，因为 YOLOv8 模型的体积更小、速度更快、效率更高。

自动标注：高效的数据集创建

自动标注是SAM 2的一个强大功能，它使用户能够利用预训练模型快速且准确地生成分割数据集。这一功能在无需大量手动工作的情况下创建大型、高质量数据集时尤其有用。

使用SAM 2自动标注您的数据集，代码如下：

from ultralytics.data.annotator import auto_annotate

auto_annotate(data="path/to/images", det_model="yolov8x.pt", sam_model="sam2_b.pt")

参数	类型	描述	默认值
data	str	包含待标注图像的文件夹的路径。
det_model	str, optional	预训练的YOLO检测模型。默认为'yolov8x.pt'。	'yolov8x.pt'
sam_model	str, optional	预训练的SAM 2分割模型。默认为'sam2_b.pt'。	'sam2_b.pt'
device	str, optional	运行模型的设备。默认为空字符串（如果有的话，则为CPU或GPU）。
output_dir	str, None, optional	保存标注结果的目录。默认为与'data'同目录下的'labels'文件夹。	None