YOLOv8 分割任务中获取多类别实例名称的实用指南

import os
from ultralytics import YOLO
import numpy as np

# 假设您的模型路径和图片路径
model_path = "path/to/best.pt"  # 替换为您的模型路径
image_directory = "path/to/images" # 替换为您的图片目录
image_name = "test_image.jpg" # 替换为您的图片名称

# 1. 加载YOLOv8分割模型
try:
    model = YOLO(model_path)
except FileNotFoundError:
    print(f"错误: 无法找到模型文件 '{model_path}'。请检查路径是否正确。")
    exit()

# 构建完整的图片路径
image_path = os.path.join(image_directory, image_name)

# 2. 进行预测
try:
    # save_conf=True 保存置信度, show=True 可视化显示结果
    results = model.predict(image_path, save_conf=True, show=True)
except FileNotFoundError:
    print(f"错误: 无法找到图片文件 '{image_path}'。请检查路径是否正确。")
    exit()
except Exception as e:
    print(f"预测过程中发生错误: {e}")
    exit()

# 检查是否有结果返回
if results:
    result = results[0] # 通常对于单张图片，结果在列表的第一个元素

    # 3. 检查是否存在分割掩码
    if result.masks is not None:
        print(f"\n在图片 '{image_name}' 中检测到 {len(result.masks.data)} 个分割实例。")

        # 4. 遍历每个检测到的分割实例
        for counter, detection_mask_tensor in enumerate(result.masks.data):
            # 获取分割掩码（转换为NumPy数组，如果需要进一步处理）
            detected_mask_np = np.asarray(detection_mask_tensor.cpu())

            # 5. 获取对应的类别ID
            # result.boxes[counter].cls 是一个张量，需要用 .item() 转换为Python整数
            cls_id = int(result.boxes[counter].cls.item())

            # 6. 使用model.names映射类别ID到类别名称
            cls_name = model.names[cls_id]

            # 打印或进一步处理结果
            print(f"  实例 {counter + 1}:")
            print(f"    类别ID: {cls_id}")
            print(f"    类别名称: {cls_name}")
            # print(f"    分割掩码形状: {detected_mask_np.shape}") # 掩码形状通常是 HxW

            # 您可以在这里对 detected_mask_np 进行进一步处理，
            # 例如保存为图像、计算面积、提取轮廓等

    else:
        print(f"在图片 '{image_name}' 中未检测到任何分割实例。")
else:
    print(f"未从图片 '{image_name}' 获取到任何预测结果。")