龙空技术网

SAM + YOLO = 自动抠图

新缸中之脑 2179

前言:

如今朋友们对“js抠图”可能比较讲究,咱们都想要分析一些“js抠图”的相关内容。那么小编同时在网摘上汇集了一些对于“js抠图””的相关内容,希望看官们能喜欢,大家快快来了解一下吧!

在计算机视觉领域,对象检测和实例分割是使机器能够理解视觉数据并与之交互的关键任务。 准确识别和隔离图像中的物体的能力具有许多实际应用,从自动驾驶车辆到医学成像。

在这篇博文中,我们将探索如何在 Roboflow 和 Ultralytics YOLOv8 的帮助下使用 Jupyter 笔记本将边界框转换为分割掩模并删除图像的背景。

你可以在 Google Colab 上找到本教程的完整notebook。

在线工具推荐: Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器

1、分割掩模代替边界框的好处

想象一下,你有一个包含感兴趣对象的图像数据集,每个图像都用边界框标注。 虽然边界框提供了有关对象的位置信息,但它们缺乏更高级的计算机视觉任务(例如实例分割或背景去除)所需的精细细节。

将边界框转换为分割掩模使我们能够提取准确的对象边界并将其与背景分开,从而为分析和操作开辟了新的机会。

2、用YOLOv8 和 SAM 创建实例分割数据集

为了解决将边界框转换为分割掩模的挑战,我们将在 Jupyter 笔记本环境中使用 Roboflow 和 Ultralytics 库。 Roboflow 简化了数据准备和标注,而 Ultralytics 提供了最先进的对象检测模型和实用程序。

2.1 搭建notebook环境

pip install roboflow ultralytics 'git+;

我们首先导入必要的包并设置笔记本环境。 下面的代码片段演示了初始设置:

import ultralyticsfrom IPython.display import display, Imagefrom roboflow import Roboflowimport cv2import sysimport numpy as npimport matplotlib.pyplot as plt# Set the device for GPU accelerationdevice = "cuda"# Check Ultralytics version and setup completionultralytics.checks()# Set the first_run flag to False after the initial runfirst_run = Falseif first_run:    !{sys.executable} -m pip install 'git+;    !wget 
2.2 加载数据集

接下来,使用 Roboflow API 加载数据集以访问和管理数据集。 以下代码片段演示了如何从特定项目加载数据集版本:

from roboflow import Roboflowrf = Roboflow(api_key="YOUR_API_KEY")project = rf.workspace("vkr-v2").project("vkrrr")dataset = project.version(5).download("yolov8")
2.3 运行 YOLOv8 推理

为了使用 YOLOv8 执行目标检测,我们运行以下代码:

from ultralytics import YOLO# Load the YOLOv8 modelmodel = YOLO('yolov8n.pt')# Perform object detection on the imageresults = model.predict(source='PATH_TO_IMAGE', conf=0.25)
2.4 提取边界框

一旦我们得到 YOLOv8 的结果,就可以提取检测到的对象的边界框坐标:

for result in results:    boxes = result.boxesbbox = boxes.xyxy.tolist()[0]print(bbox)

结果如下:

[746.568603515625, 40.80133056640625, 1142.08056640625, 712.3660888671875]
2.5 将边界框转换为分割掩模

让我们加载 SAM 模型并将其设置为推理:

from segment_anything import sam_model_registry, image = cv2.cvtColor(cv2.imread('PATH_TO_IMAGE'), cv2.COLOR_BGR2RGB)SamAutomaticMaskGenerator, SamPredictorsam_checkpoint = "sam_vit_h_4b8939.pth"model_type = "vit_h"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)sam.to(device=device)predictor = SamPredictor(sam)predictor.set_image(image)

接下来定义几个用于显示的辅助函数:

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

接下来,我们使用 SAM 模型将边界框坐标转换为分割掩模:

input_box = np.array(bbox)masks, _, _ = predictor.predict(    point_coords=None,    point_labels=None,    box=bbox[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()
2.6 背景去除

最后,我们可以使用分割掩模从图像中去除背景。 以下代码片段演示了该过程:

segmentation_mask = masks[0]# Convert the segmentation mask to a binary maskbinary_mask = np.where(segmentation_mask > 0.5, 1, 0)white_background = np.ones_like(image) * 255# Apply the binary masknew_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]plt.imshow(new_image.astype(np.uint8))plt.axis('off')plt.show()

然后显示生成的图像,结果如下,背景已经剔除干净:

3、结束语

在这篇博文中,我们探索了如何使用 Jupyter Notebook 将边界框转换为分割蒙版并删除背景。

通过利用 Roboflow 和 Ultralytics 库的功能,我们可以轻松执行对象检测、生成分割掩模并操作图像。 这为高级计算机视觉任务(例如实例分割和背景去除)开辟了可能性,并使我们能够从视觉数据中提取有价值的见解。

原文链接:

标签: #js抠图