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"D:\yolo11\ultralytics\runs\detect\train3\weights\best.pt" # best.pt替换成自己的
video_path = r"C:\Users\Administrator\Desktop\Counter-strike 2 2024.12.02 - 19.58.00.11.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)