import os
import cv2
from ultralytics import YOLO
def detect_objects_in_video(best_pt_path, video_path, output_video_name):
output_video_path = video_path.rsplit('.', 1)[0] + '_' + output_video_name + '.mp4'
model = YOLO(best_pt_path)
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
while cap.isOpened():
success, frame = cap.read()
if success:
results = model(frame)
annotated_frame = results[0].plot()
out.write(annotated_frame)
cv2.imshow('YOLO Detection', annotated_frame)
if cv2.waitKey(1) & 0xFF == ord('q'): #退出循环的话按“q”
break
else:
break
cap.release()
out.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
best_pt_path = r"C:\Users\DELL\Desktop\best.pt" #best.pt替换成自己的
video_path = r"C:\Users\DELL\Desktop\a.mp4" #原视频路径
output_video_name = "out"
detect_objects_in_video(best_pt_path, video_path, output_video_name)
output_video_path = video_path.rsplit('.', 1)[0] + '_' + output_video_name + '.mp4'
os.startfile(output_video_path)
使用方法:将YOLO训练好的best.pt文件路径和你的视频路径放入代码中,会在可视化结束后输出视频到原视频的同级文件夹
YOLO使用教程:YOLOv8改进专栏|专栏介绍&目录-优快云博客
具体效果请看:
YOLO本地检测并保存