main.py 2.76 KB
Newer Older
Pavlo Beylin's avatar
Pavlo Beylin committed
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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()