如何从 yolo8 分割结果创建二进制掩码

Question

我想使用 yolo8 分割图像，然后为图像中具有特定类别的所有对象创建一个掩码。

我开发了这段代码：

img=cv2.imread('images/bus.jpg')
model = YOLO('yolov8m-seg.pt')
results = model.predict(source=img.copy(), save=False, save_txt=False)
class_ids = np.array(results[0].boxes.cls.cpu(), dtype="int")
for i in range(len(class_ids)):
    if class_ids[i]==0:
         empty_image = np.zeros((height, width,3), dtype=np.uint8)
         res_plotted = results[0][i].plot(boxes=0, img=empty_image)

在上面的代码中，res_plotted是一个对象的遮罩，采用 RGB 格式。我想将所有这些图像相互添加，并为所有 0 类对象创建一个掩码（在本例中是行人）

我的问题：

1. 我怎样才能完成这段代码？
2. 有没有更好的方法可以在没有循环的情况下实现这一目标？
3. yolo8 库中有任何实用函数可以做到这一点吗？

Answer 1

使用 bbox 类提取人员分段。你将得到一个 shape 数组[channels, w, h]。然后，您可以使用any通道维度（等于人数）将多通道数组展平为单通道数组。

import cv2
from ultralytics import YOLO
import numpy as np
import torch


img= cv2.imread('ultralytics/assets/bus.jpg')
model = YOLO('yolov8m-seg.pt')
results = model.predict(source=img.copy(), save=True, save_txt=False, stream=True)
for result in results:
    # get array results
    masks = result.masks.data
    boxes = result.boxes.data
    # extract classes
    clss = boxes[:, 5]
    # get indices of results where class is 0 (people in COCO)
    people_indices = torch.where(clss == 0)
    # use these indices to extract the relevant masks
    people_masks = masks[people_indices]
    # scale for visualizing results
    people_mask = torch.any(people_masks, dim=0).int() * 255
    # save to file
    cv2.imwrite(str(model.predictor.save_dir / 'merged_segs.jpg'), people_mask.cpu().numpy())

输入 w bbox 和分段/输出：

一切都在 GPU 上通过内部火炬操作进行计算，以实现最佳性能