Segment Anything使用

Segment Anything使用记录

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

一、Segment Anything介绍

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

二、SAM安装

参考官网，很简单

三、利用代码进行介绍

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

fromsegment_anythingimportSamPredictor,sam_model_registry

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

defshow_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

defappend_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)

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

defprediction_with_point(points,labels,results):foriinrange(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)foriinrange(len(masks)):masks=masks[i]show_mask(masks,plt.gca(),color)plt.axis('off')plt.show()