mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 10:36:35 +00:00
Alex Bezdieniezhnykh
45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import sys
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
# from vidgear.gears import CamGear
|
|
import cv2
|
|
from time import sleep
|
|
|
|
from predictor import OnnxPredictor, YoloPredictor
|
|
|
|
# video_url = 'https://www.youtube.com/watch?v=d1n2fDOSo8c'
|
|
# stream = CamGear(source=video_url, stream_mode=True, logging=True).start()
|
|
write_output = False
|
|
|
|
predictor = YoloPredictor()
|
|
|
|
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
|
|
|
|
input_name = 'ForAI_test.mp4'
|
|
output_name = Path(input_name).stem + '_recognised.mp4'
|
|
|
|
v_input = cv2.VideoCapture(input_name)
|
|
|
|
if write_output:
|
|
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
|
|
|
|
frame_detected = predictor.predict(frame)
|
|
frame_detected = cv2.resize(frame_detected, (640, 480))
|
|
cv2.imshow('Video', frame_detected)
|
|
sleep(0.01)
|
|
|
|
if write_output:
|
|
v_output.write(frame_detected)
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
break
|
|
|
|
v_input.release()
|
|
if write_output:
|
|
v_output.release()
|
|
cv2.destroyAllWindows()
|