From 72c86443f9e23f6f7231b639b6f61845936c429a Mon Sep 17 00:00:00 2001
From: Oleksandr Bezdieniezhnykh <oleksandr.bez@2c2p.com>
Date: Thu, 15 Aug 2024 23:38:08 +0300
Subject: [PATCH] add test-prediction.py

---
 test-prediction.py | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 test-prediction.py

diff --git a/test-prediction.py b/test-prediction.py
new file mode 100644
index 0000000..1d034ae
--- /dev/null
+++ b/test-prediction.py
@@ -0,0 +1,40 @@
+import sys
+from pathlib import Path
+from ultralytics import YOLO
+# from vidgear.gears import CamGear
+import cv2
+from time import sleep
+from os import path
+
+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(path.join('tests', input_name))
+v_output = cv2.VideoWriter('', 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.03)
+
+    v_output.write(frame_detected)
+    if cv2.waitKey(1) & 0xFF == ord('q'):
+        break
+
+v_input.release()
+v_output.release()
+cv2.destroyAllWindows()