
Segment Anything(SAM)
用 Segment Anything(SAM)做提示式图像分割:点、框与多框批量分割
面向 Python 开发者的 SAM 实操教程:从加载 sam_vit_b 权重、生成图像嵌入,到用正负点、矩形框、点框组合以及多框批量预测分割掩膜,并附完整可运行示例与常见坑说明。
Segment Anything(SAM)是 Meta AI 提出的可提示分割模型:给定一张图像,再用点、矩形框这类提示告诉它「要哪块」,它就会返回对应的掩膜。它解决的是「不想为每个类别单独训练分割模型」的问题——同一套权重可以分割任意物体,适合做标注辅助、抠图预处理、遥感或科学图像里的目标提取,也适合作为目标检测结果的下游分割环节。本文用 Python 调用 segment_anything 包,把单点、多点、单框、点框组合、多框批量这几种提示方式逐一跑通。站内也有对应的工具条目可供参考:Segment Anything。
需要先明确一点:SAM 本身不是「输入一句话就出结果」的模型。它接受的是几何提示(点坐标、框坐标)以及可选的上一轮掩膜 logits,文本提示需要借助其他模型把文字转成框或点再喂给它。本文只讲几何提示这条主线。
准备工作
运行环境
需要 Python 环境,以及以下依赖:torch、numpy、opencv-python、matplotlib,以及 segment_anything 包本身。示例代码里还用到 sys 和 datetime,这两个是标准库,无需额外安装。
pip install torch numpy opencv-python matplotlib
pip install segment-anything
如果机器上有可用的 NVIDIA GPU,建议安装对应 CUDA 版本的 PyTorch,并把 device 设为 "cuda";没有 GPU 时设为 "cpu" 也能跑,只是推理会慢一些。具体安装命令与版本对应关系以 PyTorch 官网当前信息为准。
模型权重与模型类型
SAM 提供三种规模的权重,分别对应三种模型类型:vit_b、vit_l、vit_h。文件名与类型必须匹配,例如 sam_vit_b_01ec64.pth 对应 vit_b。加载时用 sam_model_registry[model_type](checkpoint=...),如果类型和权重文件对不上,加载会失败。权重文件体积较大,下载与可用性以官方发布页当前信息为准。
把权重放在本地某个固定目录,代码里用绝对路径引用最省事。示例中写作 sam_checkpoint = "G:/.../sam_vit_b_01ec64.pth",实际使用时替换成自己的路径。
图像与坐标约定
准备一张待分割的图像,示例用的是 B:/truck.jpg。有两点必须记住:
- OpenCV 读进来是 BGR 顺序,要先用
cv2.cvtColor(image, cv2.COLOR_BGR2RGB)转成 RGB,否则颜色会偏。 - 所有提示坐标都是原图像素坐标,格式为
(x, y);框的格式是xyxy,即左上角与右下角坐标[x0, y0, x1, y1]。想确定某个点的坐标,可以先把图画出来并打开坐标轴显示,肉眼读出像素位置。
操作步骤
第一步:准备三个可视化辅助函数
分割结果需要叠加到原图上才能判断好坏,所以先写三个小函数:画掩膜、画提示点、画提示框。它们不参与预测,只负责显示。
import cv2
import sys
import torch
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
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)
def show_box(box, ax):
x0, y0 = box[0], box[1]
w, h = box[2] - box[0], box[3] - box[1]
ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green',
facecolor=(0, 0, 0, 0), lw=2))
约定:标签为 1 的点是前景点(绿色星号,表示「要这个」),标签为 0 的点是背景点(红色星号,表示「不要这个」)。掩膜默认用半透明蓝色叠加,random_color=True 时每块掩膜随机取色,适合一次显示多个目标。
第二步:读图、加载模型、生成图像嵌入
这一步是整个流程里最耗时的部分。图像编码只做一次,之后的每次提示预测都复用这份嵌入,所以同一张图上反复试不同提示是很划算的。
image = cv2.imread(r'B:/truck.jpg')
print("[%s]正在转换图片格式......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
print("[%s]正在初始化模型参数......" % datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
sam_checkpoint = "G:/models/sam_vit_b_01ec64.pth"
model_type = "vit_b"
device = "cuda" # 没有 GPU 时改为 "cpu"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
predictor.set_image(image)
SamPredictor.set_image 会计算并缓存图像嵌入,后续的 predict 调用都基于这份缓存。换一张图就要重新调用一次 set_image。
如果暂时不确定目标点的坐标,可以先单独把图画出来看坐标:
plt.figure(figsize=(10, 10))
plt.imshow(image)
plt.axis('on')
plt.show()
第三步:单点提示,分割一个目标
单点是最简单的提示方式。把 multimask_output 设为 True,模型会返回三个候选掩膜及其置信度分数,因为一个点本身有歧义——它可能落在物体局部、整体或背景上,让模型给出多个解释由人来挑。
input_point = np.array([[250, 187]])
input_label = np.array([1])
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
plt.axis('on')
plt.show()
masks, scores, logit = predictor.predict(
point_coords=input_point,
point_labels=input_label,
multimask_output=True,
)
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()
返回值的形状值得记一下:masks 是 (3, H, W),即三个候选掩膜;scores 是三个对应的置信度;logit 是低分辨率的掩膜 logits,可以留到下一轮当输入。示例图像下 masks.shape 为 (3, 2160, 3840),说明掩膜与原图同尺寸。
第四步:多点提示,用正负点收窄范围
多个点的标签组合决定语义:全部为 1 时,模型会把它们理解为同一个目标的多个证据点;一个 1 加一个 0 时,正点圈定目标、负点排除区域。还可以把上一轮最好的掩膜 logits 通过 mask_input 传回来,帮助模型细化结果。
input_point = np.array([[250, 184], [562, 322]])
input_label = np.array([1, 0])
mask_input = logit[np.argmax(scores), :, :]
masks, _, _ = predictor.predict(
point_coords=input_point,
point_labels=input_label,
mask_input=mask_input[None, :, :],
multimask_output=False,
)
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show()
注意 mask_input 需要额外加一个批次维度,所以写成 mask_input[None, :, :]。这一轮把 multimask_output 设为 False,只返回一个掩膜,形状为 (1, H, W)。
第五步:矩形框提示,分割框内主体
框提示通常比点提示更稳,因为框已经给出了目标的大致范围。传入 xyxy 格式的框,同时把点和标签设为 None。
input_box = np.array([212, 300, 350, 437])
masks, _, _ = predictor.predict(
point_coords=None,
point_labels=None,
box=input_box[None, :],
multimask_output=False,
)
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[0], plt.gca())
show_box(input_box, plt.gca())
plt.axis('off')
plt.show()
框同样需要补批次维度:input_box[None, :]。
第六步:框加点组合,排除框内的干扰部分
框和点可以同时给。典型用法是:框住整个车轮,再用一个负点把车轴部分排除掉,只留下轮胎。这里有一个硬性约束——框只能表示正类,不能给框设负标签;要排除某块区域,只能靠负点。
input_box = np.array([215, 310, 350, 430])
input_point = np.array([[287, 375]])
input_label = np.array([0])
masks, _, _ = predictor.predict(
point_coords=input_point,
point_labels=input_label,
box=input_box,
multimask_output=False,
)
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(masks[0], plt.gca())
show_box(input_box, plt.gca())
show_points(input_point, input_label, plt.gca())
plt.axis('off')
plt.show()
第七步:多框批量预测,一次分割多个目标
当框来自目标检测器的输出时,通常一次就有几十上百个框,逐个调用 predict 效率很低。这时改用 predict_torch,它接受张量形式的框,一次前向就能得到每个框对应的掩膜。
关键点在于:predict_torch 要求框已经过与图像相同的坐标变换,所以要先调用 predictor.transform.apply_boxes_torch 把原始像素坐标的框编码成模型内部的特征坐标。这一步不能省,否则框的位置会错位。
input_boxes = torch.tensor([
[75, 275, 1725, 850],
[425, 600, 700, 875],
[1375, 550, 1650, 800],
[1240, 675, 1400, 750],
], device=predictor.device)
input_boxes = input_boxes / 2
transformed_boxes = predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2])
masks, _, _ = predictor.predict_torch(
point_coords=None,
point_labels=None,
boxes=transformed_boxes,
multimask_output=False,
)
plt.figure(figsize=(10, 10))
plt.imshow(image)
for mask in masks:
show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)
for box in input_boxes:
show_box(box.cpu().numpy(), plt.gca())
plt.axis('off')
plt.show()
返回的 masks 形状为 (batch_size, num_predicted_masks_per_input, H, W),四个框、每个框一个掩膜时就是 [4, 1, H, W]。因为掩膜在 GPU 上,画图前要 .cpu().numpy()。示例里对框做了 / 2 的缩放,这是为了配合示例图像的尺寸,实际使用时按自己的坐标直接传入即可,不要照抄这个除法。
一个完整示例
下面把前面的步骤串成一个可运行脚本:读图、加载模型、单点预测取最佳掩膜、再用框加负点细化,最后把结果保存成 PNG。保存掩膜这一步官方 demo 没有提供,需要自己写。
import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt
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:]
ax.imshow(mask.reshape(h, w, 1) * color.reshape(1, 1, -1))
image = cv2.cvtColor(cv2.imread(r'B:/truck.jpg'), cv2.COLOR_BGR2RGB)
sam = sam_model_registry["vit_b"](checkpoint="G:/models/sam_vit_b_01ec64.pth")
sam.to(device="cuda")
predictor = SamPredictor(sam)
predictor.set_image(image)
# 第一轮:单点,取三个候选中分数最高的
point = np.array([[250, 187]])
label = np.array([1])
masks, scores, logits = predictor.predict(
point_coords=point, point_labels=label, multimask_output=True)
best = int(np.argmax(scores))
# 第二轮:框 + 负点,并用上一轮最佳 logits 作为 mask_input
box = np.array([215, 310, 350, 430])
neg_point = np.array([[287, 375]])
neg_label = np.array([0])
masks2, _, _ = predictor.predict(
point_coords=neg_point,
point_labels=neg_label,
box=box,
mask_input=logits[best][None, :, :],
multimask_output=False,
)
final_mask = masks2[0]
# 可视化
plt.figure(figsize=(10, 10))
plt.imshow(image)
show_mask(final_mask, plt.gca())
plt.axis('off')
plt.show()
# 单独保存二值掩膜
cv2.imwrite('mask.png', (final_mask.astype(np.uint8) * 255))
# 保存叠加效果
overlay = image.copy()
overlay[final_mask] = (overlay[final_mask] * 0.5 + np.array([30, 144, 255]) * 0.5).astype(np.uint8)
cv2.imwrite('overlay.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
脚本跑完后会得到两张图:mask.png 是纯二值掩膜,适合后续做统计或当标注;overlay.png 是叠加了半透明蓝色的效果图,适合人工检查。注意写文件时 OpenCV 期望 BGR,所以叠加图保存前要转回去。
注意事项
模型类型与权重必须一致
vit_b、vit_l、vit_h 三种类型对应不同的权重文件,混用会直接报错。切换模型规模时,model_type 和 sam_checkpoint 要一起改。
坐标是原图像素坐标
点用 (x, y),框用 [x0, y0, x1, y1],都是原图上的像素位置。如果图像被缩放或裁剪过,提示坐标也要同步换算。多框场景下 apply_boxes_torch 负责把像素坐标映射到模型内部坐标,这一步漏掉会导致框错位。
框只能表示正类
框没有负标签的概念,想排除框内某块区域只能用负点。设计提示时如果发现框内混入了不需要的部分,加一个 label=0 的点比换框更有效。
单点结果有歧义
一个点可能对应物体局部、整体或背景,所以 multimask_output=True 会返回三个候选。不要默认取第一个,应该看 scores 或人工比对后再选。多点或加框能显著降低歧义。
图像嵌入只算一次
set_image 是耗时大头,同一张图上反复调提示时不要重复调用。反过来,换图必须重新调用,否则预测结果会对应到上一张图。
官方 demo 不保存掩膜
示例代码只做可视化,没有落盘逻辑。需要掩膜文件时要自己写保存代码,并且注意区分「纯掩膜」和「叠加底图的效果图」——后者对后续处理没有用处。
提示坐标仍需人工确定
当前流程里,点或框的坐标要么靠人看图读出,要么来自上游检测器。没有检测器时,逐张手动标点是主要工作量。把 SAM 接到目标检测模型后面,是减少人工介入的常见做法。
易变信息以官网为准
权重下载地址、可用模型规模、许可证条款、依赖版本要求都可能变化,使用前请核对官方发布页的当前信息。本文不涉及任何价格或配额承诺。