YOLOv8视频帧多类别目标检测：正确解析与处理预测结果

import cv2
import numpy as np
from ultralytics import YOLO # 假设 yolov8_model_in_heat 是一个YOLO模型实例

# 假设 yolov8_model_in_heat 已经加载
# yolov8_model_in_heat = YOLO('path/to/your/model.pt')

def process_video_with_second_model(video_path, yolov8_model_in_heat):
    """
    使用YOLOv8模型处理视频，对帧进行多类别目标检测，并根据类别分类存储帧。

    Args:
        video_path (str): 视频文件的路径。
        yolov8_model_in_heat (YOLO): 预训练的YOLOv8模型实例。

    Returns:
        str: 具有更高检测计数的类别名称 ('inheat' 或 'non-inheat')。
    """
    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

    print(f"开始处理视频: {video_path}")

    while True:
        ret, frame = cap.read()
        if not ret: # 检查是否成功读取帧
            break # 视频结束或读取失败

        frame_idx += 1
        # 调整帧大小以提高处理速度，同时保持检测效果
        # 注意：如果模型是在特定尺寸下训练的，调整大小可能影响性能
        frame_small = cv2.resize(frame, (640, 480)) # 常用尺寸，可根据模型训练尺寸调整

        # 使用YOLOv8模型进行预测
        # show=False 可以避免实时显示，提高处理速度
        results = yolov8_model_in_heat.predict(source=frame_small, conf=0.5, verbose=False) # conf=0.5 是一个示例阈值

        # 遍历每个检测结果实例 (通常对单帧输入只有一个实例)
        for result_instance in results:
            # 遍历每个检测到的边界框
            for box in result_instance.boxes:
                # 获取类别ID
                class_id = int(box.cls.item()) # .item() 将张量转换为Python数字
                # 获取类别名称
                class_name = result_instance.names[class_id]

                # 确保类别在预期的计数字典中
                if class_name in class_counts:
                    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)
                else:
                    print(f"警告：检测到未知类别 '{class_name}' (ID: {class_id})")

        # 打印当前帧的类别计数
        if frame_idx % 10 == 0: # 每10帧打印一次，避免频繁输出
            print(f"帧 {frame_idx} - 当前类别计数: {class_counts}")

        # 达到指定帧数阈值时停止处理
        if class_counts['inheat'] >= 50 and class_counts['non-inheat'] >= 50:
            print("已达到目标帧数，停止视频处理。")
            break

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

    print(f"视频处理完成。最终类别计数: {class_counts}")

    # 堆叠并显示帧 (如果列表非空)
    if in_heat_frames:
        # 确保所有帧大小一致，以便堆叠
        # 获取第一帧的大小作为基准
        h_in, w_in, _ = in_heat_frames[0].shape
        resized_in_heat_frames = [cv2.resize(f, (w_in, h_in)) for f in in_heat_frames]
        stacked_in_heat_frames = np.vstack(resized_in_heat_frames)
        cv2.imshow('Stacked In-Heat Frames', stacked_in_heat_frames)
    else:
        print("没有检测到 'inheat' 帧。")

    if non_in_heat_frames:
        h_non, w_non, _ = non_in_heat_frames[0].shape
        resized_non_in_heat_frames = [cv2.resize(f, (w_non, h_non)) for f in non_in_heat_frames]
        stacked_non_in_heat_frames = np.vstack(resized_non_in_heat_frames)
        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_counts' # 或者根据需求处理相等的情况

# 示例用法 (请替换为您的模型路径和视频路径)
# if __name__ == "__main__":
#     # 假设您的模型文件名为 'best.pt' 并且在当前目录下
#     my_yolov8_model = YOLO('path/to/your/yolov8_model.pt')
#     video_file = 'path/to/your/video.mp4'
#     dominant_class = process_video_with_second_model(video_file, my_yolov8_model)
#     print(f"视频中主要的类别是: {dominant_class}")