Spaces:
Sleeping
Sleeping
| 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() | |