Spatial location calculator

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import numpy as np
6stepSize = 0.05
7
8newConfig = False
9
10# Create pipeline
11pipeline = dai.Pipeline()
12
13# Define sources and outputs
14monoLeft = pipeline.create(dai.node.MonoCamera)
15monoRight = pipeline.create(dai.node.MonoCamera)
16stereo = pipeline.create(dai.node.StereoDepth)
17spatialLocationCalculator = pipeline.create(dai.node.SpatialLocationCalculator)
18
19xoutDepth = pipeline.create(dai.node.XLinkOut)
20xoutSpatialData = pipeline.create(dai.node.XLinkOut)
21xinSpatialCalcConfig = pipeline.create(dai.node.XLinkIn)
22
23xoutDepth.setStreamName("depth")
24xoutSpatialData.setStreamName("spatialData")
25xinSpatialCalcConfig.setStreamName("spatialCalcConfig")
26
27# Properties
28monoLeft.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
29monoLeft.setCamera("left")
30monoRight.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
31monoRight.setCamera("right")
32
33stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
34stereo.setLeftRightCheck(True)
35stereo.setSubpixel(True)
36
37# Config
38topLeft = dai.Point2f(0.4, 0.4)
39bottomRight = dai.Point2f(0.6, 0.6)
40
41config = dai.SpatialLocationCalculatorConfigData()
42config.depthThresholds.lowerThreshold = 100
43config.depthThresholds.upperThreshold = 10000
44calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEDIAN
45config.roi = dai.Rect(topLeft, bottomRight)
46
47spatialLocationCalculator.inputConfig.setWaitForMessage(False)
48spatialLocationCalculator.initialConfig.addROI(config)
49
50# Linking
51monoLeft.out.link(stereo.left)
52monoRight.out.link(stereo.right)
53
54spatialLocationCalculator.passthroughDepth.link(xoutDepth.input)
55stereo.depth.link(spatialLocationCalculator.inputDepth)
56
57spatialLocationCalculator.out.link(xoutSpatialData.input)
58xinSpatialCalcConfig.out.link(spatialLocationCalculator.inputConfig)
59
60# Connect to device and start pipeline
61with dai.Device(pipeline) as device:
62
63    # Output queue will be used to get the depth frames from the outputs defined above
64    depthQueue = device.getOutputQueue(name="depth", maxSize=4, blocking=False)
65    spatialCalcQueue = device.getOutputQueue(name="spatialData", maxSize=4, blocking=False)
66    spatialCalcConfigInQueue = device.getInputQueue("spatialCalcConfig")
67
68    color = (255, 255, 255)
69
70    print("Use WASD keys to move ROI!")
71
72    while True:
73        inDepth = depthQueue.get() # Blocking call, will wait until a new data has arrived
74
75        depthFrame = inDepth.getFrame() # depthFrame values are in millimeters
76
77        depth_downscaled = depthFrame[::4]
78        if np.all(depth_downscaled == 0):
79            min_depth = 0  # Set a default minimum depth value when all elements are zero
80        else:
81            min_depth = np.percentile(depth_downscaled[depth_downscaled != 0], 1)
82        max_depth = np.percentile(depth_downscaled, 99)
83        depthFrameColor = np.interp(depthFrame, (min_depth, max_depth), (0, 255)).astype(np.uint8)
84        depthFrameColor = cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_HOT)
85
86        spatialData = spatialCalcQueue.get().getSpatialLocations()
87        for depthData in spatialData:
88            roi = depthData.config.roi
89            roi = roi.denormalize(width=depthFrameColor.shape[1], height=depthFrameColor.shape[0])
90            xmin = int(roi.topLeft().x)
91            ymin = int(roi.topLeft().y)
92            xmax = int(roi.bottomRight().x)
93            ymax = int(roi.bottomRight().y)
94
95            depthMin = depthData.depthMin
96            depthMax = depthData.depthMax
97
98            fontType = cv2.FONT_HERSHEY_TRIPLEX
99            cv2.rectangle(depthFrameColor, (xmin, ymin), (xmax, ymax), color, 1)
100            cv2.putText(depthFrameColor, f"X: {int(depthData.spatialCoordinates.x)} mm", (xmin + 10, ymin + 20), fontType, 0.5, color)
101            cv2.putText(depthFrameColor, f"Y: {int(depthData.spatialCoordinates.y)} mm", (xmin + 10, ymin + 35), fontType, 0.5, color)
102            cv2.putText(depthFrameColor, f"Z: {int(depthData.spatialCoordinates.z)} mm", (xmin + 10, ymin + 50), fontType, 0.5, color)
103        # Show the frame
104        cv2.imshow("depth", depthFrameColor)
105
106        key = cv2.waitKey(1)
107        if key == ord('q'):
108            break
109        elif key == ord('w'):
110            if topLeft.y - stepSize >= 0:
111                topLeft.y -= stepSize
112                bottomRight.y -= stepSize
113                newConfig = True
114        elif key == ord('a'):
115            if topLeft.x - stepSize >= 0:
116                topLeft.x -= stepSize
117                bottomRight.x -= stepSize
118                newConfig = True
119        elif key == ord('s'):
120            if bottomRight.y + stepSize <= 1:
121                topLeft.y += stepSize
122                bottomRight.y += stepSize
123                newConfig = True
124        elif key == ord('d'):
125            if bottomRight.x + stepSize <= 1:
126                topLeft.x += stepSize
127                bottomRight.x += stepSize
128                newConfig = True
129        elif key == ord('1'):
130            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEAN
131            print('Switching calculation algorithm to MEAN!')
132            newConfig = True
133        elif key == ord('2'):
134            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN
135            print('Switching calculation algorithm to MIN!')
136            newConfig = True
137        elif key == ord('3'):
138            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MAX
139            print('Switching calculation algorithm to MAX!')
140            newConfig = True
141        elif key == ord('4'):
142            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MODE
143            print('Switching calculation algorithm to MODE!')
144            newConfig = True
145        elif key == ord('5'):
146            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEDIAN
147            print('Switching calculation algorithm to MEDIAN!')
148            newConfig = True
149
150        if newConfig:
151            config.roi = dai.Rect(topLeft, bottomRight)
152            config.calculationAlgorithm = calculationAlgorithm
153            cfg = dai.SpatialLocationCalculatorConfig()
154            cfg.addROI(config)
155            spatialCalcConfigInQueue.send(cfg)
156            newConfig = False