RGB & Tiny YOLO

1#!/usr/bin/env python3
2
3"""
4The code is the same as for Tiny Yolo V3 and V4, the only difference is the blob file
5- Tiny YOLOv3: https://github.com/david8862/keras-YOLOv3-model-set
6- Tiny YOLOv4: https://github.com/TNTWEN/OpenVINO-YOLOV4
7"""
8
9from pathlib import Path
10import sys
11import cv2
12import depthai as dai
13import numpy as np
14import time
15
16# Get argument first
17nnPath = str((Path(__file__).parent / Path('../models/yolo-v4-tiny-tf_openvino_2021.4_6shave.blob')).resolve().absolute())
18if 1 < len(sys.argv):
19    arg = sys.argv[1]
20    if arg == "yolo3":
21        nnPath = str((Path(__file__).parent / Path('../models/yolo-v3-tiny-tf_openvino_2021.4_6shave.blob')).resolve().absolute())
22    elif arg == "yolo4":
23        nnPath = str((Path(__file__).parent / Path('../models/yolo-v4-tiny-tf_openvino_2021.4_6shave.blob')).resolve().absolute())
24    else:
25        nnPath = arg
26else:
27    print("Using Tiny YoloV4 model. If you wish to use Tiny YOLOv3, call 'tiny_yolo.py yolo3'")
28
29if not Path(nnPath).exists():
30    import sys
31    raise FileNotFoundError(f'Required file/s not found, please run "{sys.executable} install_requirements.py"')
32
33# tiny yolo v4 label texts
34labelMap = [
35    "person",         "bicycle",    "car",           "motorbike",     "aeroplane",   "bus",           "train",
36    "truck",          "boat",       "traffic light", "fire hydrant",  "stop sign",   "parking meter", "bench",
37    "bird",           "cat",        "dog",           "horse",         "sheep",       "cow",           "elephant",
38    "bear",           "zebra",      "giraffe",       "backpack",      "umbrella",    "handbag",       "tie",
39    "suitcase",       "frisbee",    "skis",          "snowboard",     "sports ball", "kite",          "baseball bat",
40    "baseball glove", "skateboard", "surfboard",     "tennis racket", "bottle",      "wine glass",    "cup",
41    "fork",           "knife",      "spoon",         "bowl",          "banana",      "apple",         "sandwich",
42    "orange",         "broccoli",   "carrot",        "hot dog",       "pizza",       "donut",         "cake",
43    "chair",          "sofa",       "pottedplant",   "bed",           "diningtable", "toilet",        "tvmonitor",
44    "laptop",         "mouse",      "remote",        "keyboard",      "cell phone",  "microwave",     "oven",
45    "toaster",        "sink",       "refrigerator",  "book",          "clock",       "vase",          "scissors",
46    "teddy bear",     "hair drier", "toothbrush"
47]
48
49syncNN = True
50
51# Create pipeline
52pipeline = dai.Pipeline()
53
54# Define sources and outputs
55camRgb = pipeline.create(dai.node.ColorCamera)
56detectionNetwork = pipeline.create(dai.node.YoloDetectionNetwork)
57xoutRgb = pipeline.create(dai.node.XLinkOut)
58nnOut = pipeline.create(dai.node.XLinkOut)
59
60xoutRgb.setStreamName("rgb")
61nnOut.setStreamName("nn")
62
63# Properties
64camRgb.setPreviewSize(416, 416)
65camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
66camRgb.setInterleaved(False)
67camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
68camRgb.setFps(40)
69
70# Network specific settings
71detectionNetwork.setConfidenceThreshold(0.5)
72detectionNetwork.setNumClasses(80)
73detectionNetwork.setCoordinateSize(4)
74detectionNetwork.setAnchors([10, 14, 23, 27, 37, 58, 81, 82, 135, 169, 344, 319])
75detectionNetwork.setAnchorMasks({"side26": [1, 2, 3], "side13": [3, 4, 5]})
76detectionNetwork.setIouThreshold(0.5)
77detectionNetwork.setBlobPath(nnPath)
78detectionNetwork.setNumInferenceThreads(2)
79detectionNetwork.input.setBlocking(False)
80
81# Linking
82camRgb.preview.link(detectionNetwork.input)
83if syncNN:
84    detectionNetwork.passthrough.link(xoutRgb.input)
85else:
86    camRgb.preview.link(xoutRgb.input)
87
88detectionNetwork.out.link(nnOut.input)
89
90# Connect to device and start pipeline
91with dai.Device(pipeline) as device:
92
93    # Output queues will be used to get the rgb frames and nn data from the outputs defined above
94    qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
95    qDet = device.getOutputQueue(name="nn", maxSize=4, blocking=False)
96
97    frame = None
98    detections = []
99    startTime = time.monotonic()
100    counter = 0
101    color2 = (255, 255, 255)
102
103    # nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height
104    def frameNorm(frame, bbox):
105        normVals = np.full(len(bbox), frame.shape[0])
106        normVals[::2] = frame.shape[1]
107        return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)
108
109    def displayFrame(name, frame):
110        color = (255, 0, 0)
111        for detection in detections:
112            bbox = frameNorm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
113            cv2.putText(frame, labelMap[detection.label], (bbox[0] + 10, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
114            cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 40), cv2.FONT_HERSHEY_TRIPLEX, 0.5, 255)
115            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
116        # Show the frame
117        cv2.imshow(name, frame)
118
119    while True:
120        if syncNN:
121            inRgb = qRgb.get()
122            inDet = qDet.get()
123        else:
124            inRgb = qRgb.tryGet()
125            inDet = qDet.tryGet()
126
127        if inRgb is not None:
128            frame = inRgb.getCvFrame()
129            cv2.putText(frame, "NN fps: {:.2f}".format(counter / (time.monotonic() - startTime)),
130                        (2, frame.shape[0] - 4), cv2.FONT_HERSHEY_TRIPLEX, 0.4, color2)
131
132        if inDet is not None:
133            detections = inDet.detections
134            counter += 1
135
136        if frame is not None:
137            displayFrame("rgb", frame)
138
139        if cv2.waitKey(1) == ord('q'):
140            break