用 Segment Anything（SAM）在 Python 中做单点提示图像分割

Segment Anything（SAM）是 Meta AI 提出的图像分割模型，它的特点是把「分割什么」交给提示来决定：给一个点、一个框，或者若干正负样本点，模型就返回对应的物体掩码。它适合需要从图片里快速抠出任意物体的场景，比如做标注辅助、图像编辑前的选区生成、科学图像里的目标提取，也适合想在自己的 Python 流程里嵌入一个通用分割能力的开发者。这篇教程讲的是本地 Python 环境下的用法，重点是单点提示分割这条最短路径，同时把模型类型、权重文件、预测参数和返回值这些容易卡住的地方讲清楚。如果你更想先在线试一下效果，可以打开 Segment Anything 工具页面对照理解提示式分割的交互方式，再回到本地写代码。

准备工作

Python 与深度学习框架

运行 SAM 的代码要求 Python 3.8 及以上版本，并且需要 PyTorch 1.7 及以上、torchvision 0.8 及以上。如果机器上有 NVIDIA 显卡并装好了对应的驱动，建议安装支持 CUDA 的 PyTorch 与 torchvision，这样推理会明显快一些；只有 CPU 也能跑，但速度会慢很多，尤其是较大的模型。

PyTorch 的安装方式与 CUDA 版本强相关，不同系统、不同显卡驱动对应的安装命令不一样，具体命令请以 PyTorch 官网当前的安装指引为准。装好之后可以在 Python 里确认一下版本：

import torch, torchvision
print(torch.__version__)
print(torchvision.__version__)
print(torch.cuda.is_available())

最后一行返回 True 说明当前环境能用到 GPU，后面代码里的 device 就可以写 "cuda"；返回 False 就写 "cpu"。

安装 segment-anything 包

SAM 的代码以 Python 包的形式发布，安装方式有两种。第一种是直接从代码仓库安装：

pip install git+https://github.com/facebookresearch/segment-anything.git

这种方式需要机器上已经装好 git，并且能访问代码托管站点。第二种是先把源码压缩包下载到本地，再在命令行里用 pip 安装这个 zip 文件，适合网络受限或者想固定某个版本的情况：

cd /d G:\downloads
pip install segment-anything-main.zip

路径按你实际存放压缩包的位置改。安装完成后，能在 Python 里成功 import 就说明装好了：

from segment_anything import sam_model_registry, SamPredictor
print("ok")

安装其余依赖

除了 torch 之外，这个项目还依赖 opencv-python、pycocotools、matplotlib、onnxruntime 等包。opencv 用来读写和缩放图片，matplotlib 用来把掩码画出来看，pycocotools 与 onnxruntime 在部分流程和导出场景中会用到。可以一次性装上：

pip install opencv-python pycocotools matplotlib onnxruntime onnx torch

如果 torch 已经按上一步装好了，这里重复写一次通常不会有问题，pip 会识别已安装的版本。

准备模型权重与模型类型

SAM 的权重文件需要单独下载，文件名与模型类型一一对应。代码里通过 model_type 指定用哪一个，通过 checkpoint 指定权重文件路径，两者必须匹配，否则加载时会报错。常见的对应关系是：

model_type	权重文件名	适用场景
vit_b	sam_vit_b_01ec64.pth	体积最小、速度最快，适合先跑通流程
vit_l	sam_vit_l_0b3195.pth	精度与速度折中
vit_h	sam_vit_h_4b8939.pth	精度最高，显存与耗时也最大

权重文件的具体下载地址、文件大小与许可证信息请以项目官方页面当前的信息为准。把下载好的 .pth 文件放在一个固定目录里，代码中用绝对路径引用最稳妥，可以避免工作目录变化导致的找不到文件。

准备一张测试图片

源码包的 notebooks/images 目录下自带示例图片，例如 truck.jpg，可以直接拿来测试。也可以换成自己的图片，但要注意后面输入点的坐标是相对于「送入模型的那张图」的像素坐标，如果对图片做了缩放，坐标也要跟着换算。

操作步骤

第一步：读入图片并转成 RGB

OpenCV 读进来的图片默认是 BGR 通道顺序，而 SAM 与 matplotlib 都按 RGB 处理，所以必须转换。如果图片太大，可以先缩放，减少显存占用和推理时间：

import cv2

image = cv2.imread(r"G:/data/images/truck.jpg")
image = cv2.resize(image, None, fx=0.5, fy=0.5)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
print(image.shape)

打印出的 shape 是 (高, 宽, 3)。记住这个高和宽，后面选提示点坐标时不能超出这个范围。

第二步：加载模型并搬到设备上

用 sam_model_registry 按 model_type 取出对应的模型构造函数，传入 checkpoint 完成权重加载，再用 .to(device) 把模型放到 GPU 或 CPU 上：

import torch
from segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "G:/models/sam_vit_b_01ec64.pth"
model_type = "vit_b"
device = "cuda" if torch.cuda.is_available() else "cpu"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)

这里把 device 写成自动判断，比硬编码 "cuda" 更省事：没有 GPU 的机器上不会因为找不到 CUDA 而直接崩掉。

第三步：把图片交给 predictor

SamPredictor 会先对整张图做一次图像编码，这一步是后续所有提示预测的公共开销。同一张图上做多次不同提示的预测时，只需要 set_image 一次：

predictor.set_image(image)

set_image 之后，predictor 内部保存了这张图的特征。换图时必须重新调用，否则预测结果会对应到上一张图。

第四步：构造提示点与标签

