#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
@File    :   prediction.py
@Time    :   2023/07/26
@Author  :   Fpt 
'''
import cv2
from ultralytics import YOLO
import config

def process_video():
    # 加载 YOLOv8 模型
    model = YOLO(config.model_file_path)
    cap = cv2.VideoCapture(config.input_video_path)

    # 获取原始视频的参数
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    codec = cv2.VideoWriter_fourcc(*'XVID')
    # 创建输出视频文件的 VideoWriter 对象
    out = cv2.VideoWriter(config.output_video_path, codec, fps, (frame_width, frame_height))

    # 遍历视频帧
    while cap.isOpened():
        success, frame = cap.read()
        if success:
            # 对帧运行 YOLOv8 推理
            results = model(frame)
            annotated_frame = results[0].plot()
            
            out.write(annotated_frame)

        else:
            break

    # 释放视频捕获对象和输出视频对象
    cap.release()
    out.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    process_video()