Commit be3cb146 authored by Pavlo Beylin's avatar Pavlo Beylin
Browse files

Add webcam livetracking.

parent d1182c4f
import torch
import cv2
import time
import matplotlib.pyplot as plt
# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5l') # or yolov5m, yolov5l, yolov5x, cu
# model = torch.hub.load('ultralytics/yolov3', 'yolov3')
def show(img):
plt.imshow(img)
plt.show()
classes = ["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"]
def debug_preds():
detected_classes = [int(results.pred[0][i][-1]) for i in range(0, len(results.pred[0]))]
print(detected_classes)
for det in results.pred[0]:
if int(det[-1]) == 0: # person
print("Person ({}):".format(float(det[-2])))
print("x1:y1 : {}:{}".format(float(det[0]), float(det[1])))
print("x2:y2 : {}:{}".format(float(det[2]), float(det[3])))
if __name__=="__main__":
# set start time to current time
start_time = time.time()
# displays the frame rate every 2 second
display_time = 2
# Set primary FPS to 0
fps = 0
# we create the video capture object cap
cap = cv2.VideoCapture(0)
if not cap.isOpened():
raise IOError("We cannot open webcam")
while True:
ret, frame = cap.read()
# resize our captured frame if we need
frame = cv2.resize(frame, None, fx=1.0, fy=1.0, interpolation=cv2.INTER_AREA)
# detect object on our frame
results = model(frame.copy())
# debug_preds()
# show us frame with detection
# cv2.imshow("Web cam input", frame)
cv2.imshow("img", results.render()[0])
if cv2.waitKey(25) & 0xFF == ord("q"):
cv2.destroyAllWindows()
break
# calculate FPS
fps += 1
TIME = time.time() - start_time
if TIME > display_time:
print("FPS:", fps / TIME)
fps = 0
start_time = time.time()
cap.release()
cv2.destroyAllWindows()
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment