基于YOLO8分割结果生成指定类别二进制掩码的实现与优化疑问

我想要使用YOLO8对图像进行分割，然后为图像中特定类别的所有对象生成掩码。我编写了如下代码：

import cv2
import numpy as np
from ultralytics import YOLO

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")
# 原代码缺失height、width定义，需补充
height, width = img.shape[:2]
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库中是否有相关工具函数可以完成此操作？

1. 补全代码实现合并掩码 #

初始化全黑画布，将每个符合条件的对象掩码叠加到画布上即可。补全后的代码如下：

import cv2
import numpy as np
from ultralytics import YOLO

img = cv2.imread('images/bus.jpg')
height, width = img.shape[:2]
model = YOLO('yolov8m-seg.pt')
results = model.predict(source=img.copy(), save=False, save_txt=False)

# 初始化全黑的合并掩码画布
combined_mask = np.zeros((height, width, 3), dtype=np.uint8)
class_ids = np.array(results[0].boxes.cls.cpu(), dtype="int")

for i in range(len(class_ids)):
    if class_ids[i] == 0:
        empty_img = np.zeros((height, width, 3), dtype=np.uint8)
        single_mask = results[0][i].plot(boxes=0, img=empty_img)
        # 用最大值叠加，避免多个对象掩码互相覆盖
        combined_mask = np.maximum(combined_mask, single_mask)

# 保存或展示结果
cv2.imwrite('combined_pedestrian_mask.jpg', combined_mask)
cv2.imshow('Combined Mask', combined_mask)
cv2.waitKey(0)
cv2.destroyAllWindows()

2. 无需循环的更优实现 #

直接从YOLO的masks属性提取原始掩码数据，跳过plot方法的额外渲染开销，实现无循环合并：

import cv2
import numpy as np
from ultralytics import YOLO

img = cv2.imread('images/bus.jpg')
height, width = img.shape[:2]
model = YOLO('yolov8m-seg.pt')
results = model.predict(source=img.copy(), save=False, save_txt=False)

# 筛选类别0的对象索引
class_ids = np.array(results[0].boxes.cls.cpu(), dtype="int")
target_indices = np.where(class_ids == 0)[0]

if len(target_indices) > 0:
    # 提取所有掩码数据（形状为[N, H, W]，N为检测对象数）
    masks = results[0].masks.data.cpu().numpy()
    # 合并目标掩码，求和后转二值
    combined_mask_data = np.sum(masks[target_indices], axis=0) > 0
    # 转为RGB格式的掩码（白色为目标区域）
    combined_mask = np.zeros((height, width, 3), dtype=np.uint8)
    combined_mask[combined_mask_data] = [255, 255, 255]

cv2.imwrite('combined_pedestrian_mask.jpg', combined_mask)

这种方式直接操作底层掩码张量，效率远高于循环调用plot方法。

3. YOLO8库的工具函数支持 #

YOLO8没有直接生成特定类别合并掩码的内置工具函数，但可以利用results对象的masks、boxes属性快速封装自定义逻辑。比如可以写一个复用性强的函数：

def get_combined_mask(result, target_class):
    class_ids = np.array(result.boxes.cls.cpu(), dtype="int")
    target_indices = np.where(class_ids == target_class)[0]
    if len(target_indices) == 0:
        return np.zeros(result.orig_img.shape, dtype=np.uint8)
    # 提取并合并目标掩码
    masks = result.masks.data.cpu().numpy()[target_indices]
    combined_data = np.sum(masks, axis=0) > 0
    combined_mask = np.zeros(result.orig_img.shape, dtype=np.uint8)
    combined_mask[combined_data] = [255, 255, 255]
    return combined_mask

# 使用示例
combined_mask = get_combined_mask(results[0], 0)

内容的提问来源于stack exchange，提问作者mans