32 lines
567 B
Python
32 lines
567 B
Python
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()
|