Segment-Anything｜Meta万物分割模型教程

分割任何物体（SA）项目提出了Segment Anything Model (SAM)，它可以通过输入提示生成高质量的对象掩码，用于图像中所有对象的分割任务。该模型在包含1100万张图像和110亿个掩码的数据集上训练，展现出强大的零样本性能。研究人员还提出了数据引擎和SAM模型，用于训练和使用该项目。该研究提供了一种全新的思路和方法，具有广泛的应用前景，并且可以有效地解决图像分割问题。最终，他们将该数据集和模型公开发布，以促进计算机视觉基础模型的研究。该模型有望取代传统的OpenCV，成为未来图像分割抠图领域的主流模型。

相关链接

项目开源网址：facebookresearch/segment-anythin

在线体验网址：Segment Anything | Meta AI (segment-anything.com)

部署教程

准备工作

如果您是初学者，对于命令行不太理解，那么请按下键盘上的Win键+R键后，在弹出的新窗口内输入CMD并按下回车，打开CMD窗口，按顺序执行如下的每一条命令。

首先我们需要确认一个工作目录，用来存放segment-anything的相关文件。本站所选择的目录为D盘的根目录下openai.wiki文件夹，完整路径为：D:\openai.wiki。

1.检测D盘是否在openai.wiki，没有则创建该文件夹。

if not exist D:\openai.wiki mkdir D:\openai.wiki

2.强制切换工作路径为D盘的openai.wiki文件夹。

cd /d D:\openai.wiki

3.拉取Github仓库文件夹，将下载至openai.wiki文件夹。

git clone https://github.com/facebookresearch/segment-anything.git

注意：如果您无法完成第3步，执行后报错或者无法下载，可以下载该文件将其解压至D:\openai.wiki即可。

环境安装

为不影响电脑中的现有环境，请一定要安装Conda，如果您不知道什么是Conda，或者未安装过Conda，请参考如下文章，安装部署Conda之后再继续以下步骤。

在CMD中执行下面的命令行，创建Conda虚拟环境至该项目的目录中，方便日后重装系统也能够正常使用，无需重新部署环境。

conda create -p D:\openai.wiki\segment-anything\ENV python=3.8

执行完成上面的命令之后，将会在CMD窗口中看到Proceed ([y]/n)?提示，我们直接按下回车即可。

初始化Conda环境，避免后续可能报错。

conda init cmd.exe

激活已创建的Conda环境，这样我们可以将我们后续所需要的所有环境依赖都安装至此环境下。

conda activate D:\openai.wiki\segment-anything\ENV

为避免出现问题，我们再次强制切换至ChatGLM-6B的项目路径，确保一切可以正常安装。

cd /d D:\openai.wiki\segment-anything

PyTorch

在安装PyTorch之前，您需要先前往官网（PyTorch）查看自己应该安装哪个版本。

请您像下图一样按顺序选择即可，依次选择Stable->Windows->Conda->Python->CUDA 11.8。

其中的Stable其实就是稳定版的PyTorch，Preview(Nightly)是每天更新的预览测试版，可能存在未知Bug，所以最好选择Stable版本。

如果您看到的不是CUDA 11.8，而是CUDA 11.9或更高的版本也没关系，只要前缀是CUDA即可。

在您按顺序选择好之后，复制下面的<Run this Command>的内容即可，例如本站所得到的为：

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

将您所复制的代码直接粘贴至CMD窗口内执行，此时将会自动安装PyTorch和对应版本的CUDA，因为PyTorch比较大，所以等待时间会比较长，请不要中断该操作。

注意：一定要复制你自己得到的，尽量不要直接使用本站所给出的示例命令行！

安装segment-anything项目的相关依赖库，此步骤可能需要稍等几分钟才能够安装完成。

python setup.py install

执行如下代码，安装segment-anything项目的额外依赖库。

pip install opencv-python pycocotools matplotlib onnxruntime onnx

模型下载

本站长并未找到三种模型的对比说明，但是自行理解可能每个模型字母分别对应的是h=high=高｜l=low=低｜b=base=基础。

官网下载

• vit_h
  • https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
• vit_l
  • https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth
• vit_b
  • https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth

网盘下载

下载该模型后，将其解压至segment-anything项目的根目录下。

模型下载解压后的正确路径以及目录结构情况如下所示：

D:\openai.wiki\segment-anything\checkpoint
└─sam_vit_b_01ec64.pth
└─sam_vit_h_4b8939.pth
└─sam_vit_l_0b3195.pth

使用教程

本站已经对所有用到的脚本进行整理和修订，解决了所有本站在使用当中遇到的已知BUG和相关问题，您只需要按教程修改部分内容即可直接得到输出结果，非常方便。

