YOLOv8 视频帧级对象分类与结果解析教程

import cv2
import numpy as np
from ultralytics import YOLO # 假设你已经安装了ultralytics库

# 假设你的YOLOv8模型已经加载
# yolov8_model_in_heat = YOLO('path/to/your/best.pt')

def process_video_with_second_model(video_path, yolov8_model_in_heat):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: Could not open video {video_path}")
        return

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

    print(f"Starting video processing for: {video_path}")

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret or frame is None:
            print(f"End of video or failed to read frame at index {frame_idx}.")
            break

        frame_idx += 1
        # 调整帧大小以适应模型输入或提高处理速度
        # 注意：模型预测时通常会自动处理图像大小，这里仅作示例
        # frame_small = cv2.resize(frame, (640, 640)) # 假设模型输入是640x640

        # 使用YOLOv8模型进行预测
        # show=True 会显示带有检测框的帧，通常在实际应用中会禁用
        results = yolov8_model_in_heat.predict(source=frame, show=False, conf=0.5, verbose=False)

        # 遍历每个预测结果实例
        for result_instance in results:
            # 遍历当前结果实例中检测到的所有边界框
            for box in result_instance.boxes:
                class_id = int(box.cls.item()) # 获取类别ID
                class_name = result_instance.names[class_id] # 通过ID获取类别名称
                confidence = box.conf.item() # 获取置信度

                # 仅处理置信度高于阈值的检测
                if confidence >= 0.8: # 使用0.8作为示例置信度阈值
                    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"Warning: Detected unexpected class: {class_name}")

        print(f"Frame {frame_idx} - Class Counts: {class_counts}")

        # 示例：达到特定帧数后停止处理
        if class_counts['inheat'] >= 50 and class_counts['non-inheat'] >= 50:
            print("Reached target frame counts for both classes. Stopping.")
            break

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

    # 堆叠帧并显示 (注意：直接堆叠原始帧可能非常占用内存，且如果帧大小不一致会报错)
    # 实际应用中，你可能需要对帧进行统一缩放或存储为视频文件
    if in_heat_frames:
        # 为了成功堆叠，确保所有帧具有相同的尺寸。这里假设它们已经相同或已被处理。
        # 如果帧大小不同，需要先统一尺寸，例如：
        # in_heat_frames_resized = [cv2.resize(f, (width, height)) for f in in_heat_frames]
        stacked_in_heat_frames = np.vstack(in_heat_frames[:50]) # 限制堆叠帧数以避免内存溢出
        cv2.imshow('Stacked In-Heat Frames', stacked_in_heat_frames)
    else:
        print("No 'inheat' frames detected to stack.")

    if non_in_heat_frames:
        stacked_non_in_heat_frames = np.vstack(non_in_heat_frames[:50])
        cv2.imshow('Stacked Non-In-Heat Frames', stacked_non_in_heat_frames)
    else:
        print("No 'non-inheat' frames detected to stack.")

    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'

# 示例用法
# 确保替换为你的模型路径和视频路径
# yolov8_model = YOLO('path/to/your/yolov8_custom_model.pt')
# result_label = process_video_with_second_model('path/to/your/video.mp4', yolov8_model)
# print(f"Overall video classification: {result_label}")