Detection network Remap

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import numpy as np
6
7def colorizeDepth(frameDepth):
8    invalidMask = frameDepth == 0
9    # Log the depth, minDepth and maxDepth
10    try:
11        minDepth = np.percentile(frameDepth[frameDepth != 0], 3)
12        maxDepth = np.percentile(frameDepth[frameDepth != 0], 95)
13        logDepth = np.log(frameDepth, where=frameDepth != 0)
14        logMinDepth = np.log(minDepth)
15        logMaxDepth = np.log(maxDepth)
16        np.nan_to_num(logDepth, copy=False, nan=logMinDepth)
17        # Clip the values to be in the 0-255 range
18        logDepth = np.clip(logDepth, logMinDepth, logMaxDepth)
19
20        # Interpolate only valid logDepth values, setting the rest based on the mask
21        depthFrameColor = np.interp(logDepth, (logMinDepth, logMaxDepth), (0, 255))
22        depthFrameColor = np.nan_to_num(depthFrameColor)
23        depthFrameColor = depthFrameColor.astype(np.uint8)
24        depthFrameColor = cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_JET)
25        # Set invalid depth pixels to black
26        depthFrameColor[invalidMask] = 0
27    except IndexError:
28        # Frame is likely empty
29        depthFrameColor = np.zeros((frameDepth.shape[0], frameDepth.shape[1], 3), dtype=np.uint8)
30    except Exception as e:
31        raise e
32    return depthFrameColor
33
34# Create pipeline
35with dai.Pipeline() as pipeline:
36    cameraNode = pipeline.create(dai.node.Camera).build()
37    detectionNetwork = pipeline.create(dai.node.DetectionNetwork).build(cameraNode, dai.NNModelDescription("yolov6-nano"))
38    labelMap = detectionNetwork.getClasses()
39    monoLeft = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_B)
40    monoRight = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_C)
41    stereo = pipeline.create(dai.node.StereoDepth)
42
43    # Linking
44    monoLeftOut = monoLeft.requestOutput((1280, 720))
45    monoRightOut = monoRight.requestOutput((1280, 720))
46    monoLeftOut.link(stereo.left)
47    monoRightOut.link(stereo.right)
48
49    stereo.setRectification(True)
50    stereo.setExtendedDisparity(True)
51    stereo.setLeftRightCheck(True)
52    stereo.setSubpixel(True)
53
54    qRgb = detectionNetwork.passthrough.createOutputQueue()
55    qDet = detectionNetwork.out.createOutputQueue()
56    qDepth = stereo.disparity.createOutputQueue()
57
58    pipeline.start()
59
60    def displayFrame(name: str, frame: dai.ImgFrame, imgDetections: dai.ImgDetections):
61        color = (0, 255, 0)
62        assert imgDetections.getTransformation() is not None
63        cvFrame = frame.getFrame() if frame.getType() == dai.ImgFrame.Type.RAW16 else frame.getCvFrame()
64        if(frame.getType() == dai.ImgFrame.Type.RAW16):
65            cvFrame = colorizeDepth(cvFrame)
66        for detection in imgDetections.detections:
67            # Get the shape of the frame from which the detections originated for denormalization
68            normShape = imgDetections.getTransformation().getSize()
69
70            # Create rotated rectangle to remap
71            # Here we use an intermediate dai.Rect to create a dai.RotatedRect to simplify construction and denormalization
72            rotRect = dai.RotatedRect(dai.Rect(dai.Point2f(detection.xmin, detection.ymin), dai.Point2f(detection.xmax, detection.ymax)).denormalize(normShape[0], normShape[1]), 0)
73            # Remap the detection rectangle to target frame
74            remapped = imgDetections.getTransformation().remapRectTo(frame.getTransformation(), rotRect)
75            # Remapped rectangle could be rotated, so we get the bounding box
76            bbox = [int(l) for l in remapped.getOuterRect()]
77            cv2.putText(
78                cvFrame,
79                labelMap[detection.label],
80                (bbox[0] + 10, bbox[1] + 20),
81                cv2.FONT_HERSHEY_TRIPLEX,
82                0.5,
83                255,
84            )
85            cv2.putText(
86                cvFrame,
87                f"{int(detection.confidence * 100)}%",
88                (bbox[0] + 10, bbox[1] + 40),
89                cv2.FONT_HERSHEY_TRIPLEX,
90                0.5,
91                255,
92            )
93            cv2.rectangle(cvFrame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
94        # Show the frame
95        cv2.imshow(name, cvFrame)
96
97    while pipeline.isRunning():
98        inRgb: dai.ImgFrame = qRgb.get()
99        inDet: dai.ImgDetections = qDet.get()
100        inDepth: dai.ImgFrame = qDepth.get()
101        hasRgb = inRgb is not None
102        hasDepth = inDepth is not None
103        hasDet = inDet is not None
104        if hasRgb:
105            displayFrame("rgb", inRgb, inDet)
106        if hasDepth:
107            displayFrame("depth", inDepth, inDet)
108        if cv2.waitKey(1) == ord("q"):
109            pipeline.stop()
110            break