SAM 2：图像与视频中的任意分割

Meta AI, 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

[论文] [项目主页] [在线演示] [数据集] [博客] [引用]

Segment Anything Model 2 (SAM 2) 是一个基础模型，旨在解决图像和视频中基于提示的可视分割任务。我们将 SAM 扩展到视频，通过将图像视为单帧视频来实现。模型设计采用简单的 Transformer 架构，并具备流式记忆能力，支持实时视频处理。我们构建了模型参与的数据引擎，通过用户交互不断优化模型与数据，收集了迄今为止最大的视频分割数据集 SA-V 数据集。基于此数据训练的 SAM 2 在多种任务和视觉领域均展现出强大性能。

最新更新

2024年11月12日 —— 完整模型编译带来显著 VOS 加速，以及新的 SAM2VideoPredictor 以更好地处理多目标跟踪

• 我们现在支持对整个 SAM 2 模型在视频上进行 torch.compile，可通过在 build_sam2_video_predictor 中设置 vos_optimized=True 来启用，从而显著加速 VOS 推理。
• 我们更新了 SAM2VideoPredictor 的实现，支持独立的逐目标推理，放宽了多目标跟踪时提示的假设，并允许在跟踪开始后添加新目标。
• 详见 RELEASE_NOTES.md。

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

• 发布了一套新的、改进后的模型检查点（标记为 SAM 2.1）。详细信息请参见模型描述。
• 要使用新的 SAM 2.1 检查点，您需要本仓库最新的模型代码。如果您已安装早期版本，请先通过 pip uninstall SAM-2 卸载，然后通过 git pull 拉取最新代码，并根据下方的安装说明重新安装。
• 训练（及微调）代码已发布。请参阅 training/README.md 了解如何开始。
• SAM 2 Web 演示的前端和后端代码已发布。详见 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 上安装，强烈建议使用 适用于 Linux 的 Windows 子系统（WSL） 并运行 Ubuntu。

为了使用 SAM 2 预测器并运行示例 notebooks，需要安装 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 内核。如果您的机器上尚未安装，请安装与您 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 检查点如下所示：

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

Segment Anything 视频数据集

详情请参见 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 Font 和 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: Segment Anything in Images and Videos},
  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 preprint arXiv:2408.00714},
  url={https://arxiv.org/abs/2408.00714},
  year={2024}
}