You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
2 years ago
|
#!/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()
|