|
|
|
@ -7,38 +7,25 @@
|
|
|
|
|
'''
|
|
|
|
|
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)
|
|
|
|
|
|
|
|
|
|
# 获取原始视频的参数
|
|
|
|
|
def process_video(input_video_path):
|
|
|
|
|
cap = cv2.VideoCapture(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)
|
|
|
|
|
|
|
|
|
|
return frame, frame_width, frame_height, fps
|
|
|
|
|
else:
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
# 释放视频捕获对象和输出视频对象
|
|
|
|
|
cap.release()
|
|
|
|
|
out.release()
|
|
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
process_video()
|
|
|
|
|
video = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\video\0711-1_3_01_0.avi"
|
|
|
|
|
process_video(video)
|