提示点用 numpy 数组表示，形状是 (N, 2)，每个元素是 (x, y) 像素坐标。标签数组形状是 (N,)，取值 1 表示这个点是前景（要分割的物体），0 表示背景（要排除的区域）。单点分割就是只给一个前景点：

import numpy as np

input_point = np.array([[250, 187]])
input_label = np.array([1])

坐标的选择很关键：点落在目标物体内部，返回的掩码就是这个物体；点落在背景上，返回的就是背景区域。如果第一次结果不理想，可以增加负样本点，把不想要的部分标成 0，模型会据此调整分割范围。

第五步：调用 predict 得到掩码

把提示点、标签传进 predict，multimask_output=True 表示让模型输出多个候选掩码，通常是一个整体、一个局部、一个更大范围，由调用方按分数挑选：

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)
print(masks.shape)
print(scores)

返回值的含义：

• masks：形状为 (候选数, 高, 宽) 的布尔数组，每个候选掩码对应一个二维布尔图，True 表示该像素属于分割结果。
• scores：每个候选掩码的质量分数，数值越高表示模型越确信。单点提示下通常取分数最高的那个。
• logits：掩码的低分辨率对数几率输出，需要做进一步处理才能当掩码用，一般用于需要精细调整阈值的场景。

如果只需要一个结果，把 multimask_output 设为 False，此时 masks 的第一维就是 1，模型直接给出它认为最合适的单个掩码。

第六步：把掩码叠加到原图上查看

掩码本身是布尔数组，直接看没有意义，需要乘上一个带透明度的颜色再叠加到原图上。下面两个辅助函数分别负责画掩码和画提示点：

import matplotlib.pyplot as plt

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)

然后逐个候选掩码画出来，标题里带上分数，方便对比：

for i, (mask, score) in enumerate(zip(masks, scores)):
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_mask(mask, plt.gca())
    show_points(input_point, input_label, plt.gca())
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
    plt.axis('off')
    plt.show()

绿色星号是前景提示点，红色星号是背景提示点。看图上掩码是否贴合目标边缘，就能判断这次提示选得好不好。

一个完整示例

下面把前面的步骤串成一个可以直接运行的脚本。只需要改三处：图片路径、权重路径、提示点坐标。其余部分不用动。

import sys
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
from segment_anything import sam_model_registry, SamPredictor

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*',
               s=marker_size, edgecolor='white', linewidth=1.25)

# 1. 读图并转 RGB
image = cv2.imread(r"G:/data/images/truck.jpg")
image = cv2.resize(image, None, fx=0.5, fy=0.5)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 2. 加载模型
sam_checkpoint = "G:/models/sam_vit_b_01ec64.pth"
model_type = "vit_b"
device = "cuda" if torch.cuda.is_available() else "cpu"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)

# 3. 图像编码
predictor.set_image(image)

# 4. 提示点
input_point = np.array([[250, 187]])
input_label = np.array([1])

# 5. 预测
masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)
print(masks.shape)

# 6. 可视化
for i, (mask, score) in enumerate(zip(masks, scores)):
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    show_mask(mask, plt.gca())
    show_points(input_point, input_label, plt.gca())
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
    plt.axis('off')
    plt.show()

运行后应该看到三张叠加了掩码的图，每张标题里有一个分数。如果三张图里有一张准确圈住了目标，说明流程已经跑通；如果三张都偏了，先检查提示点坐标是不是落在目标上，再检查图片是否做过缩放导致坐标错位。

脚本里的 sys.path.append("..") 只在源码目录结构下运行 notebook 时才需要，用 pip 安装成包之后可以删掉。

注意事项

模型类型与权重必须匹配

model_type 写 vit_b 却加载 vit_h 的权重，会在加载阶段直接报错。三个权重文件不能混用，下载时看清楚文件名里的 vit_b、vit_l、vit_h 标记。

坐标基于送入模型的图片

提示点坐标是相对于 set_image 时传入的那张图的像素坐标。如果先 resize 再 set_image，坐标就要按缩放比例换算；如果直接拿原图上的坐标去点缩放后的图，位置会整体偏移。这是单点分割最常见的出错原因。

set_image 的开销与复用

set_image 会对整张图做一次编码，是整条流程里比较重的一步。同一张图上试多个提示点时，不要每次都重新 set_image，把 predictor 复用起来即可。反过来，换了图片就必须重新调用，否则预测结果对应的是旧图。

显存与设备

vit_h 模型对显存的要求明显高于 vit_b。显存不足时会出现 CUDA out of memory，可以改用更小的模型类型、缩小输入图片尺寸，或者退回 CPU 运行（速度会慢很多）。代码里用 torch.cuda.is_available() 判断设备，可以避免在没有 GPU 的机器上直接报错。

单点提示的局限

一个点只能表达「我要这个位置的东西」，无法表达边界。当目标与背景颜色接近、或者目标由多个不连通部分组成时，单点结果往往不完整。这时可以增加负样本点排除干扰区域，或者改用框提示限定范围，再在框内补点微调。

notebook 文件需要 Jupyter 打开

源码包的 notebooks 目录下是 .ipynb 文件，用 PyCharm 这类编辑器直接打开看不到可执行内容，需要 Jupyter Notebook 或 JupyterLab 环境。如果只是想在本地跑通分割，直接用上面的 .py 脚本更省事。

版本与许可信息以官方为准

模型权重的下载方式、文件大小、可商用范围，以及依赖包之间的版本兼容关系都可能变化，使用前请以项目官方页面当前公布的信息为准。安装依赖时如果遇到版本冲突，优先保证 torch 与 torchvision 的版本匹配，再装其余包。