Cpp五条

这个屌丝很懒，什么也没留下！

热门标签

SAM（分割一切模型）的简单调用_昇腾gpu sam分割

作者：Cpp五条 | 2024-04-29 08:34:32

踩

昇腾gpu sam分割

Segment Anything Model （SAM）：Meta AI 的一种新 AI 模型，只需单击一下即可“剪切”任何图像中的任何对象.

由于本人目前的学习需要，了解了一下，sam相关的简单使用并进行了简单的总结，主要包括：指定点对其物体分割、指定框对其物体分割。

一、前期准备(安装、搭建相关环境)

第一步：进入GitHub下载整个项目：https://github.com/facebookresearch/segment-anything

第二步：在其项目环境下，安装安装 PyTorch 和 TorchVision 依赖项。（在硬件条件允许的条件下，建议安装支持 CUDA 的 PyTorch 和 TorchVision，速度更快）,下载 PyTorch 可参考此处。

第三步：确保安装了以下各个内容

pip install opencv-python pycocotools matplotlib onnxruntime onnx

第四步：下载模型检查点,可根据自己的需求选择，目前本人选择的是默认的vit_h.

以上便是准备工作的全部内容。正式开始！

二、查看官方对其物体分割

以下为官方提供内容（单点和单框）的整个过程详解，了解全部案例可访问：https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb

必要导入和帮助程序函数


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))   
 
#加载待处理图片
image = cv2.imread('图片路径')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

首先，加载 SAM 模型和预测变量。更改以下路径以指向 SAM 检查点。


import sys
from segment_anything import sam_model_registry, SamPredictor
 
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
 
device = "cuda"#如果使用gpu
 
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
 
predictor = SamPredictor(sam)

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

predictor.set_image(image)

在其上选择一个点。点以（x，y）格式输入到模型中，并带有标签 1（前景点）或 0（背景点）。可输入多个点。


# 单个点时
input_point = np.array([[500,375]])  # 为要分割的指定点
input_label = np.array([1])  # 为分割对象的性质（背景|前景）
 
 
#为单个框
input_box = np.array([425, 600, 700, 875])


#只是为了显示点，与分割无关
plt.figure(figsize=(10,10))
plt.imshow(image)
show_points(input_point, input_label, plt.gca())
plt.axis('on')
plt.show()

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


#点
masks, scores, logits = predictor.predict(
    point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
)
 
 
# 框
masks, _, _ = predictor.predict(
    point_coords=None,
    point_labels=None,
    box=input_box[None, :],
    multimask_output=False,
)

（使用 multimask_output=True（默认设置），SAM 输出 3 个掩码，其中分数给出了模型自己对这些掩码质量的估计。此设置适用于不明确的输入提示，并帮助模型消除与提示一致的不同对象的歧义。如果为 False，它将返回单个掩码。对于单点等模棱两可的提示，即使只需要一个掩码，也建议使用 multimask_output=True;可以通过选择分数最高的一个来选择最佳的单个掩模。这通常会产生更好的掩码。）——————so,后续本人使用的一直是False.


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())
    show_box(input_box, plt.gca())
    plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
    plt.axis('off')
    plt.show()

三、实践：鼠标单击、画框对物体分割

本人使用opencv来完成鼠标的监听：cv2.setMouseCallback()，点击鼠标后的回调函数.


cv2.setMouseCallback(windowName, onMouse [, userdata])
相关参数说明如下：
windowName:窗口的名字
onMouse:鼠标响应函数，回调函数
userdata：传给回调函数的参数

本人使用的回调函数是: on_mouse,以on_mouse为例介绍回调函数.


on_mouse(event, x, y, flags, param)
相关参数说明如下：
event 是 CV_EVENT_* 变量之一
x 和 y 是鼠标在图像坐标系的坐标（不是窗口坐标系）
flags 是 CV_EVENT_FLAG 的组合
param 是用户定义的传递到 setMouseCallback 函数调用的参数

以下是event、flags的可选值：

单击画框实践代码：


import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import SamPredictor, sam_model_registry
 
sam = sam_model_registry["vit_h"](checkpoint='sam_vit_h_4b8939.pth')
# sam.to(device="cuda")  # 使用gpu
predictor = SamPredictor(sam)
 
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 on_mouse(event,x,y,flags,param):
    global img, point1
    img3 = img.copy()
    if event == cv2.EVENT_LBUTTONDOWN:#左键点击
        input_point = np.array([[x, y]])
        input_label = np.array([1])
        img3 = cv2.cvtColor(img3, cv2.COLOR_BGR2RGB)
        predictor.set_image(img3)
 
        masks, scores, logits = predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=False,
        )
        # 保存
        plt.imshow(img3)
        show_mask(masks, plt.gca())
        plt.savefig("1.png")
 
 
