SAM 2：在图像和视频中进行任意分割

Meta 的人工智能，FAIR

Nikhila Ravi、Valentin Gabeur、Yuan-Ting Hu、Ronghang Hu、Chaitanya Ryali、Tengyu Ma、Haitham Khedr、Roman Rädle、Chloe Rolland、Laura Gustafson、Eric Mintun、Junting Pan、Kalyan Vasudev Alwala、Nicolas Carion、Chao-Yuan Wu、Ross Girshick、Piotr Dollár、Christoph Feichtenhofer

[论文] [项目] [演示] [数据集] [博客] [BibTeX]

Segment Anything Model 2 (SAM 2) 是一种用于解决图像和视频中可提示性视觉分割的基础模型。我们通过将图像视为仅包含一帧的视频，将 SAM 扩展到视频领域。该模型采用简单的 Transformer 架构，并配备流式内存以实现实时视频处理。我们构建了一个模型闭环数据引擎，通过用户交互不断改进模型和数据，从而收集了 我们的 SA-V 数据集，这是迄今为止最大的视频分割数据集。基于我们数据训练的 SAM 2 在广泛的任务和视觉领域中均表现出色。

最新更新

2024年12月11日——完整模型编译以大幅提升 VOS 性能，并推出新的 SAM2VideoPredictor 以更好地处理多目标跟踪

• 我们现在支持对整个 SAM 2 模型在视频上的 torch.compile，只需在 build_sam2_video_predictor 中设置 vos_optimized=True 即可启用，这将显著提升 VOS 推理速度。
• 我们更新了 SAM2VideoPredictor 的实现，支持每个目标独立推理，从而放宽了多目标跟踪必须通过提示输入的假设，并允许在跟踪开始后添加新对象。
• 完整详情请参阅 RELEASE_NOTES.md。

2024年9月30日——SAM 2.1 开发者套件（全新检查点、训练代码、网页演示）发布

• 发布了一套改进后的模型检查点（标记为 SAM 2.1）。详细信息请参阅 模型描述。
  • 要使用新的 SAM 2.1 检查点，您需要从本仓库获取最新版本的模型代码。如果您已安装过本仓库的早期版本，请先通过 pip uninstall SAM-2 卸载旧版本，然后使用 git pull 获取最新代码，并按照下方的 安装说明 重新安装本仓库。
• 训练（及微调）代码现已发布。请参阅 training/README.md，了解如何开始。
• SAM 2 网页演示的前端 + 后端代码也已发布。详细信息请参阅 demo/README.md。

安装

在使用 SAM 2 之前，需先完成安装。代码要求 python>=3.10，以及 torch>=2.5.1 和 torchvision>=0.20.1。请按照 此处 的说明安装 PyTorch 和 TorchVision 依赖项。您可以在 GPU 机器上通过以下命令安装 SAM 2：

git clone https://github.com/facebookresearch/sam2.git && cd sam2

pip install -e .

如果您在 Windows 上安装，强烈建议使用 Windows Subsystem for Linux (WSL)，并搭配 Ubuntu 系统。

要使用 SAM 2 预测器并运行示例笔记本，还需要 jupyter 和 matplotlib，可通过以下命令安装：

pip install -e ".[notebooks]"

注意：

1. 建议使用 Anaconda 创建一个新的 Python 环境进行安装，并按照 https://pytorch.org/ 使用 pip 安装 PyTorch 2.5.1 或更高版本。如果您的当前环境中 PyTorch 版本低于 2.5.1，上述安装命令将尝试通过 pip 将其升级到最新版本。
2. 上述步骤需要使用 nvcc 编译器编译自定义 CUDA 内核。如果您的机器上尚未安装 CUDA 工具包，请根据您的 PyTorch CUDA 版本下载相应的 CUDA 工具包。
3. 如果在安装过程中出现类似 Failed to build the SAM 2 CUDA extension 的提示，您可以忽略它，仍然可以正常使用 SAM 2（某些后处理功能可能会受限，但通常不会影响结果）。

有关潜在问题及解决方案的常见问题解答，请参阅 INSTALL.md。

快速入门

下载检查点

首先，我们需要下载一个模型检查点。所有模型检查点都可以通过以下命令下载：

cd checkpoints && \
./download_ckpts.sh && \
cd ..

或者单独从以下链接下载：

• sam2.1_hiera_tiny.pt
• sam2.1_hiera_small.pt
• sam2.1_hiera_base_plus.pt
• sam2.1_hiera_large.pt

（请注意，这些是标记为 SAM 2.1 的改进版检查点；详细信息请参阅 模型描述。）

随后，只需几行代码即可使用 SAM 2 进行图像和视频预测。

