-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmain.py
55 lines (46 loc) · 2.15 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from ultralytics import YOLO
import cv2
import cvzone
import math
import time
#cap = cv2.VideoCapture(0) # For Webcam
#cap.set(3, 1280)
#cap.set(4, 720)
cap = cv2.VideoCapture("../Yolo-Weights/vid.mp4") # For Video
model = YOLO("../Yolo-Weights/yolov8n.pt")
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
"traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
"dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
"handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
"baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
"carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
"diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
"teddy bear", "hair drier", "toothbrush"
]
prev_frame_time = 0
new_frame_time = 0
while True:
new_frame_time = time.time()
success, img = cap.read()
results = model(img, stream=True)
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)
# cv2.rectangle(img,(x1,y1),(x2,y2),(255,0,255),3)
w, h = x2 - x1, y2 - y1
cvzone.cornerRect(img, (x1, y1, w, h))
# Confidence
conf = math.ceil((box.conf[0] * 100)) / 100
# Class Name
cls = int(box.cls[0])
cvzone.putTextRect(img, f'{classNames[cls]} {conf}', (max(0, x1), max(35, y1)), scale=1, thickness=1)
fps = 1 / (new_frame_time - prev_frame_time)
prev_frame_time = new_frame_time
print(fps)
cv2.imshow("Image", img)
cv2.waitKey(1)