RGB Encoding & MobilenetSSD

1#!/usr/bin/env python3
2
3from pathlib import Path
4import sys
5import cv2
6import depthai as dai
7import numpy as np
8
9# Get argument first
10nnPath = str((Path(__file__).parent / Path('../models/mobilenet-ssd_openvino_2021.4_6shave.blob')).resolve().absolute())
11if len(sys.argv) > 1:
12    nnPath = sys.argv[1]
13
14if not Path(nnPath).exists():
15    import sys
16    raise FileNotFoundError(f'Required file/s not found, please run "{sys.executable} install_requirements.py"')
17
18# MobilenetSSD label texts
19labelMap = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow",
20            "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
21
22# Create pipeline
23pipeline = dai.Pipeline()
24
25# Define sources and outputs
26camRgb = pipeline.create(dai.node.ColorCamera)
27videoEncoder = pipeline.create(dai.node.VideoEncoder)
28nn = pipeline.create(dai.node.MobileNetDetectionNetwork)
29
30xoutRgb = pipeline.create(dai.node.XLinkOut)
31videoOut = pipeline.create(dai.node.XLinkOut)
32nnOut = pipeline.create(dai.node.XLinkOut)
33
34xoutRgb.setStreamName("rgb")
35videoOut.setStreamName("h265")
36nnOut.setStreamName("nn")
37
38# Properties
39camRgb.setBoardSocket(dai.CameraBoardSocket.CAM_A)
40camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
41camRgb.setPreviewSize(300, 300)
42camRgb.setInterleaved(False)
43
44videoEncoder.setDefaultProfilePreset(30, dai.VideoEncoderProperties.Profile.H265_MAIN)
45
46nn.setConfidenceThreshold(0.5)
47nn.setBlobPath(nnPath)
48nn.setNumInferenceThreads(2)
49nn.input.setBlocking(False)
50
51# Linking
52camRgb.video.link(videoEncoder.input)
53camRgb.preview.link(xoutRgb.input)
54camRgb.preview.link(nn.input)
55videoEncoder.bitstream.link(videoOut.input)
56nn.out.link(nnOut.input)
57
58# Connect to device and start pipeline
59with dai.Device(pipeline) as device, open('video.h265', 'wb') as videoFile:
60
61    # Queues
62    queue_size = 8
63    qRgb = device.getOutputQueue("rgb", queue_size)
64    qDet = device.getOutputQueue("nn", queue_size)
65    qRgbEnc = device.getOutputQueue('h265', maxSize=30, blocking=True)
66
67    frame = None
68    detections = []
69
70    def frameNorm(frame, bbox):
71        normVals = np.full(len(bbox), frame.shape[0])
72        normVals[::2] = frame.shape[1]
73        return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)
74
75    def displayFrame(name, frame):
76        color = (255, 0, 0)
77        for detection in detections:
78            bbox = frameNorm(frame, (detection.xmin, detection.ymin, detection.xmax, detection.ymax))
79            cv2.putText(frame, labelMap[detection.label], (bbox[0] + 10, bbox[1] + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
80            cv2.putText(frame, f"{int(detection.confidence * 100)}%", (bbox[0] + 10, bbox[1] + 40), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
81            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
82        # Show the frame
83        cv2.imshow(name, frame)
84
85    while True:
86        inRgb = qRgb.tryGet()
87        inDet = qDet.tryGet()
88
89        while qRgbEnc.has():
90            qRgbEnc.get().getData().tofile(videoFile)
91
92        if inRgb is not None:
93            frame = inRgb.getCvFrame()
94
95        if inDet is not None:
96            detections = inDet.detections
97
98        if frame is not None:
99            displayFrame("rgb", frame)
100
101        if cv2.waitKey(1) == ord('q'):
102            break
103
104print("To view the encoded data, convert the stream file (.h265) into a video file (.mp4), using a command below:")
105print("ffmpeg -framerate 30 -i video.h265 -c copy video.mp4")