语义分割——SAM 模型环境配置和使用教程_sam语义分割

目录

语义分割——SAM 模型环境配置和使用教程

• 【SAM 介绍】：Meta新模型“分割一切”：抠图完成究极进化，计算机视觉迎来GPT-3时刻 Meta 在论文中发布的新模型名叫 Segment Anything Model (SAM) 。他们在博客中介绍说，「S… - 雪球 (xueqiu.com)

• 【项目地址】：facebookresearch/segment-anything：存储库提供了使用 SegmentAnything 模型 （SAM） 运行推理的代码、用于下载经过训练的模型检查点的链接，以及演示如何使用该模型的示例笔记本。 (github.com)

• 【官网体验地址】：Segment Anything | Meta AI (segment-anything.com)

• 【官网体验教程】：Meta AI 开源万物可分割 AI 模型(SAM) - 努力的小雨 - 博客园 (cnblogs.com)

注：作者用的 第一个方法.

方法一：pip

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

方法二：在本地 git 存储库并使用

git clone git@github.com:facebookresearch/segment-anything.git
cd segment-anything; pip install -e .
1
2

方法三：手动下载

• 链接：https://pan.baidu.com/s/1T-OPGiWVdM18uedR9V8imQ?pwd=0vhi

• 提取码：0vhi

---

可选依赖项

• 掩码后处理、以 COCO 格式保存掩码、示例笔记本以及以 ONNX 格式导出模型时，需要这个可选依赖项：

pip install opencv-python pycocotools matplotlib onnxruntime onnx
1

---

下载模型

注：单击下面的链接下载相应模型，作者用 第一个；也可以去作者网盘下载.

• 默认或vit_h：ViT-H SAM 模型。

• vit_l：ViT-L SAM 模型。

• vit_b：ViT-B SAM 模型。

• 链接：https://pan.baidu.com/s/1Wp4_3eWJ9jhNnLeE3awpSA?pwd=5c1o

• 提取码：5c1o

---

predictor 示例

注：需要自定义修改sam_checkpoint、image = cv2.imread('images/truck.jpg')路径.

设置

• 必要的导入和辅助函数，用于显示点、框和掩码：

import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2

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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

示例图片

• 链接：https://pan.baidu.com/s/1266nx0_tqJqY7buBaNpuvw?pwd=dcxw
• 提取码：dcxw

image = cv2.imread('images/truck.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(10,10))
plt.imshow(image)
plt.axis('on')
plt.show()
1
2
3
4
5
6
7

使用 SAM 选择对象

• 首先，加载 SAM 模型和预测器。更改下面的路径以指向 SAM 检查点。建议在 CUDA 上运行并使用默认模型以获得最佳效果。

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "models/sam_vit_h_4b8939.pth"
model_type = "vit_h"

device = "cuda"

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

predictor = SamPredictor(sam)
1
2
3
4
5
6
7
8
9
10
11
12
13

• 通过调用 ‘SamPredictor.set_image’ 来处理图像以生成图像嵌入。“SamPredictor”会记住这个嵌入，并将其用于后续的掩码预测。

predictor.set_image(image)
1

• 要选择卡车，请选择卡车上的一个点。点以 (x,y) 格式输入模型，并带有标签 1(前景点) 或 0(背景点)。可输入多个点，这里我们只用一个。所选的点将在图像上显示为星形。

input_point = np.array([[500, 375]])
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()  
1
2
3
4
5
6
7
8

• 使用 ‘SamPredictor.predict’ 进行预测。该模型返回掩码，这些掩码的质量预测，以及可以传递给下一次预测迭代的低分辨率掩码逻辑。

masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)
1
2
3
4
5

• 使用multimask_output=True(默认设置)，SAM 输出 3 个掩码，其中scores给出了模型自己对这些掩码质量的估计。此设置用于歧义输入提示，并帮助模型消除与提示一致的不同对象的歧义。当False时，它将返回单个掩码。对于模棱两可的提示，例如单个点，即使只需要一个掩码，也建议使用multimask_output=True；最好的单个掩码可以通过选择在“分数”中返回的分数最高的掩码来选择。这通常会产生一个更好的 mask。

masks.shape  # (3, 1200, 1800)

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()  
1
2
3
4
5
6
7
8
9
10

指定带有附加点的特定对象

• 单个输入点是不明确的，模型返回了与它一致的多个对象。为了获得单个对象，可以提供多个点。如果可用，还可以向模型提供以前迭代的掩码，以帮助进行预测。当使用多个提示指定单个对象时，可以通过设置multimask_output=False来请求单个掩码。

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

mask_input = logits[np.argmax(scores), :, :]  # 选择模型最好的 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() 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

• 为了排除汽车并只指定窗口，可以提供一个背景点（标签为 0，这里用红色显示）。

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() 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

用方框指定一个特定的对象

• 该模型还可以采用 xyxy 格式的框作为输入。

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

结合点和框

• 积分和方框可以结合在一起，只需将两种类型的提示都包含到预测器中。在这里，这可以用来选择卡车的轮胎，而不是整个车轮。

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

批量提示输入

• SamPredictor可以使用predict_torch方法为相同的图像接受多个输入提示。该方法假定输入点已经是 torch 张量，并且已经转换为输入帧。例如，假设我们有几个来自对象检测器的框输出。

input_boxes = torch.tensor([
    [75, 275, 1725, 850],
    [425, 600, 700, 875],
    [1375, 550, 1650, 800],
    [1240, 675, 1400, 750],
], device=predictor.device)
1
2
3
4
5
6

• 将盒子转换为输入帧，然后预测蒙版。SamPredictor将必要的转换存储为transform字段，以便于访问。

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  # torch.Size([4, 1, 1200, 1800])

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

端到端批处理推理

• 如果预先提供了所有提示，则可以直接以端到端方式运行 SAM。这也允许对图像进行批处理。

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('images/groceries.jpg')
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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

• 图像和提示都是作为 PyTorch 张量输入的，这些张量已经转换为正确的帧。输入被打包成一个包含图像的列表，其中每个元素都是一个字典，接受以下键:

  • image：以 CHW 格式的 PyTorch 张量的形式输入图像。
  • original_size：图像在转换为 SAM 之前的大小，格式为(H, W)。
  • point_coords：点提示的批处理坐标。
  • point_labels：批处理点提示的标签。
  • boxes：批量输入框。
  • mask_inputs：批处理输入掩码。

  如果提示符不存在，则可以排除该键。

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]
     }
]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

• 运行模型

batched_output = sam(batched_input, multimask_output=False)

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22