空间检测网络

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5
6fps = 20
7modelDescription = dai.NNModelDescription("yolov6-nano")
8
9class SpatialVisualizer(dai.node.HostNode):
10    def __init__(self):
11        dai.node.HostNode.__init__(self)
12        self.sendProcessingToPipeline(True)
13    def build(self, depth:dai.Node.Output, detections: dai.Node.Output, rgb: dai.Node.Output):
14        self.link_args(depth, detections, rgb) # Must match the inputs to the process method
15
16    def process(self, depthPreview, detections, rgbPreview):
17        rgbPreview = rgbPreview.getCvFrame()
18        depthFrameColor = self.processDepthFrame(depthPreview)
19        self.displayResults(rgbPreview, depthFrameColor, detections.detections)
20
21    def processDepthFrame(self, depthFrame):
22        return dai.utility.colorizeDepthFrame(depthFrame, colormap=cv2.COLORMAP_HOT).getCvFrame()
23
24    def displayResults(self, rgbFrame, depthFrameColor, detections):
25        height, width, _ = rgbFrame.shape
26        for detection in detections:
27            self.drawBoundingBoxes(depthFrameColor, detection)
28            self.drawDetections(rgbFrame, detection, width, height)
29
30        cv2.imshow("Depth frame", depthFrameColor)
31        cv2.imshow("Color frame", rgbFrame)
32        if cv2.waitKey(1) == ord('q'):
33            self.stopPipeline()
34
35    def drawBoundingBoxes(self, depthFrameColor, detection):
36        roiData = detection.boundingBoxMapping
37        roi = roiData.roi
38        roi = roi.denormalize(depthFrameColor.shape[1], depthFrameColor.shape[0])
39        topLeft = roi.topLeft()
40        bottomRight = roi.bottomRight()
41        cv2.rectangle(depthFrameColor, (int(topLeft.x), int(topLeft.y)), (int(bottomRight.x), int(bottomRight.y)), (255, 255, 255), 1)
42
43    def drawDetections(self, frame, detection, frameWidth, frameHeight):
44        x1 = int(detection.xmin * frameWidth)
45        x2 = int(detection.xmax * frameWidth)
46        y1 = int(detection.ymin * frameHeight)
47        y2 = int(detection.ymax * frameHeight)
48        label = detection.labelName
49        color = (255, 255, 255)
50        cv2.putText(frame, str(label), (x1 + 10, y1 + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
51        cv2.putText(frame, "{:.2f}".format(detection.confidence * 100), (x1 + 10, y1 + 35), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
52        cv2.putText(frame, f"X: {int(detection.spatialCoordinates.x)} mm", (x1 + 10, y1 + 50), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
53        cv2.putText(frame, f"Y: {int(detection.spatialCoordinates.y)} mm", (x1 + 10, y1 + 65), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
54        cv2.putText(frame, f"Z: {int(detection.spatialCoordinates.z)} mm", (x1 + 10, y1 + 80), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
55        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 1)
56
57# Creates the pipeline and a default device implicitly
58with dai.Pipeline() as p:
59    # Define sources and outputs
60    colorSockets = p.getDefaultDevice().getConnectedCameras(dai.CameraSensorType.COLOR)
61    colorSocket = colorSockets[0] if colorSockets else dai.CameraBoardSocket.CAM_A
62    camRgb = p.create(dai.node.Camera).build(colorSocket, sensorFps=fps)
63    depth = p.create(dai.node.Depth).build(dai.node.Depth.Algorithm.AUTO, fps, (640, 400))
64
65    spatialDetectionNetwork = p.create(dai.node.SpatialDetectionNetwork).build(
66        camRgb, depth, modelDescription
67    )
68    visualizer = p.create(SpatialVisualizer)
69
70    spatialDetectionNetwork.spatialLocationCalculator.initialConfig.setSegmentationPassthrough(False)
71    spatialDetectionNetwork.input.setBlocking(False)
72    spatialDetectionNetwork.setDepthLowerThreshold(100)
73    spatialDetectionNetwork.setDepthUpperThreshold(5000)
74
75    visualizer.build(
76        spatialDetectionNetwork.passthroughDepth,
77        spatialDetectionNetwork.out,
78        spatialDetectionNetwork.passthrough,
79    )
80
81    print("Starting pipeline")
82
83    p.run()

1#include <argparse/argparse.hpp>
2#include <csignal>
3#include <iostream>
4#include <memory>
5#include <opencv2/opencv.hpp>
6#include <vector>
7
8#include "depthai/depthai.hpp"
9
10constexpr float NEURAL_FPS = 8.0f;
11constexpr float STEREO_DEFAULT_FPS = 20.0f;
12
13std::atomic<bool> quitEvent(false);
14
15void signalHandler(int) {
16    quitEvent = true;
17}
18
19// Custom host node for spatial visualization
20class SpatialVisualizer : public dai::NodeCRTP<dai::node::HostNode, SpatialVisualizer> {
21   public:
22    Input& depthInput = inputs["depth"];
23    Input& detectionsInput = inputs["detections"];
24    Input& rgbInput = inputs["rgb"];
25
26    std::vector<std::string> labelMap;
27
28    std::shared_ptr<SpatialVisualizer> build(Output& depth, Output& detections, Output& rgb) {
29        depth.link(depthInput);
30        detections.link(detectionsInput);
31        rgb.link(rgbInput);
32        sendProcessingToPipeline(true);
33        return std::static_pointer_cast<SpatialVisualizer>(this->shared_from_this());
34    }
35
36    std::shared_ptr<dai::Buffer> processGroup(std::shared_ptr<dai::MessageGroup> in) override {
37        if(quitEvent) {
38            stopPipeline();
39            return nullptr;
40        }
41
42        auto depthFrame = in->get<dai::ImgFrame>("depth");
43        auto detections = in->get<dai::SpatialImgDetections>("detections");
44        auto rgbFrame = in->get<dai::ImgFrame>("rgb");
45
46        cv::Mat rgbCv = rgbFrame->getCvFrame();
47        cv::Mat depthFrameColor = processDepthFrame(*depthFrame);
48        displayResults(rgbCv, depthFrameColor, detections->detections);
49
50        return nullptr;
51    }
52
53   private:
54    cv::Mat processDepthFrame(const dai::ImgFrame& depthFrameImg) {
55        return dai::utility::colorizeDepthFrame(depthFrameImg, 500.0f, 12000.0f, cv::COLORMAP_HOT, true).getCvFrame();
56    }
57
58    void displayResults(cv::Mat& rgbFrame, cv::Mat& depthFrameColor, const std::vector<dai::SpatialImgDetection>& detections) {
59        int height = rgbFrame.rows;
60        int width = rgbFrame.cols;
61
62        for(const auto& detection : detections) {
63            drawBoundingBoxes(depthFrameColor, detection);
64            drawDetections(rgbFrame, detection, width, height);
65        }
66
67        cv::imshow("depth", depthFrameColor);
68        cv::imshow("rgb", rgbFrame);
69
70        if(cv::waitKey(1) == 'q') {
71            stopPipeline();
72        }
73    }
74
75    void drawBoundingBoxes(cv::Mat& depthFrameColor, const dai::SpatialImgDetection& detection) {
76        auto roi = detection.boundingBoxMapping.roi;
77        roi = roi.denormalize(depthFrameColor.cols, depthFrameColor.rows);
78        auto topLeft = roi.topLeft();
79        auto bottomRight = roi.bottomRight();
80        cv::rectangle(depthFrameColor,
81                      cv::Point(static_cast<int>(topLeft.x), static_cast<int>(topLeft.y)),
82                      cv::Point(static_cast<int>(bottomRight.x), static_cast<int>(bottomRight.y)),
83                      cv::Scalar(255, 255, 255),
84                      1);
85    }
86
87    void drawDetections(cv::Mat& frame, const dai::SpatialImgDetection& detection, int frameWidth, int frameHeight) {
88        int x1 = static_cast<int>(detection.xmin * frameWidth);
89        int x2 = static_cast<int>(detection.xmax * frameWidth);
90        int y1 = static_cast<int>(detection.ymin * frameHeight);
91        int y2 = static_cast<int>(detection.ymax * frameHeight);
92
93        std::string label;
94        try {
95            label = labelMap[detection.label];
96        } catch(...) {
97            label = std::to_string(detection.label);
98        }
99
100        cv::Scalar color(255, 255, 255);
101        cv::putText(frame, label, cv::Point(x1 + 10, y1 + 20), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
102        cv::putText(frame, std::to_string(detection.confidence * 100), cv::Point(x1 + 10, y1 + 35), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
103        cv::putText(frame,
104                    "X: " + std::to_string(static_cast<int>(detection.spatialCoordinates.x)) + " mm",
105                    cv::Point(x1 + 10, y1 + 50),
106                    cv::FONT_HERSHEY_TRIPLEX,
107                    0.5,
108                    color);
109        cv::putText(frame,
110                    "Y: " + std::to_string(static_cast<int>(detection.spatialCoordinates.y)) + " mm",
111                    cv::Point(x1 + 10, y1 + 65),
112                    cv::FONT_HERSHEY_TRIPLEX,
113                    0.5,
114                    color);
115        cv::putText(frame,
116                    "Z: " + std::to_string(static_cast<int>(detection.spatialCoordinates.z)) + " mm",
117                    cv::Point(x1 + 10, y1 + 80),
118                    cv::FONT_HERSHEY_TRIPLEX,
119                    0.5,
120                    color);
121        cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, 1);
122    }
123};
124
125int main(int argc, char** argv) {
126    signal(SIGTERM, signalHandler);
127    signal(SIGINT, signalHandler);
128
129    // Initialize argument parser
130    argparse::ArgumentParser program("spatial_detection", "1.0.0");
131    program.add_description("Spatial detection network example with configurable depth source");
132    program.add_argument("--depthSource").default_value(std::string("stereo")).help("Depth source: stereo, neural, tof");
133
134    try {
135        // Parse arguments
136        program.parse_args(argc, argv);
137    } catch(const std::runtime_error& err) {
138        std::cerr << err.what() << '\n';
139        std::cerr << program;
140        return EXIT_FAILURE;
141    }
142
143    // Get arguments
144    std::string depthSourceArg = program.get<std::string>("--depthSource");
145
146    // Validate depth source argument
147    if(depthSourceArg != "stereo" && depthSourceArg != "neural" && depthSourceArg != "tof") {
148        std::cerr << "Invalid depth source: " << depthSourceArg << '\n';
149        std::cerr << "Valid options are: stereo, neural, tof" << '\n';
150        return EXIT_FAILURE;
151    }
152
153    try {
154        float fps = STEREO_DEFAULT_FPS;
155        if(depthSourceArg == "neural") {
156            fps = NEURAL_FPS;
157        }
158
159        // Create pipeline
160        dai::Pipeline pipeline;
161
162        // Define sources and outputs
163        auto colorSockets = pipeline.getDefaultDevice()->getConnectedCameras(dai::CameraSensorType::COLOR);
164        auto colorSocket = colorSockets.empty() ? dai::CameraBoardSocket::CAM_A : colorSockets.front();
165        auto camRgb = pipeline.create<dai::node::Camera>();
166        camRgb->build(colorSocket, std::nullopt, fps);
167
168        // Create depth source based on argument
169        dai::node::DepthSource depthSource;
170
171        if(depthSourceArg == "stereo") {
172            auto depth = pipeline.create<dai::node::Depth>();
173            depth->build(dai::node::Depth::Algorithm::AUTO, fps, std::make_pair(640u, 400u));
174
175            depthSource = depth;
176        } else if(depthSourceArg == "neural") {
177            auto monoLeft = pipeline.create<dai::node::Camera>();
178            auto monoRight = pipeline.create<dai::node::Camera>();
179
180            monoLeft->build(dai::CameraBoardSocket::CAM_B, std::nullopt, fps);
181            monoRight->build(dai::CameraBoardSocket::CAM_C, std::nullopt, fps);
182
183            auto neuralDepth = pipeline.create<dai::node::NeuralDepth>();
184            neuralDepth->build(*monoLeft->requestFullResolutionOutput(), *monoRight->requestFullResolutionOutput(), dai::DeviceModelZoo::NEURAL_DEPTH_LARGE);
185
186            depthSource = neuralDepth;
187        } else if(depthSourceArg == "tof") {
188            auto tof = pipeline.create<dai::node::ToF>();
189            depthSource = tof;
190        }
191
192        // Create spatial detection network using the unified build method with DepthSource variant
193        auto spatialDetectionNetwork = pipeline.create<dai::node::SpatialDetectionNetwork>();
194        auto visualizer = pipeline.create<SpatialVisualizer>();
195
196        // Configure spatial detection network
197        spatialDetectionNetwork->input.setBlocking(false);
198        spatialDetectionNetwork->setBoundingBoxScaleFactor(0.5f);
199        spatialDetectionNetwork->setDepthLowerThreshold(100);
200        spatialDetectionNetwork->setDepthUpperThreshold(5000);
201
202        // Set up model and build with DepthSource variant
203        dai::NNModelDescription modelDesc;
204        // For better results on OAK4, use a segmentation model like "luxonis/yolov8-instance-segmentation-large:coco-640x480"
205        // for depth estimation over the objects mask instead of the full bounding box.
206        modelDesc.model = "yolov6-nano";
207        spatialDetectionNetwork->build(camRgb, depthSource, modelDesc);
208
209        // Set label map
210        visualizer->labelMap = spatialDetectionNetwork->getClasses().value();
211        spatialDetectionNetwork->spatialLocationCalculator->initialConfig->setSegmentationPassthrough(false);
212
213        // Linking
214        visualizer->build(spatialDetectionNetwork->passthroughDepth, spatialDetectionNetwork->out, spatialDetectionNetwork->passthrough);
215
216        std::cout << "Pipeline starting with depth source: " << depthSourceArg << '\n';
217
218        // Start pipeline
219        pipeline.run();
220
221    } catch(const std::exception& e) {
222        std::cerr << "Error: " << e.what() << '\n';
223        return EXIT_FAILURE;
224    }
225
226    return EXIT_SUCCESS;
227}