YOLOv8视频帧目标检测：精确类别提取与处理指南

import cv2
import numpy as np
from ultralytics import YOLO # 确保已安装ultralytics库

# 假设 yolov8_model_in_heat 已经加载并初始化
# 例如：yolov8_model_in_heat = YOLO('path/to/your/model.pt')

def process_video_with_yolov8_corrected(video_path, yolov8_model):
    """
    使用YOLOv8模型处理视频，并根据检测到的类别对帧进行分类和计数。

    Args:
        video_path (str): 待处理视频文件的路径。
        yolov8_model (YOLO): 已加载的YOLOv8模型实例。

    Returns:
        str: 具有更高检测计数的类别名称 ('inheat', 'non-inheat')，
             如果计数相等或无检测，则返回 'equal_or_no_detection'。
    """
    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 = []

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

    while True:
        ret, frame = cap.read()
        if not ret:
            # 视频帧读取完毕或发生错误
            break

        # 可以选择性地调整帧大小以加快推理速度，
        # 但YOLOv8模型本身支持不同尺寸的输入，并会在内部进行调整。
        # 这里为了保持与原始问题代码一致，保留了手动resize。
        # 注意：如果后续需要堆叠帧，确保所有帧具有相同的尺寸。
        frame_resized = cv2.resize(frame, (640, 480)) # 示例尺寸，可根据模型输入调整

        # 使用YOLOv8模型进行预测
        # conf=0.8 设置了置信度阈值，只有高于此阈值的检测才会被考虑。
        # verbose=False 可以抑制控制台输出的预测细节，使日志更简洁。
        results = yolov8_model.predict(source=frame_resized, conf=0.8, verbose=False)

        # 遍历当前帧的检测结果
        for result in results:
            boxes = result.boxes # 获取所有检测到的边界框

            if len(boxes) > 0:
                # 遍历每个检测到的对象
                for box in boxes:
                    class_id = int(box.cls.item()) # 获取类别索引并转换为整数
                    class_name = result.names[class_id] # 根据索引从names字典获取类别名称

                    # 根据实际检测到的类别名称更新计数和收集帧
                    if class_name == 'non-inheat':
                        class_counts['non-inheat'] += 1
                        # 收集原始帧，如果需要堆叠，确保尺寸一致性
                        non_in_heat_frames.append(frame.copy()) 
                    elif class_name == 'inheat':
                        class_counts['inheat'] += 1
                        in_heat_frames.append(frame.copy())

        # 打印当前帧的类别计数，用于监控进度
        # print(f"当前帧类别计数: {class_counts}")

        # 可选：如果达到特定帧数，提前停止处理
        if class_counts['inheat'] >= 50 and class_counts['non-inheat'] >= 50:
            print("已达到目标帧计数，提前停止视频处理。")
            break

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

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

    # --- 显示堆叠帧（可选，仅用于演示）---
    # 为了堆叠，确保所有帧具有相同的尺寸。这里我们将它们统一调整为400x400。
    display_width, display_height = 400, 400

    if in_heat_frames:
        # 限制堆叠帧的数量以避免内存过载，例如最多堆叠前50帧
        frames_to_stack_inheat = [cv2.resize(f, (display_width, display_height)) for f in in_heat_frames[:50]]
        if frames_to_stack_inheat:
            stacked_in_heat_frames = np.vstack(frames_to_stack_inheat)
            cv2.imshow('Stacked In-Heat Frames', stacked_in_heat_frames)
        else:
            print("没有足够的'inheat'帧用于堆叠显示。")
    else:
        print("未检测到或收集到任何'inheat'帧。")

    if non_in_heat_frames:
        frames_to_stack_noninheat = [cv2.resize(f, (display_width, display_height)) for f in non_in_heat_frames[:50]]
        if frames_to_stack_noninheat:
            stacked_non_in_heat_frames = np.vstack(frames_to_stack_noninheat)
            cv2.imshow('Stacked Non-In-Heat Frames', stacked_non_in_heat_frames)
        else:
            print("没有足够的'non-inheat'帧用于堆叠显示。")
    else:
        print("未检测到或收集到任何'non-inheat'帧。")

    # 等待按键，然后关闭显示窗口
    if in_heat_frames or non_in_heat_frames:
        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_detection' # 处理计数相等或无检测的情况

# 示例用法 (需要替换为实际的模型路径和视频路径)
# if __name__ == '__main__':
#     # 假设你有一个名为 'best.pt' 的YOLOv8模型
#     # yolov8_model_instance = YOLO('best.pt') 
#     # video_file = 'your_video.mp4'
#     # result_label = process_video_with_yolov8_corrected(video_file, yolov8_model_instance)
#     # print(f"视频主要类别: {result_label}")