无论是全局分割还是局部分割，本站建议您直接复制以下脚本，然后将其粘贴自VS Code代码编辑器内，将代码保存至segment-anything项目的根目录，最后在VS Code代码编辑器的右下角，选择运行的Python版本为3.8，Python路径为D:\openai.wiki\segment-anything\ENV即可。

全局分割

方法说明

使用 SAM 自动生成对象掩码，由于SAM可以有效地处理提示，因此可以通过在图像上采样大量提示来生成整个图像的掩码。这种方法被用来生成SA-1B数据集。

类SamAutomaticMaskGenerator实现了这种能力。它通过在图像上的网格中对单点输入提示进行采样工作，从每个采样点中 SAM 可以预测多个掩码。然后，对掩码进行质量过滤和非最大值抑制去重。其他选项允许进一步提高掩码的质量和数量，例如在图像的多个裁剪区域上运行预测，或对掩码进行后处理以消除小的断开区域和孔洞。

效果展示

本站所提供的脚本内置三种输出方式，得到的效果分别如下：

修改配置

• 第12行
  • 您所下载的模型路径，如果您下载本站所提供的模型则无需修改。
• 第21行
  • 您想使用的模型名称，默认h模型效果最好，建议不要修改。
• 第22行
  • 您需要进行全局自动分割的原图文件路径
• 第23行
  • 如果您需要使用OutMask方法进行分割，则需要填写输出路径，也可以保持默认自动生成。
• 第173行
  • 您想以哪种方式进行自动分割？直接修改为您想分割的方法名称即可。

代码示例

# 此脚本由openai.wiki提供，转载请注明出处。

# 导入相关模块
import os,time,cv2,torch,torch.nn
import numpy as np
import matplotlib.pyplot as plt
import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

# 模型路径
model_Path = './checkpoint'

# 定义模型路径
models = {'b':['%s/sam_vit_b_01ec64.pth' % model_Path,'vit_b'],
          'h':['%s/sam_vit_h_4b8939.pth' % model_Path,'vit_h'],
          'l':['%s/sam_vit_l_0b3195.pth' % model_Path,'vit_l'] 
}

# 必要参数
model = models['h']                           # 直接输入想使用的模型字母即可，可用修改参数为'b'|'h'|'l'
image_path = './notebooks/images/dog.jpg'     # 输入的图片路径
output_folder = './Mask_Folder'               # 输出 Mask 的文件夹

# 窗口定义
def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

# 通用定义
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

sam_checkpoint = model[0]
model_type = model[1]

device = "cuda"

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

# 自动掩码生成
def mask_generator():
    mask_generator = SamAutomaticMaskGenerator(sam)

    masks = mask_generator.generate(image)

    print('mask总数:%s' % len(masks))
    # 掩码生成返回一个掩码列表，其中每个掩码都是一个包含有关掩码的各种数据的字典。 这些键是：

    # 分割：面具
    # area ：掩码的面积（以像素为单位）
    # bbox : XYWH 格式的掩码边界框
    # predicted_iou ：模型自己对掩码质量的预测
    # point_coords ：生成此掩码的采样输入点
    # stability_score ：掩模质量的附加度量
    # crop_box ：用于生成 XYWH 格式蒙版的图像裁剪
    
    print(masks[0].keys())

    plt.figure(figsize=(20,20))
    plt.imshow(image)
    show_anns(masks)
    plt.axis('off')
    plt.show()

# 自动掩码生成选项
# 自动掩码生成中有几个可调参数，用于控制采样点的密度以及去除低质量或重复掩码的阈值。
# 此外，生成可以在图像的裁剪上自动运行以提高较小对象的性能，并且后处理可以去除杂散像素和孔洞。
# 以下是对更多掩码进行采样的示例配置：
def mask_generator_2():
    mask_generator_2 = SamAutomaticMaskGenerator(
        model=sam,
        points_per_side=32,
        pred_iou_thresh=0.86,
        stability_score_thresh=0.92,
        crop_n_layers=1,
        crop_n_points_downscale_factor=2,
        min_mask_region_area=100,           # 需要 opencv 运行后处理
    )

    masks2 = mask_generator_2.generate(image)

    plt.figure(figsize=(20,20))
    plt.imshow(image)
    show_anns(masks2)
    plt.axis('off')
    plt.show() 

