基本介绍#

SAM2 (Segment Anything Model 2) 是 Meta AI 研究院开发的一款强大的图像分割模型，能够在各种图像上进行高效的分割任务。

• SAM2 论文
• 代码

Okay, clone 完项目的代码后，在 README.md 中可以看到 SAM2 的基本结构图。

模型结构#

README.md 中是这样介绍 SAM2 的：

Segment Anything Model 2 (SAM 2) is a foundation model towards solving promptable visual segmentation in images and videos. We extend SAM to video by considering images as a video with a single frame. The model design is a simple transformer architecture with streaming memory for real-time video processing. We build a model-in-the-loop data engine, which improves model and data via user interaction, to collect our SA-V dataset, the largest video segmentation dataset to date. SAM 2 trained on our data provides strong performance across a wide range of tasks and visual domains.

可以看出 SAM2 相较于 SAM 的主要改进在于其引入了一个流式内存（记忆）机制，使其能够更高效地处理视频数据。

配置环境#

参考 README.md 中的说明，作如下配置

conda create -n sam2 python=3.12 -y
conda activate sam2

cd sam2
pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu124
pip install -e .
pip install -e ".[notebooks]"

预训练模型下载#

官方的预训练模型有 SAM2 和 SAM2.1 两个版本，其中 SAM 2.1 是 SAM 2 的全面升级版，在训练的时候引入了额外的数据增强技术，因此 SAM 2.1 更加适合处理复杂场景（遮挡、小物体、低光照）。下面两个是官方提供的预训练模型的性能对比表格：

模型加载#

官方提供了 build_sam2 函数来加载预训练模型。和 SAM 不同，这次不是用字典做映射，而是需要提供模型的配置文件路径和权重文件路径来加载模型:

from sam2.build_sam import build_sam2

ckpt_path = "your/path/to/sam2.1_hiera_tiny.pt"
model_cfg_path = "your/path/to/sam2.1_hiera_t.yaml"

sam2_model = build_sam2(model_cfg_path, ckpt_path)

SAM2ImagePredictor 类#

import torch
from sam2.sam2_image_predictor import SAM2ImagePredictor

predictor = SAM2ImagePredictor(sam2_model)

with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    predictor.set_image(img)    # img 是一个 np.ndarray 对象
    masks, iou_predictions, low_res_masks = predictor.predict(mask_input=None, point_coords=np.array([[1280, 720]]), point_labels=np.array([1]), multimask_output=False)

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(img)
    masks, iou_predictions, low_res_masks = predictor.predict(mask_input=None, point_coords=np.array([[1280, 720]]), point_labels=np.array([1]), multimask_output=False)

build_sam2_video_predictor 函数#

build_sam2_video_predictor 函数是 SAM2 中用于构建视频分割预测器的函数。它接受一个预训练的 SAM2 模型作为输入，并返回一个 SAM2VideoPredictor 对象，该对象可以用于处理视频数据并生成分割结果:

import torch
import numpy as np
import cv2
from sam2.build_sam import build_sam2_video_predictor

# 1. 模型权重和配置路径
ckpt_path = "your/path/to/sam2.1_hiera_tiny.pt"
model_cfg_path = "your/path/to/sam2.1_hiera_t.yaml"

# 2. 构建视频预测器
predictor = build_sam2_video_predictor(model_cfg_path, ckpt_path)

# 3. 初始化推理状态（直接传入 MP4 视频路径）
video_path = r"your/video/path.mp4"
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    inference_state = predictor.init_state(
        video_path=video_path,
        offload_video_to_cpu=False,  # False: 帧加载到 GPU, True: 帧加载到 CPU 节省显存
    )

# 4. 添加 prompt（例如：在第 0 帧添加一个正样本点）
ann_frame_idx = 0  # 交互的帧索引
ann_obj_id = 1     # 物体 ID（任意唯一整数）

# 添加一个正样本点击 (x, y) = (210, 350)，label=1 表示正样本
points = np.array([[210, 350]], dtype=np.float32)
labels = np.array([1], np.int32)

_, out_obj_ids, out_mask_logits = predictor.add_new_points_or_box(
    inference_state=inference_state,
    frame_idx=ann_frame_idx,
    obj_id=ann_obj_id,
    points=points,
    labels=labels,
)

# 5. 传播到整个视频
video_segments = {}
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
    video_segments[out_frame_idx] = {
        out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
        for i, out_obj_id in enumerate(out_obj_ids)
    }

print(f"处理完成！共处理 {len(video_segments)} 帧")

# 6. 可视化并保存视频
output_video_path = "your/output/video/path.mp4"
cap = cv2.VideoCapture(video_path)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

