Detection network Replay

1#!/usr/bin/env python3
2
3from pathlib import Path
4import cv2
5import depthai as dai
6import numpy as np
7import time
8
9# Create pipeline
10with dai.Pipeline() as pipeline:
11    cameraNode = pipeline.create(dai.node.Camera).build()
12    detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build(cameraNode, dai.NNModelDescription("yolov6-nano"))
13    labelMap = detectionNetwork.getClasses()
14
15    qRgb = detectionNetwork.passthrough.createOutputQueue()
16    qDet = detectionNetwork.out.createOutputQueue()
17
18    pipeline.start()
19
20    frame = None
21    detections = []
22    startTime = time.monotonic()
23    counter = 0
24    color2 = (255, 255, 255)
25
26    # nn data, being the bounding box locations, are in <0..1> range - they need to be normalized with frame width/height
27    def frameNorm(frame, bbox):
28        normVals = np.full(len(bbox), frame.shape[0])
29        normVals[::2] = frame.shape[1]
30        return (np.clip(np.array(bbox), 0, 1) * normVals).astype(int)
31
32    def displayFrame(name, frame):
33        color = (255, 0, 0)
34        for detection in detections:
35            bbox = frameNorm(
36                frame,
37                (detection.xmin, detection.ymin, detection.xmax, detection.ymax),
38            )
39            cv2.putText(
40                frame,
41                labelMap[detection.label],
42                (bbox[0] + 10, bbox[1] + 20),
43                cv2.FONT_HERSHEY_TRIPLEX,
44                0.5,
45                255,
46            )
47            cv2.putText(
48                frame,
49                f"{int(detection.confidence * 100)}%",
50                (bbox[0] + 10, bbox[1] + 40),
51                cv2.FONT_HERSHEY_TRIPLEX,
52                0.5,
53                255,
54            )
55            cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
56        # Show the frame
57        cv2.imshow(name, frame)
58
59    while pipeline.isRunning():
60        inRgb: dai.ImgFrame = qRgb.get()
61        inDet: dai.ImgDetections = qDet.get()
62        if inRgb is not None:
63            frame = inRgb.getCvFrame()
64            cv2.putText(
65                frame,
66                "NN fps: {:.2f}".format(counter / (time.monotonic() - startTime)),
67                (2, frame.shape[0] - 4),
68                cv2.FONT_HERSHEY_TRIPLEX,
69                0.4,
70                color2,
71            )
72
73        if inDet is not None:
74            detections = inDet.detections
75            counter += 1
76
77        if frame is not None:
78            displayFrame("rgb", frame)
79            print("FPS: {:.2f}".format(counter / (time.monotonic() - startTime)))
80        if cv2.waitKey(1) == ord("q"):
81            pipeline.stop()
82            break