YOLOv8视频帧目标分类：正确提取预测类别与帧处理实践

import cv2
import numpy as np
from ultralytics import YOLO

# 假设已经加载了YOLOv8模型
# yolov8_model_in_heat = YOLO('path/to/your/model.pt') 
# 示例：为了演示，我们假设模型能够正确加载
# 在实际应用中，请替换为你的模型路径
class MockYOLOModel:
    def predict(self, source, show=False, conf=0.8):
        # 模拟YOLOv8的预测结果
        # 返回一个包含模拟Results对象的列表
        class MockBox:
            def __init__(self, cls_id):
                self.cls = np.array([cls_id]) # 模拟box.cls是一个包含类别ID的数组
            def item(self):
                return self.cls[0]

        class MockResult:
            def __init__(self, detected_classes):
                self.names = {0: 'inheat', 1: 'non-inheat'}
                self.boxes = [MockBox(cls_id) for cls_id in detected_classes]

        # 模拟检测到不同类别的场景
        if np.random.rand() < 0.5: # 50%概率检测到'inheat'
            return [MockResult(detected_classes=[0])]
        else: # 50%概率检测到'non-inheat'
            return [MockResult(detected_classes=[1])]

yolov8_model_in_heat = MockYOLOModel() # 替换为你的实际YOLO模型加载

# 核心修正逻辑
def extract_class_names_correctly(results):
    detected_classes_in_frame = []
    for result in results:
        # 检查是否有检测框
        if result.boxes:
            for box in result.boxes:
                # 获取类别ID
                class_id = int(box.cls.item())
                # 使用类别ID从names字典中获取类别名称
                class_name = result.names[class_id]
                detected_classes_in_frame.append(class_name)
    return detected_classes_in_frame

# 示例使用
# frame_example = np.zeros((400, 400, 3), dtype=np.uint8) # 模拟一个视频帧
# results_example = yolov8_model_in_heat.predict(source=frame_example, show=False, conf=0.8)
# correct_names = extract_class_names_correctly(results_example)
# print(f"正确提取的类别名称: {correct_names}")

def process_video_with_yolov8(video_path, yolov8_model):
    """
    使用YOLOv8模型处理视频，并根据检测到的类别（inheat/non-inheat）
    将帧分类存储，直到满足特定条件。

    Args:
        video_path (str): 输入视频文件的路径。
        yolov8_model (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_count = 0
    max_frames_to_collect = 50 # 每个类别收集的目标帧数

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

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

        frame_count += 1
        # 示例：每隔N帧处理一次，避免处理所有帧导致性能问题
        # if frame_count % 5 != 0: 
        #    continue

        # 将帧缩放到模型期望的大小或更小的尺寸以提高推理速度
        # 注意：模型训练时使用的图像尺寸会影响检测效果
        frame_resized = cv2.resize(frame, (640, 640)) # 常见YOLOv8输入尺寸

        # 使用YOLOv8模型进行预测
        # show=False 避免在处理大量帧时弹出过多窗口
        results = yolov8_model.predict(source=frame_resized, show=False, conf=0.5)

        # 遍历预测结果，正确提取类别信息
        current_frame_detected_classes = set() # 使用set避免同一帧重复计数
        for result in results:
            if result.boxes: # 确保有检测框
                for box in result.boxes:
                    class_id = int(box.cls.item())
                    # 确保class_id在names字典的有效范围内
                    if class_id in result.names:
                        class_name = result.names[class_id]
                        current_frame_detected_classes.add(class_name)
                    else:
                        print(f"警告: 检测到未知类别ID {class_id}")

        # 根据当前帧检测到的类别更新计数和存储帧
        for detected_class_name in current_frame_detected_classes:
            if detected_class_name in class_counts:
                class_counts[detected_class_name] += 1
                # 仅当该类别帧数未达到上限时才存储
                if detected_class_name == 'non-inheat' and len(non_in_heat_frames) < max_frames_to_collect:
                    non_in_heat_frames.append(frame) # 存储原始大小的帧
                elif detected_class_name == 'inheat' and len(in_heat_frames) < max_frames_to_collect:
                    in_heat_frames.append(frame) # 存储原始大小的帧
            else:
                print(f"警告: 模型检测到未预期的类别: {detected_class_name}")

        print(f"当前帧计数: {frame_count}, 类别计数: {class_counts}")

        # 检查是否已收集到足够数量的帧
        if len(in_heat_frames) >= max_frames_to_collect and \
           len(non_in_heat_frames) >= max_frames_to_collect:
            print(f"已收集到足够数量的帧 ({max_frames_to_collect}每类)。停止处理。")
            break

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

    print(f"最终类别计数: {class_counts}")

    # 堆叠和显示帧 (如果收集到足够帧)
    if in_heat_frames:
        # 确保所有帧尺寸相同才能堆叠
        # 这里为了简化，假设所有in_heat_frames尺寸相同，否则需要resize
        # 示例：取第一个帧的尺寸作为基准
        if in_heat_frames:
            first_in_heat_h, first_in_heat_w, _ = in_heat_frames[0].shape
            resized_in_heat_frames = [cv2.resize(f, (first_in_heat_w, first_in_heat_h)) 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("没有收集到In-Heat帧。")

    if non_in_heat_frames:
        if non_in_heat_frames:
            first_non_in_heat_h, first_non_in_heat_w, _ = non_in_heat_frames[0].shape
            resized_non_in_heat_frames = [cv2.resize(f, (first_non_in_heat_w, first_non_in_heat_h)) 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-In-Heat帧。")

    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__":
    # 替换为你的YOLOv8模型路径
    # 例如: yolov8_model = YOLO('yolov8n.pt') 或 yolov8_model = YOLO('runs/detect/train/weights/best.pt')
    # 为了演示，我们使用上面定义的MockYOLOModel
    yolov8_model_instance = MockYOLOModel() 

    # 替换为你的视频文件路径
    # video_file_path = 'path/to/your/video.mp4' 
    # 假设有一个名为 'test_video.mp4' 的视频文件
    # 为了使示例可运行，我们创建一个虚拟视频文件
    # 这部分仅用于测试，实际使用时请提供真实视频文件
    try:
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter('test_video.mp4', fourcc, 20.0, (640, 480))
        for _ in range(200): # 创建200帧的虚拟视频
            frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
            out.write(frame)
        out.release()
        print("虚拟视频 'test_video.mp4' 创建成功。")
        video_file_path = 'test_video.mp4'
    except Exception as e:
        print(f"创建虚拟视频失败: {e}")
        video_file_path = None # 如果创建失败，则不运行视频处理

    if video_file_path:
        final_prediction = process_video_with_yolov8(video_file_path, yolov8_model_instance)
        print(f"视频处理结果: {final_prediction}")
    else:
        print("未提供有效的视频文件路径，跳过视频处理。")