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()