1.导入需要的库
#coding:utf-8 import numpy as np import torch import matplotlib.pyplot as plt import cv2 import os import copy
2. 定义相关绘图函数
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))
3. 加载图片与SAM模型
image = cv2.imread('notebooks/images/truck.jpg') image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) import sys sys.path.append("..") from segment_anything import sam_model_registry, SamPredictor sam_checkpoint = "./models/sam_vit_b_01ec64.pth" model_type = "vit_b" device = "cpu" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) predictor = SamPredictor(sam) predictor.set_image(image)
4.指定单个提示点，并显示该提示点
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()
5.使用SAM模型对上面的提示点进行目标分割
masks, scores, logits = 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()
6.保存分割结果的mask
由于我使用的是最小的一个SAM预训练模型：sam_vit_b_01ec64.pth，所以分割出来的结果没有那么完美。如果想得到更好的结果，可以使用较大的SAM预训练模型即可。
for i, (mask, score) in enumerate(zip(masks, scores)): mask = mask + 255 plt.imshow(mask, cmap='gray') plt.savefig('pic-{}.png'.format(i + 1)) plt.show()
7.提取并保存mask对应的图片区域
for i, (mask, score) in enumerate(zip(masks, scores)): mask = ~mask mask = mask + 255 mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) mask = mask.astype(np.uint8) res = cv2.bitwise_and(image, mask) res[res == 0] = 255 plt.imshow(res) plt.savefig('res-{}.png'.format(i + 1)) plt.show()
