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

1#include <atomic>
2#include <iostream>
3#include <memory>
4#include <opencv2/opencv.hpp>
5
6#include "depthai/depthai.hpp"
7
8int main() {
9    try {
10        // Create pipeline
11        dai::Pipeline pipeline;
12
13        // Define sources and outputs
14        auto monoLeft = pipeline.create<dai::node::Camera>();
15        monoLeft->build(dai::CameraBoardSocket::CAM_B);
16
17        auto monoRight = pipeline.create<dai::node::Camera>();
18        monoRight->build(dai::CameraBoardSocket::CAM_C);
19
20        auto stereo = pipeline.create<dai::node::StereoDepth>();
21        auto spatialLocationCalculator = pipeline.create<dai::node::SpatialLocationCalculator>();
22
23        // Configure stereo
24        stereo->setRectification(true);
25        stereo->setExtendedDisparity(true);
26
27        // Initial ROI configuration
28        dai::Point2f topLeft(0.4f, 0.4f);
29        dai::Point2f bottomRight(0.6f, 0.6f);
30        float stepSize = 0.05f;
31
32        // Configure spatial location calculator
33        dai::SpatialLocationCalculatorConfigData config;
34        config.depthThresholds.lowerThreshold = 10;
35        config.depthThresholds.upperThreshold = 10000;
36        auto calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MEDIAN;
37        config.roi = dai::Rect(topLeft, bottomRight);
38
39        spatialLocationCalculator->inputConfig.setWaitForMessage(false);
40        spatialLocationCalculator->initialConfig->addROI(config);
41
42        // Create output queues
43        auto xoutSpatialQueue = spatialLocationCalculator->out.createOutputQueue();
44        auto outputDepthQueue = spatialLocationCalculator->passthroughDepth.createOutputQueue();
45        auto inputConfigQueue = spatialLocationCalculator->inputConfig.createInputQueue();
46
47        // Linking
48        monoLeft->requestOutput(std::make_pair(640, 400))->link(stereo->left);
49        monoRight->requestOutput(std::make_pair(640, 400))->link(stereo->right);
50        stereo->depth.link(spatialLocationCalculator->inputDepth);
51
52        // Start pipeline
53        pipeline.start();
54
55        cv::Scalar color(255, 255, 255);
56
57        while(pipeline.isRunning()) {
58            auto spatialData = xoutSpatialQueue->get<dai::SpatialLocationCalculatorData>();
59            std::cout << "Use WASD keys to move ROI!" << std::endl;
60
61            auto outputDepthImage = outputDepthQueue->get<dai::ImgFrame>();
62            cv::Mat frameDepth = outputDepthImage->getCvFrame();
63
64            // Calculate median depth
65            std::vector<float> depthValues;
66            for(int i = 0; i < frameDepth.rows; i++) {
67                for(int j = 0; j < frameDepth.cols; j++) {
68                    uint16_t val = frameDepth.at<uint16_t>(i, j);
69                    if(val > 0) depthValues.push_back(val);
70                }
71            }
72            std::sort(depthValues.begin(), depthValues.end());
73            float medianDepth = depthValues[depthValues.size() / 2];
74            std::cout << "Median depth value: " << medianDepth << std::endl;
75
76            // Process depth frame for visualization
77            cv::Mat depthFrameColor;
78            cv::normalize(frameDepth, depthFrameColor, 255, 0, cv::NORM_INF, CV_8UC1);
79            cv::equalizeHist(depthFrameColor, depthFrameColor);
80            cv::applyColorMap(depthFrameColor, depthFrameColor, cv::COLORMAP_HOT);
81
82            // Draw spatial data
83            for(const auto& depthData : spatialData->spatialLocations) {
84                auto roi = depthData.config.roi;
85                roi = roi.denormalize(depthFrameColor.cols, depthFrameColor.rows);
86                int xmin = static_cast<int>(roi.topLeft().x);
87                int ymin = static_cast<int>(roi.topLeft().y);
88                int xmax = static_cast<int>(roi.bottomRight().x);
89                int ymax = static_cast<int>(roi.bottomRight().y);
90
91                float depthMin = depthData.depthMin;
92                float depthMax = depthData.depthMax;
93
94                cv::rectangle(depthFrameColor, cv::Point(xmin, ymin), cv::Point(xmax, ymax), color, 1);
95                cv::putText(depthFrameColor,
96                            "X: " + std::to_string(static_cast<int>(depthData.spatialCoordinates.x)) + " mm",
97                            cv::Point(xmin + 10, ymin + 20),
98                            cv::FONT_HERSHEY_TRIPLEX,
99                            0.5,
100                            color);
101                cv::putText(depthFrameColor,
102                            "Y: " + std::to_string(static_cast<int>(depthData.spatialCoordinates.y)) + " mm",
103                            cv::Point(xmin + 10, ymin + 35),
104                            cv::FONT_HERSHEY_TRIPLEX,
105                            0.5,
106                            color);
107                cv::putText(depthFrameColor,
108                            "Z: " + std::to_string(static_cast<int>(depthData.spatialCoordinates.z)) + " mm",
109                            cv::Point(xmin + 10, ymin + 50),
110                            cv::FONT_HERSHEY_TRIPLEX,
111                            0.5,
112                            color);
113            }
114
115            cv::imshow("depth", depthFrameColor);
116
117            int key = cv::waitKey(1);
118            bool newConfig = false;
119
120            if(key == 'q') {
121                break;
122            } else if(key == 'w') {
123                if(topLeft.y - stepSize >= 0) {
124                    topLeft.y -= stepSize;
125                    bottomRight.y -= stepSize;
126                    newConfig = true;
127                }
128            } else if(key == 'a') {
129                if(topLeft.x - stepSize >= 0) {
130                    topLeft.x -= stepSize;
131                    bottomRight.x -= stepSize;
132                    newConfig = true;
133                }
134            } else if(key == 's') {
135                if(bottomRight.y + stepSize <= 1) {
136                    topLeft.y += stepSize;
137                    bottomRight.y += stepSize;
138                    newConfig = true;
139                }
140            } else if(key == 'd') {
141                if(bottomRight.x + stepSize <= 1) {
142                    topLeft.x += stepSize;
143                    bottomRight.x += stepSize;
144                    newConfig = true;
145                }
146            } else if(key == '1') {
147                calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MEAN;
148                std::cout << "Switching calculation algorithm to MEAN!" << std::endl;
149                newConfig = true;
150            } else if(key == '2') {
151                calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MIN;
152                std::cout << "Switching calculation algorithm to MIN!" << std::endl;
153                newConfig = true;
154            } else if(key == '3') {
155                calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MAX;
156                std::cout << "Switching calculation algorithm to MAX!" << std::endl;
157                newConfig = true;
158            } else if(key == '4') {
159                calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MODE;
160                std::cout << "Switching calculation algorithm to MODE!" << std::endl;
161                newConfig = true;
162            } else if(key == '5') {
163                calculationAlgorithm = dai::SpatialLocationCalculatorAlgorithm::MEDIAN;
164                std::cout << "Switching calculation algorithm to MEDIAN!" << std::endl;
165                newConfig = true;
166            }
167
168            if(newConfig) {
169                config.roi = dai::Rect(topLeft, bottomRight);
170                config.calculationAlgorithm = calculationAlgorithm;
171                std::shared_ptr<dai::SpatialLocationCalculatorConfig> cfg = std::make_shared<dai::SpatialLocationCalculatorConfig>();
172                cfg->addROI(config);
173                inputConfigQueue->send(cfg);
174            }
175        }
176
177        // Cleanup
178        pipeline.stop();
179        pipeline.wait();
180
181    } catch(const std::exception& e) {
182        std::cerr << "Error: " << e.what() << std::endl;
183        return 1;
184    }
185
186    return 0;
187}