# Spatial Detection Network

The example creates a pipeline to perform YOLOv6-Nano spatial object detection using RGB and stereo depth streams, visualizes
results with bounding boxes and spatial coordinates on both colorized depth and RGB frames, and uses a custom visualization node.

This example requires the DepthAI v3 API, see [installation instructions](https://docs.luxonis.com/software-v3/depthai.md).

## Pipeline

## Source code

#### Python

```python
#!/usr/bin/env python3

import argparse
from pathlib import Path
import cv2
import depthai as dai
import numpy as np

NEURAL_FPS = 8
STEREO_DEFAULT_FPS = 20

parser = argparse.ArgumentParser()
parser.add_argument(
    "--depthSource", type=str, default="stereo", choices=["stereo", "neural"]
)
args = parser.parse_args()
# For better results on OAK4, use a segmentation model like "luxonis/yolov8-instance-segmentation-large:coco-640x480"
# for depth estimation over the objects mask instead of the full bounding box.
modelDescription = dai.NNModelDescription("yolov6-nano")
size = (640, 400)

if args.depthSource == "stereo":
    fps = STEREO_DEFAULT_FPS
else:
    fps = NEURAL_FPS

class SpatialVisualizer(dai.node.HostNode):
    def __init__(self):
        dai.node.HostNode.__init__(self)
        self.sendProcessingToPipeline(True)
    def build(self, depth:dai.Node.Output, detections: dai.Node.Output, rgb: dai.Node.Output):
        self.link_args(depth, detections, rgb) # Must match the inputs to the process method

    def process(self, depthPreview, detections, rgbPreview):
        depthPreview = depthPreview.getCvFrame()
        rgbPreview = rgbPreview.getCvFrame()
        depthFrameColor = self.processDepthFrame(depthPreview)
        self.displayResults(rgbPreview, depthFrameColor, detections.detections)

    def processDepthFrame(self, depthFrame):
        depthDownscaled = depthFrame[::4]
        if np.all(depthDownscaled == 0):
            minDepth = 0
        else:
            minDepth = np.percentile(depthDownscaled[depthDownscaled != 0], 1)
        maxDepth = np.percentile(depthDownscaled, 99)
        depthFrameColor = np.interp(depthFrame, (minDepth, maxDepth), (0, 255)).astype(np.uint8)
        return cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_HOT)

    def displayResults(self, rgbFrame, depthFrameColor, detections):
        height, width, _ = rgbFrame.shape
        for detection in detections:
            self.drawBoundingBoxes(depthFrameColor, detection)
            self.drawDetections(rgbFrame, detection, width, height)

        cv2.imshow("Depth frame", depthFrameColor)
        cv2.imshow("Color frame", rgbFrame)
        if cv2.waitKey(1) == ord('q'):
            self.stopPipeline()

    def drawBoundingBoxes(self, depthFrameColor, detection):
        roiData = detection.boundingBoxMapping
        roi = roiData.roi
        roi = roi.denormalize(depthFrameColor.shape[1], depthFrameColor.shape[0])
        topLeft = roi.topLeft()
        bottomRight = roi.bottomRight()
        cv2.rectangle(depthFrameColor, (int(topLeft.x), int(topLeft.y)), (int(bottomRight.x), int(bottomRight.y)), (255, 255, 255), 1)

    def drawDetections(self, frame, detection, frameWidth, frameHeight):
        x1 = int(detection.xmin * frameWidth)
        x2 = int(detection.xmax * frameWidth)
        y1 = int(detection.ymin * frameHeight)
        y2 = int(detection.ymax * frameHeight)
        label = detection.labelName
        color = (255, 255, 255)
        cv2.putText(frame, str(label), (x1 + 10, y1 + 20), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
        cv2.putText(frame, "{:.2f}".format(detection.confidence * 100), (x1 + 10, y1 + 35), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
        cv2.putText(frame, f"X: {int(detection.spatialCoordinates.x)} mm", (x1 + 10, y1 + 50), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
        cv2.putText(frame, f"Y: {int(detection.spatialCoordinates.y)} mm", (x1 + 10, y1 + 65), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
        cv2.putText(frame, f"Z: {int(detection.spatialCoordinates.z)} mm", (x1 + 10, y1 + 80), cv2.FONT_HERSHEY_TRIPLEX, 0.5, color)
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 1)

# Creates the pipeline and a default device implicitly
with dai.Pipeline() as p:
    # Define sources and outputs
    platform = p.getDefaultDevice().getPlatform()

    camRgb = p.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_A, sensorFps=fps)
    monoLeft = p.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_B, sensorFps=fps)
    monoRight = p.create(dai.node.Camera).build(dai.CameraBoardSocket.CAM_C, sensorFps=fps)
    if args.depthSource == "stereo":
        depthSource = p.create(dai.node.StereoDepth)
        depthSource.setExtendedDisparity(True)
        monoLeft.requestOutput(size).link(depthSource.left)
        monoRight.requestOutput(size).link(depthSource.right)
    elif args.depthSource == "neural":
        depthSource = p.create(dai.node.NeuralDepth).build(
            monoLeft.requestFullResolutionOutput(),
            monoRight.requestFullResolutionOutput(),
            dai.DeviceModelZoo.NEURAL_DEPTH_LARGE,
        )
    else:
        raise ValueError(f"Invalid depth source: {args.depthSource}")

    spatialDetectionNetwork = p.create(dai.node.SpatialDetectionNetwork).build(
        camRgb, depthSource, modelDescription
    )
    visualizer = p.create(SpatialVisualizer)

    spatialDetectionNetwork.spatialLocationCalculator.initialConfig.setSegmentationPassthrough(False)
    spatialDetectionNetwork.input.setBlocking(False)
    spatialDetectionNetwork.setDepthLowerThreshold(100)
    spatialDetectionNetwork.setDepthUpperThreshold(5000)

    visualizer.build(
        spatialDetectionNetwork.passthroughDepth,
        spatialDetectionNetwork.out,
        spatialDetectionNetwork.passthrough,
    )

    print("Starting pipeline with depth source: ", args.depthSource)

    p.run()
```

#### C++

```cpp
#include <argparse/argparse.hpp>
#include <csignal>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
#include <vector>

#include "depthai/depthai.hpp"

constexpr float NEURAL_FPS = 8.0f;
constexpr float STEREO_DEFAULT_FPS = 20.0f;

std::atomic<bool> quitEvent(false);

void signalHandler(int) {
    quitEvent = true;
}

// Custom host node for spatial visualization
class SpatialVisualizer : public dai::NodeCRTP<dai::node::HostNode, SpatialVisualizer> {
   public:
    Input& depthInput = inputs["depth"];
    Input& detectionsInput = inputs["detections"];
    Input& rgbInput = inputs["rgb"];

    std::vector<std::string> labelMap;

    std::shared_ptr<SpatialVisualizer> build(Output& depth, Output& detections, Output& rgb) {
        depth.link(depthInput);
        detections.link(detectionsInput);
        rgb.link(rgbInput);
        sendProcessingToPipeline(true);
        return std::static_pointer_cast<SpatialVisualizer>(this->shared_from_this());
    }

    std::shared_ptr<dai::Buffer> processGroup(std::shared_ptr<dai::MessageGroup> in) override {
        if(quitEvent) {
            stopPipeline();
            return nullptr;
        }

        auto depthFrame = in->get<dai::ImgFrame>("depth");
        auto detections = in->get<dai::SpatialImgDetections>("detections");
        auto rgbFrame = in->get<dai::ImgFrame>("rgb");

        cv::Mat depthCv = depthFrame->getCvFrame();
        cv::Mat rgbCv = rgbFrame->getCvFrame();
        cv::Mat depthFrameColor = processDepthFrame(depthCv);
        displayResults(rgbCv, depthFrameColor, detections->detections);

        return nullptr;
    }

   private:
    cv::Mat processDepthFrame(const cv::Mat& depthFrame) {
        // Downscale depth frame
        cv::Mat depthDownscaled;
        cv::resize(depthFrame, depthDownscaled, cv::Size(), 0.25, 0.25);

        // Find min and max depth values
        double minDepth = 0, maxDepth = 0;
        cv::Mat mask = (depthDownscaled != 0);
        if(cv::countNonZero(mask) > 0) {
            cv::minMaxLoc(depthDownscaled, &minDepth, &maxDepth, nullptr, nullptr, mask);
        }

        // Normalize depth frame
        cv::Mat depthFrameColor;
        depthFrame.convertTo(depthFrameColor, CV_8UC1, 255.0 / (maxDepth - minDepth), -minDepth * 255.0 / (maxDepth - minDepth));

        // Apply color map
        cv::Mat colorized;
        cv::applyColorMap(depthFrameColor, colorized, cv::COLORMAP_HOT);
        return colorized;
    }

    void displayResults(cv::Mat& rgbFrame, cv::Mat& depthFrameColor, const std::vector<dai::SpatialImgDetection>& detections) {
        int height = rgbFrame.rows;
        int width = rgbFrame.cols;

        for(const auto& detection : detections) {
            drawBoundingBoxes(depthFrameColor, detection);
            drawDetections(rgbFrame, detection, width, height);
        }

        cv::imshow("depth", depthFrameColor);
        cv::imshow("rgb", rgbFrame);

        if(cv::waitKey(1) == 'q') {
            stopPipeline();
        }
    }

    void drawBoundingBoxes(cv::Mat& depthFrameColor, const dai::SpatialImgDetection& detection) {
        auto roi = detection.boundingBoxMapping.roi;
        roi = roi.denormalize(depthFrameColor.cols, depthFrameColor.rows);
        auto topLeft = roi.topLeft();
        auto bottomRight = roi.bottomRight();
        cv::rectangle(depthFrameColor,
                      cv::Point(static_cast<int>(topLeft.x), static_cast<int>(topLeft.y)),
                      cv::Point(static_cast<int>(bottomRight.x), static_cast<int>(bottomRight.y)),
                      cv::Scalar(255, 255, 255),
                      1);
    }

    void drawDetections(cv::Mat& frame, const dai::SpatialImgDetection& detection, int frameWidth, int frameHeight) {
        int x1 = static_cast<int>(detection.xmin * frameWidth);
        int x2 = static_cast<int>(detection.xmax * frameWidth);
        int y1 = static_cast<int>(detection.ymin * frameHeight);
        int y2 = static_cast<int>(detection.ymax * frameHeight);

        std::string label;
        try {
            label = labelMap[detection.label];
        } catch(...) {
            label = std::to_string(detection.label);
        }

        cv::Scalar color(255, 255, 255);
        cv::putText(frame, label, cv::Point(x1 + 10, y1 + 20), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
        cv::putText(frame, std::to_string(detection.confidence * 100), cv::Point(x1 + 10, y1 + 35), cv::FONT_HERSHEY_TRIPLEX, 0.5, color);
        cv::putText(frame,
                    "X: " + std::to_string(static_cast<int>(detection.spatialCoordinates.x)) + " mm",
                    cv::Point(x1 + 10, y1 + 50),
                    cv::FONT_HERSHEY_TRIPLEX,
                    0.5,
                    color);
        cv::putText(frame,
                    "Y: " + std::to_string(static_cast<int>(detection.spatialCoordinates.y)) + " mm",
                    cv::Point(x1 + 10, y1 + 65),
                    cv::FONT_HERSHEY_TRIPLEX,
                    0.5,
                    color);
        cv::putText(frame,
                    "Z: " + std::to_string(static_cast<int>(detection.spatialCoordinates.z)) + " mm",
                    cv::Point(x1 + 10, y1 + 80),
                    cv::FONT_HERSHEY_TRIPLEX,
                    0.5,
                    color);
        cv::rectangle(frame, cv::Point(x1, y1), cv::Point(x2, y2), color, 1);
    }
};

int main(int argc, char** argv) {
    signal(SIGTERM, signalHandler);
    signal(SIGINT, signalHandler);

    // Initialize argument parser
    argparse::ArgumentParser program("spatial_detection", "1.0.0");
    program.add_description("Spatial detection network example with configurable depth source");
    program.add_argument("--depthSource").default_value(std::string("stereo")).help("Depth source: stereo, neural, tof");

    // Parse arguments
    try {
        program.parse_args(argc, argv);
    } catch(const std::runtime_error& err) {
        std::cerr << err.what() << '\n';
        std::cerr << program;
        return EXIT_FAILURE;
    }

    // Get arguments
    std::string depthSourceArg = program.get<std::string>("--depthSource");

    // Validate depth source argument
    if(depthSourceArg != "stereo" && depthSourceArg != "neural" && depthSourceArg != "tof") {
        std::cerr << "Invalid depth source: " << depthSourceArg << '\n';
        std::cerr << "Valid options are: stereo, neural, tof" << '\n';
        return EXIT_FAILURE;
    }

    try {
        float fps = STEREO_DEFAULT_FPS;
        if(depthSourceArg == "neural") {
            fps = NEURAL_FPS;
        }

        // Create pipeline
        dai::Pipeline pipeline;

        const std::pair<int, int> size = {640, 400};

        // Define sources and outputs
        auto camRgb = pipeline.create<dai::node::Camera>();
        camRgb->build(dai::CameraBoardSocket::CAM_A, std::nullopt, fps);

        auto platform = pipeline.getDefaultDevice()->getPlatform();

        // Create depth source based on argument
        dai::node::DepthSource depthSource;

        if(depthSourceArg == "stereo") {
            auto monoLeft = pipeline.create<dai::node::Camera>();
            auto monoRight = pipeline.create<dai::node::Camera>();
            auto stereo = pipeline.create<dai::node::StereoDepth>();

            monoLeft->build(dai::CameraBoardSocket::CAM_B, std::nullopt, fps);
            monoRight->build(dai::CameraBoardSocket::CAM_C, std::nullopt, fps);

            stereo->setExtendedDisparity(true);
            monoLeft->requestOutput(size, std::nullopt, dai::ImgResizeMode::CROP)->link(stereo->left);
            monoRight->requestOutput(size, std::nullopt, dai::ImgResizeMode::CROP)->link(stereo->right);

            depthSource = stereo;
        } else if(depthSourceArg == "neural") {
            auto monoLeft = pipeline.create<dai::node::Camera>();
            auto monoRight = pipeline.create<dai::node::Camera>();

            monoLeft->build(dai::CameraBoardSocket::CAM_B, std::nullopt, fps);
            monoRight->build(dai::CameraBoardSocket::CAM_C, std::nullopt, fps);

            auto neuralDepth = pipeline.create<dai::node::NeuralDepth>();
            neuralDepth->build(*monoLeft->requestFullResolutionOutput(), *monoRight->requestFullResolutionOutput(), dai::DeviceModelZoo::NEURAL_DEPTH_LARGE);

            depthSource = neuralDepth;
        } else if(depthSourceArg == "tof") {
            auto tof = pipeline.create<dai::node::ToF>();
            depthSource = tof;
        }

        // Create spatial detection network using the unified build method with DepthSource variant
        auto spatialDetectionNetwork = pipeline.create<dai::node::SpatialDetectionNetwork>();
        auto visualizer = pipeline.create<SpatialVisualizer>();

        // Configure spatial detection network
        spatialDetectionNetwork->input.setBlocking(false);
        spatialDetectionNetwork->setBoundingBoxScaleFactor(0.5f);
        spatialDetectionNetwork->setDepthLowerThreshold(100);
        spatialDetectionNetwork->setDepthUpperThreshold(5000);

        // Set up model and build with DepthSource variant
        dai::NNModelDescription modelDesc;
        // For better results on OAK4, use a segmentation model like "luxonis/yolov8-instance-segmentation-large:coco-640x480"
        // for depth estimation over the objects mask instead of the full bounding box.
        modelDesc.model = "yolov6-nano";
        spatialDetectionNetwork->build(camRgb, depthSource, modelDesc);

        // Set label map
        visualizer->labelMap = spatialDetectionNetwork->getClasses().value();
        spatialDetectionNetwork->spatialLocationCalculator->initialConfig->setSegmentationPassthrough(false);

        // Linking
        visualizer->build(spatialDetectionNetwork->passthroughDepth, spatialDetectionNetwork->out, spatialDetectionNetwork->passthrough);

        std::cout << "Pipeline starting with depth source: " << depthSourceArg << '\n';

        // Start pipeline
        pipeline.run();

    } catch(const std::exception& e) {
        std::cerr << "Error: " << e.what() << '\n';
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}
```

### Need assistance?

Head over to [Discussion Forum](https://discuss.luxonis.com/) for technical support or any other questions you might have.