图像预测

SAM 2 具备 SAM 在静态图像上的所有能力，并且我们提供了与 SAM 非常相似的图像预测 API，以满足图像使用场景的需求。SAM2ImagePredictor 类提供了一个简单的接口来处理图像提示。

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

checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_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>)

有关静态图像使用场景的示例，请参阅 image_predictor_example.ipynb（也可在 Colab 中查看 这里）。

SAM 2 还支持像 SAM 一样在图像上自动生成掩码。有关图像中自动掩码生成的详细信息，请参阅 automatic_mask_generator_example.ipynb（也可在 Colab 中查看 这里）。

视频预测

对于视频中的可提示分割和跟踪，我们提供了一个视频预测器，其 API 允许用户添加提示并在整个视频中传播小掩码。SAM 2 支持多对象视频推理，并使用推理状态来跟踪每个视频中的交互过程。

import torch
from sam2.build_sam import build_sam2_video_predictor

checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_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>)

    # 添加新提示并立即在同一帧上获得输出
    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):

    # 传播提示以在整个视频中获取小掩码
    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
        ...

有关如何在视频中添加点击或框选提示、进行细化以及跟踪多个对象的详细信息，请参阅 video_predictor_example.ipynb（也可在 Colab 中查看 这里）。

从 🤗 Hugging Face 加载

此外，模型也可以从 Hugging Face 加载（需要安装 pip install huggingface_hub）。

对于图像预测：

import torch
from sam2.sam2_image_predictor import SAM2ImagePredictor

predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")

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

对于视频预测：

import torch
from sam2.sam2_video_predictor import SAM2VideoPredictor

predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-large")

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

    # 添加新提示并立即在同一帧上获得输出
    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>):

    # 传播提示以在整个视频中获取小掩码
    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
        ...

模型说明

SAM 2.1 检查点

下表展示了于 2024 年 9 月 29 日发布的改进版 SAM 2.1 检查点。

模型	参数量 (M)	速度 (FPS)	SA-V 测试 (J&F)	MOSE 验证集 (J&F)	LVOS v2 (J&F)
sam2.1_hiera_tiny (配置, 检查点)	38.9	91.2	76.5	71.8	77.3
sam2.1_hiera_small (配置, 检查点)	46	84.8	76.6	73.5	78.3
sam2.1_hiera_base_plus (配置, 检查点)	80.8	64.1	78.2	73.7	78.2
sam2.1_hiera_large (配置, 检查点)	224.4	39.5	79.5	74.6	80.6

SAM 2 检查点

此前于 2024 年 7月 29 日发布的 SAM 2 检查点如下：

模型	参数量 (M)	速度 (FPS)	SA-V 测试 (J&F)	MOSE 验证集 (J&F)	LVOS v2 (J&F)
sam2_hiera_tiny (配置文件, 检查点)	38.9	91.5	75.0	70.9	75.3
sam2_hiera_small (配置文件, 检查点)	46	85.6	74.9	71.5	76.4
sam2_hiera_base_plus (配置文件, 检查点)	80.8	64.8	74.7	72.8	75.8
sam2_hiera_large (配置文件, 检查点)	224.4	39.7	76.0	74.6	79.8

速度是在配备 torch 2.5.1, cuda 12.4 的 A100 上测量的。有关基准测试的示例（编译所有模型组件），请参阅 benchmark.py。仅编译图像编码器可能更加灵活，并且也能带来（较小的）速度提升（在配置中设置 compile_image_encoder: True）。

Segment Anything Video 数据集

详情请参阅 sav_dataset/README.md。

训练 SAM 2

您可以在自定义的图像、视频或两者兼有的数据集上训练或微调 SAM 2。请查看训练 README，了解如何开始。

SAM 2 的 Web 演示

我们已发布了 SAM 2 Web 演示的前端和后端代码（一个可本地部署的版本，类似于 https://sam2.metademolab.com/demo）。详细信息请参阅 Web 演示 README。

许可证

SAM 2 模型检查点、SAM 2 演示代码（前端和后端）以及 SAM 2 训练代码均采用 Apache 2.0 许可证，然而 SAM 2 演示代码中使用的 Inter 字体 和 Noto Color Emoji 则根据 SIL 开源字体许可证，第 1.1 版 提供。

贡献

请参阅 贡献说明 和 行为准则。

贡献者

SAM 2 项目得以实现，离不开众多贡献者的帮助（按字母顺序排列）：

