YOLOv8视频帧多类别检测：正确提取预测类别名称的实践指南

import cv2
from ultralytics import YOLO
import numpy as np

# 假设你已经加载了YOLOv8模型
# yolov8_model_in_heat = YOLO('your_model_path.pt') 

# 模拟一个YOLOv8模型和预测结果
class MockBox:
    def __init__(self, cls_id):
        self.cls = cls_id # cls是一个tensor，实际使用时需要.item()

class MockResult:
    def __init__(self, boxes_data, names_map):
        self.boxes = [MockBox(cls_id) for cls_id in boxes_data]
        self.names = names_map

    def __iter__(self):
        # 模拟results列表的迭代行为
        yield self # 在实际YOLOv8中，model.predict返回的是一个Results对象列表

# 假设模型定义了以下类别
mock_model_names = {0: 'inheat', 1: 'non-inheat'}

def process_frame_for_classes(frame, model, class_names_map):
    """
    处理单个视频帧，并返回检测到的所有对象的类别名称。
    """
    detected_class_names = []
    # 假设model.predict返回一个Results对象列表
    # 实际YOLOv8中，results = model.predict(source=frame_small, show=True, conf=0.8)
    # 这里的mock是为了演示内部逻辑

    # 模拟一个预测结果，假设检测到 'non-inheat' 和 'inheat'
    if np.random.rand() > 0.5: # 随机模拟检测结果
        mock_results = [MockResult(boxes_data=[0, 1], names_map=class_names_map)] # 模拟检测到inheat和non-inheat
    else:
        mock_results = [MockResult(boxes_data=[0], names_map=class_names_map)] # 模拟只检测到inheat

    results = mock_results # 实际代码中替换为 model.predict(frame)

    for result_instance in results:
        # 遍历每个检测到的边界框
        for box in result_instance.boxes:
            # 获取类别ID
            class_id = int(box.cls) # 实际YOLOv8中，box.cls是一个tensor，需要box.cls.item()
            # 根据类别ID从names字典中获取类别名称
            class_name = result_instance.names[class_id]
            detected_class_names.append(class_name)
    return detected_class_names

# 示例使用
# detected_classes = process_frame_for_classes(None, None, mock_model_names)
# print(f"Detected classes in frame: {detected_classes}")

import cv2
from ultralytics import YOLO
import numpy as np

# 假设您已经加载了YOLOv8模型
# yolov8_model_in_heat = YOLO('path/to/your/yolov8_model.pt') 
# 为了演示，这里使用一个占位符
class MockYOLOModel:
    def __init__(self, names_map):
        self._names_map = names_map

    def predict(self, source, show=False, conf=0.8):
        # 模拟YOLOv8的predict方法
        # 在实际应用中，这里会调用真正的模型进行预测
        # 假设根据某种逻辑生成检测结果
        detected_class_ids = []
        if np.random.rand() > 0.7: # 模拟检测到 'inheat'
            detected_class_ids.append(0)
        if np.random.rand() > 0.7: # 模拟检测到 'non-inheat'
            detected_class_ids.append(1)

        # 如果什么都没检测到，随机添加一个
        if not detected_class_ids and np.random.rand() > 0.5:
             detected_class_ids.append(np.random.choice([0, 1]))

        # 构造模拟的Results对象
        boxes_list = [MockBox(cls_id) for cls_id in detected_class_ids]
        mock_result_instance = MockResult(boxes_data=detected_class_ids, names_map=self._names_map)

        # predict返回的是一个Results对象列表
        return [mock_result_instance]

# 实际使用时，请替换为您的模型加载代码
yolov8_model_in_heat = MockYOLOModel(names_map={0: 'inheat', 1: 'non-inheat'}) 

def process_video_with_yolov8_model(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"错误：无法打开视频文件 {video_path}")
        return None

    class_counts = {'inheat': 0, 'non-inheat': 0}
    in_heat_frames = []
    non_in_heat_frames = []

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret: # 当没有更多帧或读取失败时退出
            break

        frame_idx += 1
        # 缩小帧尺寸以提高处理速度，并作为模型输入
        # 注意：模型训练时使用的输入尺寸应与此处保持一致或进行适当调整
        frame_small = cv2.resize(frame, (400, 400))

        # 使用YOLOv8模型进行预测
        # show=True 会在窗口中显示带有边界框的帧，调试时很有用
        results = yolov8_model_in_heat.predict(source=frame_small, show=False, conf=0.5) # 降低conf用于模拟，实际可根据需求设置

        # 遍历每个预测结果实例（通常只有一个）
        for result_instance in results:
            # 遍历每个检测到的边界框
            for box in result_instance.boxes:
                # 获取类别ID（box.cls是一个Tensor，需要使用.item()获取Python数值）
                class_id = int(box.cls.item())
                # 根据类别ID从模型定义的names字典中获取类别名称
                class_name = result_instance.names[class_id]

                # 更新类别计数
                class_counts[class_name] += 1

                # 将帧添加到对应的列表中
                if class_name == 'non-inheat':
                    non_in_heat_frames.append(frame)
                elif class_name == 'inheat':
                    in_heat_frames.append(frame)

        # 打印当前帧的检测计数
        print(f"Frame {frame_idx} - Class Counts: {class_counts}")

        # 达到特定帧数阈值后停止处理（可选，用于控制处理量）
        if class_counts['inheat'] >= 50 and class_counts['non-inheat'] >= 50:
            print("达到指定帧数阈值，停止处理。")
            break

    # 释放视频捕获对象和所有OpenCV窗口
    cap.release()
    cv2.destroyAllWindows()

    # 堆叠帧并显示（如果收集到足够帧）
    if in_heat_frames:
        # 为了显示，需要确保所有帧尺寸相同，这里假设原始帧尺寸已保留
        # 如果需要堆叠不同尺寸的帧，需要先统一尺寸
        stacked_in_heat_frames = np.vstack(in_heat_frames[:50]) # 限制显示前50帧
        cv2.imshow('Stacked In-Heat Frames', stacked_in_heat_frames)
    else:
        print("没有收集到 'inheat' 帧。")

    if non_in_heat_frames:
        stacked_non_in_heat_frames = np.vstack(non_in_heat_frames[:50]) # 限制显示前50帧
        cv2.imshow('Stacked Non-In-Heat Frames', stacked_non_in_heat_frames)
    else:
        print("没有收集到 'non-inheat' 帧。")

    cv2.waitKey(0) # 等待按键，然后关闭显示窗口
    cv2.destroyAllWindows()

    # 比较计数并返回出现次数更多的类别标签
    if class_counts['inheat'] > class_counts['non-inheat']:
        return 'inheat'
    elif class_counts['non-inheat'] > class_counts['inheat']:
        return 'non-inheat'
    else:
        return 'equal_or_no_detections' # 或者根据实际情况处理平局或无检测的情况

# 示例调用 (请替换为实际的视频路径)
# video_file_path = 'your_video.mp4'
# result_label = process_video_with_yolov8_model(video_file_path)
# print(f"Video analysis result: {result_label}")