mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 10:26:36 +00:00
Alex Bezdieniezhnykh
40 lines
1.0 KiB
Python
40 lines
1.0 KiB
Python
import sys
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
# from vidgear.gears import CamGear
|
|
import cv2
|
|
from time import sleep
|
|
|
|
model = YOLO('azaion-2024-08-13.pt')
|
|
|
|
# video_url = 'https://www.youtube.com/watch?v=d1n2fDOSo8c'
|
|
# stream = CamGear(source=video_url, stream_mode=True, logging=True).start()
|
|
|
|
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
|
|
|
|
input_name = sys.argv[1]
|
|
output_name = Path(input_name).stem + '_recognised.mp4'
|
|
|
|
v_input = cv2.VideoCapture(input_name)
|
|
v_output = cv2.VideoWriter(output_name, fourcc, 20.0, (640, 480))
|
|
|
|
while v_input.isOpened():
|
|
ret, frame = v_input.read()
|
|
if frame is None:
|
|
break
|
|
|
|
results = model.track(frame, persist=True, tracker='bytetrack.yaml')
|
|
frame_detected = results[0].plot()
|
|
|
|
frame_detected = cv2.resize(frame_detected, (640, 480))
|
|
cv2.imshow('Video', frame_detected)
|
|
sleep(0.01)
|
|
|
|
v_output.write(frame_detected)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
v_input.release()
|
|
v_output.release()
|
|
cv2.destroyAllWindows()
|