Karen Bergan, Daniel Bolya, Alex Bosenberg, Kai Brown, Vispi Cassod, Christopher Chedeau, Ida Cheng, Luc Dahlin, Shoubhik Debnath, Rene Martinez Doehner, Grant Gardner, Sahir Gomez, Rishi Godugu, Baishan Guo, Caleb Ho, Andrew Huang, Somya Jain, Bob Kamma, Amanda Kallet, Jake Kinney, Alexander Kirillov, Shiva Koduvayur, Devansh Kukreja, Robert Kuo, Aohan Lin, Parth Malani, Jitendra Malik, Mallika Malhotra, Miguel Martin, Alexander Miller, Sasha Mitts, William Ngan, George Orlin, Joelle Pineau, Kate Saenko, Rodrick Shepard, Azita Shokrpour, David Soofian, Jonathan Torres, Jenny Truong, Sagar Vaze, Meng Wang, Claudette Ward, Pengchuan Zhang。

第三方代码：我们使用基于 GPU 的连通分量算法，该算法改编自 cc_torch（其许可证见 LICENSE_cctorch），作为掩码预测的可选后处理步骤。

引用 SAM 2

如果您在研究中使用了 SAM 2 或 SA-V 数据集，请使用以下 BibTeX 条目。

@article{ravi2024sam2,
  title={SAM 2：图像与视频中的任意对象分割},
  author={Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and Hu, Ronghang and Ryali, Chaitanya and Ma, Tengyu and Khedr, Haitham and R{\"a}dle, Roman and Rolland, Chloe and Gustafson, Laura and Mintun, Eric and Pan, Junting and Alwala, Kalyan Vasudev and Carion, Nicolas and Wu, Chao-Yuan and Girshick, Ross and Doll{\'a}r, Piotr and Feichtenhofer, Christoph},
  journal={arXiv 预印本 arXiv:2408.00714},
  url={https://arxiv.org/abs/2408.00714},
  year={2024}
}

SAM 2 快速上手指南

SAM 2 (Segment Anything Model 2) 是 Meta FAIR 推出的新一代基础模型，支持对图像和视频进行提示式分割。相比前代，SAM 2 引入了流式记忆机制，能够实时处理视频并跟踪多个对象。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: 推荐 Linux (Ubuntu)。Windows 用户强烈建议使用 WSL2 (Windows Subsystem for Linux)。
• Python: 版本 >= 3.10
• PyTorch: 版本 >= 2.5.1
• TorchVision: 版本 >= 0.20.1
• GPU: 需要 NVIDIA GPU 及对应的 CUDA 工具包（用于编译自定义 CUDA 内核）。

注意：建议先访问 PyTorch 官网 安装匹配您 CUDA 版本的 PyTorch 和 TorchVision。国内用户可使用清华或中科大镜像源加速安装。

安装步骤

1. 克隆代码库

git clone https://github.com/facebookresearch/sam2.git && cd sam2

2. 安装依赖

使用 pip 以可编辑模式安装 SAM 2：

pip install -e .

如果需要运行示例 Notebook，请额外安装：

pip install -e ".[notebooks]"

提示：安装过程中若出现 Failed to build the SAM 2 CUDA extension 警告，通常可以忽略，不影响核心功能使用。

3. 下载模型权重

进入检查点目录并运行下载脚本，获取最新的 SAM 2.1 改进版模型：

cd checkpoints && ./download_ckpts.sh && cd ..

或者手动下载单个模型（推荐 sam2.1_hiera_large.pt 以获得最佳效果）：

• sam2.1_hiera_tiny.pt
• sam2.1_hiera_small.pt
• sam2.1_hiera_base_plus.pt
• sam2.1_hiera_large.pt

基本使用

场景一：图像分割预测

SAM 2 完全兼容静态图像分割，接口与 SAM v1 类似。

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

# 配置模型路径
checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_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>)

场景二：视频分割与跟踪

SAM 2 的核心优势在于视频处理，支持添加提示点/框并在整个视频中传播掩码。

import torch
from sam2.build_sam import build_sam2_video_predictor

# 配置模型路径
checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_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>)

    # 添加新提示（点或框），立即获取当前帧结果
    frame_idx, object_ids, masks = predictor.add_new_points_or_box(state, <your_prompts>)

    # 在视频中传播提示，获取全程掩码
    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
        # 在此处处理每一帧的 masks
        ...

替代方案：从 Hugging Face 加载

如果您无法直接下载权重文件，可以使用 Hugging Face 自动加载（需先安装 huggingface_hub）：

from sam2.sam2_image_predictor import SAM2ImagePredictor

# 图像预测
predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")

# 视频预测
from sam2.sam2_video_predictor import SAM2VideoPredictor
predictor = SAM2VideoPredictor.from_pretrained("facebook/sam2-hiera-large")