SAM + YOLO 智能抠图

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

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

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

在线工具推荐： - - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器

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

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

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

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

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

2.1 搭建notebook环境

pip install roboflow ultralytics 'git+https://github.com/facebookresearch/segment-anything.git'

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

import ultralytics
from IPython.display import display, Image
from roboflow import Roboflow
import cv2
import sys
import numpy as np
import matplotlib.pyplot as plt

# Set the device for GPU acceleration
device = "cuda"

# Check Ultralytics version and setup completion
ultralytics.checks()

# Set the first_run flag to False after the initial run
first_run = False

if first_run:
    !{sys.executable} -m pip install 'git+https://github.com/facebookresearch/segment-anything.git'
    !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

2.2 加载数据集

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

from roboflow import Roboflow
rf = 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 model
model = YOLO('yolov8n.pt')

# Perform object detection on the image
results = model.predict(source='PATH_TO_IMAGE', conf=0.25)

2.4 提取边界框

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

for result in results:
    boxes = result.boxes
bbox = 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, SamPredictor
sam_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 mask
binary_mask = np.where(segmentation_mask > 0.5, 1, 0)
white_background = np.ones_like(image) * 255

# Apply the binary mask
new_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()

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