SamPredictor

相关源文件

• notebooks/automatic_mask_generator_example.ipynb
• notebooks/onnx_model_example.ipynb
• notebooks/predictor_example.ipynb
• segment_anything/predictor.py

The SamPredictor 类提供了一个交互式界面，用于处理 Segment Anything Model (SAM)，允许用户基于各种提示高效地预测分割掩码。它会一次性计算图像嵌入，并能够使用不同的提示进行快速、重复的掩码预测，而无需重新计算图像嵌入。

有关无需提示的自动掩码生成，请参阅 Automatic Mask Generator。

概述

The SamPredictor 是 SAM 模型的一个便捷封装，专为交互式分割场景设计，用户希望通过提供点、框或之前的掩码作为提示来快速生成掩码。

目的和功能

SamPredictor 负责

• 处理输入图像以计算嵌入
• 处理不同的提示类型（点、框、掩码）
• 转换坐标并为模型准备输入
• 生成具有质量得分的分割掩码
• 支持掩码的高效迭代精炼

架构与工作流程

工作流程图

来源：segment_anything/predictor.py34-166

对象交互

来源：segment_anything/predictor.py17-270

关键组件

初始化

The SamPredictor 使用预训练的 SAM 模型进行初始化

来源：segment_anything/predictor.py17-32

图像处理

在预测掩码之前，必须使用 set_image 方法设置图像。此步骤会计算并存储图像嵌入，这在计算上成本很高，但对于每张图像只需执行一次。

该方法执行以下操作

1. 检查图像格式
2. 使用 ResizeLongestSide 转换图像
3. 将图像转换为 PyTorch 张量
4. 调用模型的图像编码器来计算嵌入

来源：segment_anything/predictor.py34-60

提示类型和掩码预测

设置图像后，您可以使用不同类型的提示来预测掩码

1. 点提示：指定前景（label=1）和背景（label=0）点
2. 框提示：提供 XYXY 格式的边界框
3. 掩码提示：提供先前预测的低分辨率掩码

The predict 方法处理这些提示并返回预测的掩码

来源：segment_anything/predictor.py92-166

输出格式

The predict 方法返回三个项目

1. masks: 二进制掩码，格式为 CxHxW（C 是掩码数量，H 和 W 是原始图像尺寸）
2. scores: 每个掩码的质量得分
3. logits: 低分辨率掩码 logits，可用作后续预测的输入

来源：segment_anything/predictor.py124-131

实现细节

图像转换

The SamPredictor 使用 ResizeLongestSide 转换来调整图像大小以达到标准尺寸，同时保持纵横比。它还提供了在原始图像空间和模型输入空间之间转换坐标和框的方法。

来源：segment_anything/predictor.py31 segment_anything/predictor.py142-148

处理 PyTorch 张量

对于高级用户，SamPredictor 提供直接与 PyTorch 张量进行操作的低级方法

• set_torch_image: 设置一个已转换的图像
• predict_torch: 从张量输入预测掩码

这些方法避免了冗余转换，对于与 PyTorch 工作流的集成非常有用。

来源：segment_anything/predictor.py63-90 segment_anything/predictor.py169-243

嵌入重用

The key efficiency advantage of SamPredictor is its ability to reuse the same image embeddings for multiple predictions. Once set_image is called, the embeddings are stored and can be reused for any number of subsequent predict calls with different prompts.

来源：segment_anything/predictor.py89 segment_anything/predictor.py230-231

使用示例

这是一个使用 SamPredictor 的简化示例

来源：notebooks/predictor_example.ipynb

参数和返回值

构造函数参数

参数	类型	描述
sam_model	Sam	要用于掩码预测的 Segment Anything Model 实例

set_image 参数

参数	类型	描述
image	np.ndarray	HWC uint8 格式的图像，像素值在 [0, 255] 范围内
image_format	str	图像的颜色格式，可以是 'RGB' 或 'BGR'

predict 参数

参数	类型	描述
point_coords	np.ndarray 或 None	(X,Y) 像素坐标的 N x 2 点提示数组
point_labels	np.ndarray 或 None	用于点的长度为 N 的标签数组（1=前景，0=背景）
box	np.ndarray 或 None	给出 XYXY 格式框提示的长度为 4 的数组
mask_input	np.ndarray 或 None	先前预测的低分辨率掩码输入
multimask_output	bool	如果为 True，则返回三个掩码而不是一个
return_logits	bool	如果为 True，则返回未阈值的掩码 logits 而不是二进制掩码

predict 返回值

返回值	类型	描述
masks	np.ndarray	CxHxW 格式的二进制掩码（C = 掩码数量，H,W = 原始图像尺寸）
scores	np.ndarray	每个掩码的质量预测，形状为 (C,)
logits	np.ndarray	CxHxW 格式的低分辨率 logits（H=W=256），可用作掩码输入

来源：segment_anything/predictor.py17-32 segment_anything/predictor.py34-60 segment_anything/predictor.py92-166

最佳实践

1. 一次设置图像：始终一次设置图像，然后使用不同的提示进行多次预测。避免对同一图像重复调用 set_image。

2. 明智地使用多掩码输出:

   • 对于模糊提示（例如单个点），请将 multimask_output=True 设置为获取多个掩码选项
   • 对于清晰的提示（例如多个点或框），请将 multimask_output=False 设置为获得更好的结果

3. 迭代精炼：使用先前生成的掩码作为后续预测的输入以精炼分割结果。

4. 质量得分：使用 multimask_output=True 时，使用返回的质量得分来选择最佳掩码。

来源：segment_anything/predictor.py115-122 notebooks/predictor_example.ipynb

与其他组件的关系

来源： segment_anything/predictor.py17-32 notebooks/automatic_mask_generator_example.ipynb

总结

SamPredictor 类提供了一个高效的接口，用于与 Segment Anything Model 进行交互式分割。它处理了图像预处理、提示编码和掩码解码的复杂细节，让用户能够专注于提供提示以获得所需的分割结果。通过仅计算一次图像嵌入并在多次预测中重复使用它们，它可以实现分割掩码的快速迭代优化。