frame_idx = 0
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    if frame_idx in video_segments:
        # 获取当前帧的所有掩码
        for obj_id, mask in video_segments[frame_idx].items():
            # mask 的形状通常是 (1, H, W) 或 (H, W)
            mask = mask.squeeze()
            # 创建一个彩色叠加层 (例如 蓝色)
            color_mask = np.zeros_like(frame, dtype=np.uint8)
            color_mask[mask] = [255, 0, 0] # BGR: Blue
            # 叠加到原图
            frame = cv2.addWeighted(frame, 1.0, color_mask, 0.5, 0)
            
            # 可选：绘制掩码轮廓
            contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(frame, contours, -1, (255, 255, 255), 2)

    out.write(frame)
    frame_idx += 1

cap.release()
out.release()
print(f"可视化视频已保存至: {output_video_path}")

从视频的结果可以看出，SAM2 的分割是追踪 prompts 所提示的物体的，因此在视频中同一物体的分割结果是连续的，且能够适应物体的形变和运动。例如，如果下一帧中这个物体消失了，可能会导致后续所有帧中都没有分割内容了。

SAM2VideoPredictor 类#

SAM2VideoPredictor 类就是前面 build_sam2_video_predictor 函数返回的对象，它提供了一个接口来处理视频数据并生成分割结果。它的使用方式与前面介绍的 SAM2ImagePredictor 类类似，但它专门针对视频数据进行了优化，能够处理视频中的连续帧，并且支持在视频中添加交互式提示（如点、框等）来指导分割过程。

很遗憾地是，官方给出的接口并不能像 SAM2ImagePredictor 那样直接自行从模型权重加载，而是需要通过 build_sam2_video_predictor 函数来构建视频预测器对象。

但是，官方提供了使用 hugginface 权重加载的方式：

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

    # add new prompts and instantly get the output on the same frame
    frame_idx, object_ids, masks = predictor.add_new_points_or_box(
            state, 
            inference_state=inference_state,
            frame_idx=ann_frame_idx,
            obj_id=ann_obj_id,
            points=points,
            labels=labels
        )

    # propagate the prompts to get masklets throughout the video
    for frame_idx, object_ids, masks in predictor.propagate_in_video(state):
        video_segments[out_frame_idx] = {
            out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
            for i, out_obj_id in enumerate(out_obj_ids)
        }

SAM2AutomaticMaskGenerator 类#

类似 SAM 的 SAMAutomaticMaskGenerator，SAM2AutomaticMaskGenerator 是 SAM2 提供的全自动实例分割工具，无需任何人工提示，即可在单张图像上自动检测并分割出所有可能的物体对象。

from sam2.build_sam import build_sam2
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator

# 加载模型
sam2_model = build_sam2("configs/sam2.1/sam2.1_hiera_l.yaml", "checkpoints/sam2.1_hiera_large.pt")

# 初始化自动掩码生成器
mask_generator = SAM2AutomaticMaskGenerator(sam2_model)

# 生成所有掩码
masks = mask_generator.generate(image)

# masks 是一个列表，每个元素包含 segmentation, bbox, area, predicted_iou 等字段

SAM2Base#

文件路径: sam2/modeling/sam2_base.py

类作用: SAM2 模型的基础类，负责图像编码、记忆融合、SAM 头预测和记忆编码的核心流程。它集成了图像主干网络、记忆注意力机制、记忆编码器以及 SAM 的提示编码器和掩码解码器，是视频目标跟踪的核心模型。

属性表格:

方法表格:

SAM2VideoPredictor#

文件路径: sam2/sam2_video_predictor.py

类作用: 继承自 SAM2Base，负责视频推理的状态管理、用户交互处理和视频传播。它管理推理状态（图像缓存、对象映射、每对象输出字典），提供添加点/框/掩码提示、传播跟踪、清除提示等功能。

属性表格:

方法表格:

SAM2ImagePredictor#

文件路径: sam2/sam2_image_predictor.py

类作用: 用于单张图像的交互式分割预测器。它缓存图像嵌入，支持点、框、掩码提示，并返回分割掩码和 IoU 分数。

属性表格:

方法表格:

PositionEmbeddingSine#

文件路径: sam2/modeling/position_encoding.py

类作用: 生成正弦位置编码，用于图像特征的位置嵌入。支持对点、框的编码，以及缓存的 2D 位置编码图。

属性表格:

方法表格:

MemoryEncoder#

文件路径: sam2/modeling/memory_encoder.py

类作用: 将视觉特征和预测掩码融合编码为记忆特征。包含掩码下采样器、特征投影、融合模块和位置编码。

属性表格:

方法表格:

MemoryAttentionLayer#

文件路径: sam2/modeling/memory_attention.py

类作用: 记忆注意力层，包含自注意力和交叉注意力（图像到记忆），以及前馈网络。用于在 MemoryAttention 中堆叠。

属性表格:

方法表格:

MemoryAttention#

