cv_workshop_1 / video.py
kostya-cholak's picture
Init repo
d1f37f2
import numpy as np
from ultralytics import YOLO
import cv2
import math
cap = cv2.VideoCapture('CityUtdR video.mp4')
if not cap.isOpened():
print("Error: Could not open video file.")
exit()
# model
model = YOLO('my_model.pt')
# object classes
classNames = ["ball", ]
while True:
success, img = cap.read()
results = model(img, stream=True)
# coordinates
for r in results:
boxes = r.boxes
for box in boxes:
# bounding box
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # convert to int values
# put box in cam
cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 3)
# object details
org = [x1, y1]
font = cv2.FONT_HERSHEY_SIMPLEX
fontScale = 1
color = (255, 0, 0)
thickness = 2
cv2.putText(img, classNames[np.argmax(box.cls)], org, font, fontScale, color, thickness)
cv2.imshow('Output', img)
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()