def OutMask():
    # 检测输出文件夹是否存在，不存在则自动创建。
    os.makedirs(output_folder, exist_ok=True)

    # 载入模型
    print("%s 模型正在载入" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    sam = sam_model_registry["vit_h"](checkpoint=model)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    sam = sam.to(device)

    # 输出模型加载完成的当前时间
    print("%s 模型载入完成" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

    # 这里是加载图片
    image = cv2.imread(image_path)
    # 输出图片加载完成的current时间
    print("%s 图片加载完成" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

    # 这里是加载图片，这里的image_path是图片的路径
    print("%s 正在分割图像" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    # 这里是预测,不用提示词,进行全图分割
    mask_generator = SamAutomaticMaskGenerator(sam)
    masks = mask_generator.generate(image)

    # 使用提示词,进行局部分割
    # predictor = SamPredictor(sam)
    # predictor.set_image(image)
    # masks, _, _ = predictor.predict(point_coords=None, point_labels=None, box=None, mask_input=None, multimask_output=True, return_logits=False)

    print('%s 图像分割完成' % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))

    # 循环保存mask
    # 遍历 masks 列表并保存每个掩码
    for i, mask in enumerate(masks):
        mask_array = mask['segmentation']
        mask_uint8 = (mask_array * 255).astype(np.uint8)

        # 为每个掩码生成一个唯一的文件名
        output_file = os.path.join(output_folder, f"Mask_{i+1}.png")

        # 保存掩码
        cv2.imwrite(output_file, mask_uint8)


    # 输出完整的mask
    # 获取输入图像的尺寸
    height, width, _ = image.shape

    # 创建一个全零数组，用于合并掩码
    merged_mask = np.zeros((height, width), dtype=np.uint8)

    # 遍历 masks 列表并合并每个掩码
    for i, mask in enumerate(masks):
        mask_array = mask['segmentation']
        mask_uint8 = (mask_array * 255).astype(np.uint8)

        # 为每个掩码生成一个唯一的文件名
        output_file = os.path.join(output_folder, f"Mask_{i+1}.png")

        # 保存掩码
        cv2.imwrite(output_file, mask_uint8)

        # 将当前掩码添加到合并掩码上
        merged_mask = np.maximum(merged_mask, mask_uint8)

    # 保存合并后的掩码
    merged_output_file = os.path.join(output_folder, "Mask_All.png")
    cv2.imwrite(merged_output_file, merged_mask)


# 本站提供三种输方式，分别为mask_generator|mask_generator_2|OutMask三种，请根据自身情况修改。
run_function = mask_generator()               # 运行模式

局部分割

方法说明

使用SAM从提示中获取对象掩码

Segment Anything Model (SAM) 是一种能够根据指定对象的提示，预测对象掩码的强大模型。该模型首先将图像转换为一个图像嵌入，使得可以从提示中高效地生成高质量的掩码。

SamPredictor类提供了一个易于使用的接口，用于与该模型进行交互。它允许用户首先使用set_image方法设置一张图像，并计算出必要的图像嵌入。然后，通过predict方法提供提示，从这些提示中高效地预测出掩码。该模型可以接受点提示、框提示、以及上一次预测的掩码作为输入。

效果展示

本站所提供的脚本内置五种输出方式，得到的效果分别如下：

修改配置

• 第13行
  • 您所下载的模型路径，如果您下载本站所提供的模型则无需修改。
• 第22行
  • 您想使用的模型名称，默认h模型效果最好，建议不要修改。
• 第23行
  • 您需要进行全局自动分割的原图文件路径
• 第251行
  • 您想以哪种方式进行自动分割？直接修改为您想分割的方法名称即可。

代码示例

# 此脚本由openai.wiki提供，转载请注明出处。

# 导入相关模块
import cv2,torch,torch.nn,sys
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

# 模型路径
model_Path = './checkpoint'

# 定义模型路径
models = {'b':['%s/sam_vit_b_01ec64.pth' % model_Path,'vit_b'],
          'h':['%s/sam_vit_h_4b8939.pth' % model_Path,'vit_h'],
          'l':['%s/sam_vit_l_0b3195.pth' % model_Path,'vit_l'] 
}

# 必要参数
model = models['h']                             # 直接输入想使用的模型字母即可，可用修改参数为'b'|'h'|'l'
image_path = './notebooks/images/truck.jpg'     # 输入的图片路径

# 窗口定义
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))  

# 通用定义
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

sam_checkpoint = model[0]
model_type = model[1]

device = "cuda"

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

predictor = SamPredictor(sam)
predictor.set_image(image)


input_point = np.array([[500, 375]])
input_label = np.array([1])

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

# 示例图像 | Example image
def EI():
    masks.shape  # (number_of_masks) x H x W

    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() 
    
# 选择对象的 SAM | Selecting objects with SAM
# 首先，加载SAM模型和预测器。将下面的路径更改为指向SAM检查点。
# 建议在CUDA上运行并使用默认模型以获得最佳结果。
def SOWS():
    input_point = np.array([[500, 375], [1125, 625]])
    input_label = np.array([1, 1])

    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
    masks, _, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        mask_input=mask_input[None, :, :],
        multimask_output=False,
    )
    masks.shape
    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() 
    input_point = np.array([[500, 375], [1125, 625]])
    input_label = np.array([1, 0])

    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
    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()

