from ultralytics import YOLO
import json
import cv2 as cv
from sys import argv, exit

if len(argv) != 2:
    exit(1)

m = YOLO("best.pt")

ci = cv.VideoCapture(argv[1])
target_dir = argv[1].rsplit(".", 1)[0] + "-p"

k = 0

while True:
    ret, frame = ci.read()

    if not ret:
        break

    res = m(frame)

    for i in json.loads(res[0].to_json()):
        i = i['box']
        cv.rectangle(frame, (round(i['x1']), round(i['y1'])), (round(i['x2']), round(i['y2'])), (0, 0, 255), 3)

    cv.imwrite(f"{target_dir}/{k:07d}.jpg", frame)
    k += 1

ci.release()