Object Tracker 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    objectTracker = pipeline.create(dai.node.ObjectTracker)
39    labelMap = detectionNetwork.getClasses()
40    monoLeft = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_B)
41    monoRight = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_C)
42    stereo = pipeline.create(dai.node.StereoDepth)
43
44    # Linking
45    monoLeftOut = monoLeft.requestOutput((1280, 720), type=dai.ImgFrame.Type.NV12)
46    monoRightOut = monoRight.requestOutput((1280, 720), type=dai.ImgFrame.Type.NV12)
47    monoLeftOut.link(stereo.left)
48    monoRightOut.link(stereo.right)
49
50    detectionNetwork.out.link(objectTracker.inputDetections)
51    detectionNetwork.passthrough.link(objectTracker.inputDetectionFrame)
52    detectionNetwork.passthrough.link(objectTracker.inputTrackerFrame)
53
54    stereo.setRectification(True)
55    stereo.setExtendedDisparity(True)
56    stereo.setLeftRightCheck(True)
57    stereo.setSubpixel(True)
58
59
60    qRgb = detectionNetwork.passthrough.createOutputQueue()
61    qTrack = objectTracker.out.createOutputQueue()
62    qDepth = stereo.disparity.createOutputQueue()
63
64    pipeline.start()
65
66    def displayFrame(name: str, frame: dai.ImgFrame, tracklets: dai.Tracklets):
67        color = (0, 255, 0)
68        assert tracklets.getTransformation() is not None
69        cvFrame = frame.getFrame() if frame.getType() == dai.ImgFrame.Type.RAW16 else frame.getCvFrame()
70        if(frame.getType() == dai.ImgFrame.Type.RAW16):
71            cvFrame = colorizeDepth(cvFrame)
72        for tracklet in tracklets.tracklets:
73            # Get the shape of the frame from which the detections originated for denormalization
74            normShape = tracklets.getTransformation().getSize()
75
76            # Create rotated rectangle to remap
77            # Here we use an intermediate dai.Rect to create a dai.RotatedRect to simplify construction and denormalization
78            rotRect = dai.RotatedRect(tracklet.roi.denormalize(normShape[0], normShape[1]), 0)
79            # Remap the detection rectangle to target frame
80            remapped = tracklets.getTransformation().remapRectTo(frame.getTransformation(), rotRect)
81            # Remapped rectangle could be rotated, so we get the bounding box
82            bbox = [int(l) for l in remapped.getOuterRect()]
83            cv2.putText(
84                cvFrame,
85                labelMap[tracklet.label],
86                (bbox[0] + 10, bbox[1] + 20),
87                cv2.FONT_HERSHEY_TRIPLEX,
88                0.5,
89                255,
90            )
91            cv2.putText(
92                cvFrame,
93                f"{int(tracklet.srcImgDetection.confidence * 100)}%",
94                (bbox[0] + 10, bbox[1] + 40),
95                cv2.FONT_HERSHEY_TRIPLEX,
96                0.5,
97                255,
98            )
99            cv2.rectangle(cvFrame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
100        # Show the frame
101        cv2.imshow(name, cvFrame)
102
103    while pipeline.isRunning():
104        inRgb: dai.ImgFrame = qRgb.get()
105        inTrack: dai.Tracklets = qTrack.get()
106        inDepth: dai.ImgFrame = qDepth.get()
107        hasRgb = inRgb is not None
108        hasDepth = inDepth is not None
109        hasTrack = inTrack is not None
110        if hasRgb:
111            displayFrame("rgb", inRgb, inTrack)
112        if hasDepth:
113            displayFrame("depth", inDepth, inTrack)
114        if cv2.waitKey(1) == ord("q"):
115            pipeline.stop()
116            break