AB
AiBoss
Tutorials

Segment Anything(SAM)

Tutorials

用 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 接到目标检测模型后面,是减少人工介入的常见做法。

易变信息以官网为准

权重下载地址、可用模型规模、许可证条款、依赖版本要求都可能变化,使用前请核对官方发布页的当前信息。本文不涉及任何价格或配额承诺。