文件路径: sam2/modeling/memory_attention.py

类作用: 记忆注意力模块，堆叠多个 MemoryAttentionLayer，实现自注意力与交叉注意力的多层融合。支持位置编码在输入时添加、batch_first/seq_first 格式转换，以及 RoPE 注意力中的对象指针 token 排除。

属性表格:

方法表格:

ImageEncoder#

文件路径: sam2/modeling/backbones/image_encoder.py

类作用: 图像编码器，包含主干网络（如 Hiera）、FPN 颈部（多尺度特征融合）和可选的 scalp（最高分辨率特征）。输出多尺度视觉特征和位置编码。

属性表格:

方法表格:

MaskDecoder#

文件路径: sam2/modeling/sam/mask_decoder.py

类作用: SAM 掩码解码器，基于 Transformer 和超网络生成掩码预测。支持多掩码输出、IoU 预测和对象分数预测。

属性表格:

方法表格:

PromptEncoder#

文件路径: sam2/modeling/sam/prompt_encoder.py

类作用: SAM 提示编码器，将点、框、掩码提示编码为稀疏嵌入和稠密嵌入。包含点嵌入、框嵌入、掩码下采样和位置编码。

属性表格:

方法表格:

Hiera#

文件路径: sam2/modeling/backbones/hieradet.py

类作用: 分层视觉 Transformer 主干网络，支持多尺度注意力、窗口划分和全局注意力块。用于提取图像的多尺度特征。

属性表格:

方法表格:

形状与执行路径总结#

• 图像编码: ImageEncoder.forward -> 多尺度特征 + 位置编码
  • 输入: [B, 3, H, W] -> 输出: 多尺度特征列表 [B, C, H//s, W//s]
• 记忆融合: SAM2Base._prepare_memory_conditioned_features -> MemoryAttention.forward
  • 输入: 当前特征 [HW, B, C] + 记忆 -> 输出: 融合特征 [B, C, H, W]
• 掩码预测: SAM2Base._forward_sam_heads -> PromptEncoder.forward + MaskDecoder.forward/predict_masks
  • 输入: 特征 [B, C, H, W] + 提示 -> 输出: 掩码 [B, M, H*4, W*4], IoU [B, M], 对象指针 [B, C]
• 记忆写回: SAM2Base._encode_new_memory -> MemoryEncoder.forward
  • 输入: 特征 [B, C, H, W] + 掩码 [B, 1, H*16, W*16] -> 输出: 记忆特征 [B, C, H, W], 位置编码 [1, B, C, H, W]

1. Vision Tokens (视觉特征 Token 向量)#

• 来源: ImageEncoder 输出，代表图像空间网格上的特征采样。
• 本质: 一个 $C$ 维特征向量，代表图像在该坐标点的纹理、颜色等高维信息。
• 形状: 常态下为序列 [HW, B, C]。其中每个 $1 \times C$ 的向量即为一个 Token。
• 特性: 携带 2D RoPE。在 MemoryAttention 中作为 Query 向量 参与点积运算。

2. Prompt Tokens (提示 Token 向量)#

• 来源: PromptEncoder 将用户交互位置编码至嵌入空间。
• 本质:
  • 点/框 Token: 为每个交互点生成一个 $C$ 维位置向量。
  • 掩码 Token: 掩码下采样后，每个网格单元也是一个 $C$ 维描述向量。
• 作用: 在计算相似度时，通过向量空间中的距离来吸引或排斥特定的 Vision Tokens。

3. Memory Tokens (记忆 Token 向量)#

• 来源: MemoryEncoder。
• 本质: 经过掩码加权后的 $C$ 维视觉记忆向量。
• 作用: 存储物体历史状态的向量集合。
• 位置: 在 MemoryAttention 中作为 Key/Value 向量库。

4. Object Pointers (对象指针向量)#

• 来源: MaskDecoder 的输出 Token。
• 本质: 对整帧目标的 $C$ 维语义降维向量。它将整帧物体的信息通过 Attention Pooling 压缩进几个点向量中（通常 16x256）。
• 作用: 作为目标的“向量指纹”，用于跨帧的最高层匹配。
• 特殊处理: 拼接到序列末尾参与矩阵乘法，但不带空间含义。

5. Decoder Queries (解码器查询向量)#

• 注意: 这些在严谨术语中更接近 Queries 或 Task Embeddings。
• 来源: MaskDecoder 的初始可学习 Embedding。
• 本质: 它们是进入双向 Transformer 的“初始占位符”。
• 包含:
  • Mask Queries: 并不是最终掩码，而是通过注意力机制搜集特征，最终其向量值与特征图做点积来生成掩码。
  • IoU Query: 负责汇聚整组特征，最后输入 MLP 回归出一个标量分数。
  • Object Score Query: 辅助判断物体是否在当前帧消失。