YOLOv8多类别分割预测时，如何获取分割实例的类名？

在YOLOv8分割任务中，分割掩码（masks）的顺序和检测框（boxes）的顺序是一一对应的，你可以通过boxes里的类别索引关联到类名，具体实现步骤如下：

1. 从检测结果的boxes中提取每个实例的类别索引：result[0].boxes.cls会返回所有实例的类别张量，每个元素对应一个分割掩码的类别。
2. 利用模型自带的names字典，将类别索引转换为类名：model.names[索引值]即可得到对应类名。

修改你的代码如下：

from ultralytics import YOLO
import numpy as np
import os

model = YOLO("path/to/best.pt")
result = model.predict(os.path.join(cut_dir, im_name), save_conf=True, show=True)

if result[0].masks is not None:
    # 将类别索引张量转为numpy数组，方便遍历
    cls_indices = result[0].boxes.cls.cpu().numpy()
    # 同时遍历掩码和对应的类别索引
    for counter, (detection_mask, cls_idx) in enumerate(zip(result[0].masks.data, cls_indices)):
        detected_mask = np.asarray(detection_mask.cpu())
        # 获取当前分割实例的类名
        class_name = model.names[int(cls_idx)]
        print(f"分割实例 {counter+1} 的类名: {class_name}")

关键说明 #

• result[0].boxes包含了所有检测实例的类别、置信度等元数据，它和masks的元素顺序完全匹配，因此可以直接一一对应关联。
• model.names是模型训练时生成的类别映射字典，键为整数类型的类别索引，值为对应的类名字符串。

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