Segment Anything使用

Segment Anything使用记录

科研要用，总是忘了怎么用，记录一下

---

一、Segment Anything介绍

一个分割大模型，可以分割anything

二、SAM安装

参考官网，很简单

三、利用代码进行介绍

按照官网的安装方法，针对该conda环境下的python可以任意使用SAM，通过引入

from segment_anything import SamPredictor, sam_model_registry

即可。
SAM的作用是通过一张图片分割成为多个mask，所以我们需要构建绘制mask的函数

def show_mask(mask, ax, color_input):
    color = color_input
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)**

而SAM可以指定prompt，这里介绍使用点的prompt

def append_point(record_point_list,record_label_list, input_point):
    input_label = np.arange(1, len(input_point)+1) 
    record_point_list.append(input_point)
    record_label_list.append(input_label)

那么分割函数就可以写为：

def prediction_with_point(points, labels,results):
    for i in range(len(points)):
        masks, scores, logits = predictor.predict(
            point_coords=points[i],
            point_labels=labels[i],
            multimask_output=False,
        )
        results.append(masks)

主函数需要先初始化

sam_checkpoints = "xxx.pth"
model_type = "vit_h"
device = "cuda"
sam = sam_model_registry[model_type](checkpoint = sam_checkpoints)
sam.to(device = device)
predictor = SamPredictor(sam)

接着就可以进行图片读取并分割，展示

img = cv2.imread('xxx.png')
img = img[:, :, ::-1]
predictor.set_image(img)
points = []
labels = []
# 自主指定二位点，这里的示例为点(a，b)和(c,d)，具体自己给定
# prompt可以指定多个，多次分割后统一上色
append_point(points ,labels ,np.array([[a,b],[c,d]]))
append_point(points ,labels ,np.array([[e,f]]))
masks = []
prediction_with_point(points , labels , masks )

# 指定mask的颜色
color = np.array([144/255, 30/255, 255/255, 0.6])
# 画图
plt.figure(figsize=(10,10))
plt.imshow(img)
for i in range(len(masks)):
    masks = masks[i]
    show_mask(masks, plt.gca(),color)
plt.axis('off')
plt.show()