Spatial Location Calculator

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import numpy as np
6from pathlib import Path
7
8color = (255, 255, 255)
9
10# Create pipeline
11pipeline = dai.Pipeline()
12# Config
13topLeft = dai.Point2f(0.4, 0.4)
14bottomRight = dai.Point2f(0.6, 0.6)
15
16# Define sources and outputs
17monoLeft = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_B)
18monoRight = pipeline.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_C)
19stereo = pipeline.create(dai.node.StereoDepth)
20spatialLocationCalculator = pipeline.create(dai.node.SpatialLocationCalculator)
21
22# Linking
23monoLeftOut = monoLeft.requestOutput((640, 400))
24monoRightOut = monoRight.requestOutput((640, 400))
25monoLeftOut.link(stereo.left)
26monoRightOut.link(stereo.right)
27
28stereo.setRectification(True)
29stereo.setExtendedDisparity(True)
30
31stepSize = 0.05
32
33config = dai.SpatialLocationCalculatorConfigData()
34config.depthThresholds.lowerThreshold = 10
35config.depthThresholds.upperThreshold = 10000
36calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEDIAN
37config.roi = dai.Rect(topLeft, bottomRight)
38
39spatialLocationCalculator.inputConfig.setWaitForMessage(False)
40spatialLocationCalculator.initialConfig.addROI(config)
41
42
43xoutSpatialQueue = spatialLocationCalculator.out.createOutputQueue()
44outputDepthQueue = spatialLocationCalculator.passthroughDepth.createOutputQueue()
45
46stereo.depth.link(spatialLocationCalculator.inputDepth)
47
48
49inputConfigQueue = spatialLocationCalculator.inputConfig.createInputQueue()
50
51with pipeline:
52    pipeline.start()
53    while pipeline.isRunning():
54        spatialData = xoutSpatialQueue.get().getSpatialLocations()
55
56        print("Use WASD keys to move ROI!")
57        outputDepthIMage : dai.ImgFrame = outputDepthQueue.get()
58
59        frameDepth = outputDepthIMage.getCvFrame()
60        frameDepth = outputDepthIMage.getFrame()
61        print("Median depth value: ", np.median(frameDepth))
62
63        depthFrameColor = cv2.normalize(frameDepth, None, 255, 0, cv2.NORM_INF, cv2.CV_8UC1)
64        depthFrameColor = cv2.equalizeHist(depthFrameColor)
65        depthFrameColor = cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_HOT)
66        for depthData in spatialData:
67            roi = depthData.config.roi
68            roi = roi.denormalize(width=depthFrameColor.shape[1], height=depthFrameColor.shape[0])
69            xmin = int(roi.topLeft().x)
70            ymin = int(roi.topLeft().y)
71            xmax = int(roi.bottomRight().x)
72            ymax = int(roi.bottomRight().y)
73
74            depthMin = depthData.depthMin
75            depthMax = depthData.depthMax
76
77            fontType = cv2.FONT_HERSHEY_TRIPLEX
78            cv2.rectangle(depthFrameColor, (xmin, ymin), (xmax, ymax), color, cv2.FONT_HERSHEY_SCRIPT_SIMPLEX)
79            cv2.putText(depthFrameColor, f"X: {int(depthData.spatialCoordinates.x)} mm", (xmin + 10, ymin + 20), fontType, 0.5, color)
80            cv2.putText(depthFrameColor, f"Y: {int(depthData.spatialCoordinates.y)} mm", (xmin + 10, ymin + 35), fontType, 0.5, color)
81            cv2.putText(depthFrameColor, f"Z: {int(depthData.spatialCoordinates.z)} mm", (xmin + 10, ymin + 50), fontType, 0.5, color)
82        # Show the frame
83        cv2.imshow("depth", depthFrameColor)
84
85        key = cv2.waitKey(1)
86        if key == ord('q'):
87            pipeline.stop()
88            break
89
90        stepSize = 0.05
91
92        newConfig = False
93
94        if key == ord('q'):
95            break
96        elif key == ord('w'):
97            if topLeft.y - stepSize >= 0:
98                topLeft.y -= stepSize
99                bottomRight.y -= stepSize
100                newConfig = True
101        elif key == ord('a'):
102            if topLeft.x - stepSize >= 0:
103                topLeft.x -= stepSize
104                bottomRight.x -= stepSize
105                newConfig = True
106        elif key == ord('s'):
107            if bottomRight.y + stepSize <= 1:
108                topLeft.y += stepSize
109                bottomRight.y += stepSize
110                newConfig = True
111        elif key == ord('d'):
112            if bottomRight.x + stepSize <= 1:
113                topLeft.x += stepSize
114                bottomRight.x += stepSize
115                newConfig = True
116        elif key == ord('1'):
117            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEAN
118            print('Switching calculation algorithm to MEAN!')
119            newConfig = True
120        elif key == ord('2'):
121            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MIN
122            print('Switching calculation algorithm to MIN!')
123            newConfig = True
124        elif key == ord('3'):
125            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MAX
126            print('Switching calculation algorithm to MAX!')
127            newConfig = True
128        elif key == ord('4'):
129            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MODE
130            print('Switching calculation algorithm to MODE!')
131            newConfig = True
132        elif key == ord('5'):
133            calculationAlgorithm = dai.SpatialLocationCalculatorAlgorithm.MEDIAN
134            print('Switching calculation algorithm to MEDIAN!')
135            newConfig = True
136
137        if newConfig:
138            config.roi = dai.Rect(topLeft, bottomRight)
139            config.calculationAlgorithm = calculationAlgorithm
140            cfg = dai.SpatialLocationCalculatorConfig()
141            cfg.addROI(config)
142            inputConfigQueue.send(cfg)
143            newConfig = False