if __name__ == '__main__':
    path = "图片路径"
    img = cv2.imread(path)
    cv2.namedWindow('image')
    cv2.setMouseCallback('image',on_mouse)
    cv2.imshow('image', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

画框实践代码：


import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import SamPredictor, sam_model_registry
 
sam = sam_model_registry["vit_h"](checkpoint='sam_vit_h_4b8939.pth')
# sam.to(device="cuda")  # gpu
predictor = SamPredictor(sam)
 
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 on_mouse(event,x,y,flags,param):
    global img, point1, point2
    img2 = img.copy()
    img3 = img.copy()
    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        point1 = (x, y)
        cv2.circle(img2, point1, 10, (0, 255, 0), 5)
        cv2.imshow('image', img2)
 
    elif event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_LBUTTON):  # 移动鼠标，左键拖拽
        cv2.rectangle(img2, point1, (x, y), (255, 0, 0), 15)  # 需要确定的就是矩形的两个点（左上角与右下角），颜色红色，线的类型（不设置就默认）。
        cv2.imshow('image', img2)
 
    elif event == cv2.EVENT_LBUTTONUP:  # 左键释放
        point2 = (x, y)
        cv2.rectangle(img2, point1, point2, (0, 0, 255), 5)  # 需要确定的就是矩形的两个点（左上角与右下角），颜色蓝色，线的类型（不设置就默认）。
        cv2.imshow('image', img2)
        min_x = min(point1[0], point2[0])
        min_y = min(point1[1], point2[1])
        width = abs(point1[0] - point2[0])
        height = abs(point1[1] - point2[1])
        roi = [min_x, min_y, min_x + width, min_y + height]
        roi = np.array(roi).astype(dtype=int).tolist()
 
        image = cv2.cvtColor(img3, cv2.COLOR_BGR2RGB)
        predictor.set_image(image)
 
        # 选区
        input_box = np.array(roi)
        masks, _, _ = predictor.predict(
            point_coords=None,
            point_labels=None,
            box=input_box[None, :],
            multimask_output=False,
        )
        # 保存
        plt.imshow(img3)
        show_mask(masks, plt.gca())
        plt.savefig("1.png")
 
 
if __name__ == '__main__':
    path = "图片路径"
    img = cv2.imread(path)
    cv2.namedWindow('image')
    cv2.setMouseCallback('image',on_mouse)
    cv2.imshow('image', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

注：若想将图片以cv的方式处理，需要进行以下的转换


# 处理show_mask内结果
mask_image *= 250.0
mask2 = mask_image.astype(np.uint8)

四、将分割的内容单独剪切出来

本人最终需要将上述分割图片进行粘贴，以此进行了图像的处理以便于进行图像融合

1、灰度、二值图像


image_gray = cv2.cvtColor(mask2, cv2.COLOR_BGR2GRAY)  # 以灰度方法读取图像
ret1, th1 = cv2.threshold(image_gray, 0, 255, cv2.THRESH_OTSU)  # 方法选择为THRESH_OTSU

2、将黑色非背景区域使用掩膜去除


# 按位与运算
pic = cv2.bitwise_and(img3, img3, mask=th1)

3、按其掩膜区域进行剪裁

streamline_pic = pic[roi[1]:roi[3], roi[0]:roi[2]]

五、分割图像的背景融合

图像融合，本人使用cv2.seamlessClone：


cv2.seamlessClone(src, dst, mask, center, flags)
相关参数说明如下：
src：目标影像
dst：背景图像
mask：目标影像上的mask，表示目标影像上那些区域是感兴趣区域。
center：目标影像的中心在背景图像上的坐标！注意是目标影像的中心！
flags：选择融合的方式：
    NORMAL_CLONE（不保留dst 图像的texture细节）
    MIXED_CLONE（保留dest图像的texture 细节）
    MONOCHROME_TRANSFER（ 不保留src图像的颜色细节，只有src图像的质地，颜色和目标图像一样，可以用来进行皮肤质地填充）

具体使用：


img = cv2.imread(path_back)  # 大图（背景）
image = cv2.imread(path_stick)  # 小图（前景）
 
h_img, w_img = img.shape[:2]
h_image, w_image = image.shape[:2]
mask = 255 * np.ones(image.shape, image.dtype)
w = random.randrange(int(0.5 * w_image), int(w_img - 0.5 * w_image - 1))
h = random.randrange(int(0.5 * h_image), int(h_img - 0.5 * h_image - 1))
center = (w, h)
normal_clone = cv2.seamlessClone(image, img, mask, center, cv2.MIXED_CLONE)
 
cv2.imwrite("2.jpg", normal_clone)

六、目前的SAM衍生的标注工具

可参考以下文章完成：https://blog.csdn.net/weixin_45977690/article/details/130088039文章浏览阅读1.5w次，点赞51次，收藏238次。详细教程，使用Segment Anything(SAM)模型当工具进行自己数据的自动标注。https://blog.csdn.net/weixin_45977690/article/details/130088039

利用SAM实现自动标注-CSDN博客文章浏览阅读943次。（2）检测图像的文件(可调整后面的图片高/宽)：python helpers/generate_onnx.py --checkpoint-path sam_vit_h_4b8939.pth --onnx-model-path ./sam_onnx.onnx --orig-im-size 720 1280。/segment-anything/dataset/ -a …（3）运行完会有对应的sam_onnx.onnx文件，将其移到SAM工具主文件夹中：cp sam_onnx.onnx …/SAM-Tool/https://blog.csdn.net/qq_37249793/article/details/131956211本人是自己实践编写了自己的工具，未实践上述方法，以上便是全部内容，感谢您可以看到这里，手动鞠躬~

本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】