更新视频分帧代码

tool/Video-framing.py

@@ -0,0 +1,44 @@
+#!/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()