# 使用框指定特定对象 | Specifying a specific object with a box
# 该模型还可以接受框作为输入，框以xyxy格式提供。
def SASOWAB():
    input_box = np.array([425, 600, 700, 875])
    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()

# 结合点和框 | Combining points and boxes
# 可以通过将两种类型的提示都包含在预测器中来结合点和框。这可以用于选择卡车的轮胎而不是整个车轮。
def CPAB():
    input_box = np.array([425, 600, 700, 875])
    input_point = np.array([[575, 750]])
    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()

# 批量提示输入 | Batched prompt inputs
# SamPredictor可以使用predict_torch方法为同一图像输入多个提示。该方法假定输入点已经是torch张量，并且已经转换为输入帧。例如，假设我们有来自目标检测器的多个框输出。
def BPI():
    input_boxes = torch.tensor([
        [75, 275, 1725, 850],
        [425, 600, 700, 875],
        [1375, 550, 1650, 800],
        [1240, 675, 1400, 750],
    ], device=predictor.device)

    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,
    )
    masks.shape  # (batch_size) x (num_predicted_masks_per_input) x H x W
    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()

# 端到端的批量推断 | End-to-end batched inference
# 如果所有提示都提前可用，则可以直接以端到端的方式运行SAM。这也允许在图像上进行批处理。
def ETEBI():
    image1 = image  # truck.jpg from above
    image1_boxes = torch.tensor([
        [75, 275, 1725, 850],
        [425, 600, 700, 875],
        [1375, 550, 1650, 800],
        [1240, 675, 1400, 750],
    ], device=sam.device)

    image2 = cv2.imread(image_path)
    image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
    image2_boxes = torch.tensor([
        [450, 170, 520, 350],
        [350, 190, 450, 350],
        [500, 170, 580, 350],
        [580, 170, 640, 350],
    ], device=sam.device)
    from segment_anything.utils.transforms import ResizeLongestSide
    resize_transform = ResizeLongestSide(sam.image_encoder.img_size)

    def prepare_image(image, transform, device):
        image = transform.apply_image(image)
        image = torch.as_tensor(image, device=device.device) 
        return image.permute(2, 0, 1).contiguous()
    batched_input = [
        {
            'image': prepare_image(image1, resize_transform, sam),
            'boxes': resize_transform.apply_boxes_torch(image1_boxes, image1.shape[:2]),
            'original_size': image1.shape[:2]
        },
        {
            'image': prepare_image(image2, resize_transform, sam),
            'boxes': resize_transform.apply_boxes_torch(image2_boxes, image2.shape[:2]),
            'original_size': image2.shape[:2]
        }
    ]
    batched_output = sam(batched_input, multimask_output=False)
    # 输出是针对每个输入图像的结果列表，其中列表元素是具有以下键的字典：

    # masks：预测的二进制掩码的批处理torch张量，大小与原始图像相同。
    # iou_predictions：模型对每个掩码的质量预测。
    # low_res_logits：每个掩码的低分辨率logits，可以在稍后的迭代中作为掩码输入传回模型。
    batched_output[0].keys()
    fig, ax = plt.subplots(1, 2, figsize=(20, 20))

    ax[0].imshow(image1)
    for mask in batched_output[0]['masks']:
        show_mask(mask.cpu().numpy(), ax[0], random_color=True)
    for box in image1_boxes:
        show_box(box.cpu().numpy(), ax[0])
    ax[0].axis('off')

    ax[1].imshow(image2)
    for mask in batched_output[1]['masks']:
        show_mask(mask.cpu().numpy(), ax[1], random_color=True)
    for box in image2_boxes:
        show_box(box.cpu().numpy(), ax[1])
    ax[1].axis('off')

    plt.tight_layout()
    plt.show()

# 本站提供五种输方式，分别为SOWS|SASOWAB|CPAB|BPI|ETEBI五种，请根据自身情况修改。
run_function = SOWS()               # 运行模式

总结

经本站测试，暂时并未发现较为明显的问题，但是分割效果真的非常不错。

这个世界正在改变，我在建站之前还焦虑，会不会有一天没有AI相关文章或教程可写了？现在看来，是我想太多了，我每天即使更新三篇也写不完。