Neural Depth Align

1#!/usr/bin/env python3
2
3import numpy as np
4import cv2
5import depthai as dai
6import time
7from datetime import timedelta
8FPS = 10
9
10RGB_SOCKET = dai.CameraBoardSocket.CAM_A
11LEFT_SOCKET = dai.CameraBoardSocket.CAM_B
12RIGHT_SOCKET = dai.CameraBoardSocket.CAM_C
13
14class FPSCounter:
15    def __init__(self):
16        self.frameTimes = []
17
18    def tick(self):
19        now = time.time()
20        self.frameTimes.append(now)
21        self.frameTimes = self.frameTimes[-10:]
22
23    def getFps(self):
24        if len(self.frameTimes) <= 1:
25            return 0
26        return (len(self.frameTimes) - 1) / (self.frameTimes[-1] - self.frameTimes[0])
27
28pipeline = dai.Pipeline()
29
30platform = pipeline.getDefaultDevice().getPlatform()
31
32# Define sources and outputs
33camRgb = pipeline.create(dai.node.Camera).build(RGB_SOCKET)
34left = pipeline.create(dai.node.Camera).build(LEFT_SOCKET)
35right = pipeline.create(dai.node.Camera).build(RIGHT_SOCKET)
36stereo = pipeline.create(dai.node.NeuralDepth)
37sync = pipeline.create(dai.node.Sync)
38align = pipeline.create(dai.node.ImageAlign)
39
40sync.setSyncThreshold(timedelta(seconds=1/(2*FPS)))
41
42rgbOut = camRgb.requestOutput(size = (1280, 960), fps = FPS, enableUndistortion=True)
43leftOut = left.requestFullResolutionOutput(fps = FPS)
44rightOut = right.requestFullResolutionOutput(fps = FPS)
45stereo.build(leftOut, rightOut, dai.DeviceModelZoo.NEURAL_DEPTH_LARGE)
46# Linking
47stereo.depth.link(align.input)
48rgbOut.link(align.inputAlignTo)
49rgbOut.link(sync.inputs["rgb"])
50align.outputAligned.link(sync.inputs["depth_aligned"])
51
52queue = sync.out.createOutputQueue()
53
54def colorizeDepth(frameDepth):
55    invalidMask = frameDepth == 0
56    # Log the depth, minDepth and maxDepth
57    try:
58        minDepth = np.percentile(frameDepth[frameDepth != 0], 3)
59        maxDepth = np.percentile(frameDepth[frameDepth != 0], 95)
60        logDepth = np.log(frameDepth, where=frameDepth != 0)
61        logMinDepth = np.log(minDepth)
62        logMaxDepth = np.log(maxDepth)
63        np.nan_to_num(logDepth, copy=False, nan=logMinDepth)
64        # Clip the values to be in the 0-255 range
65        logDepth = np.clip(logDepth, logMinDepth, logMaxDepth)
66
67        # Interpolate only valid logDepth values, setting the rest based on the mask
68        depthFrameColor = np.interp(logDepth, (logMinDepth, logMaxDepth), (0, 255))
69        depthFrameColor = np.nan_to_num(depthFrameColor)
70        depthFrameColor = depthFrameColor.astype(np.uint8)
71        depthFrameColor = cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_JET)
72        # Set invalid depth pixels to black
73        depthFrameColor[invalidMask] = 0
74    except IndexError:
75        # Frame is likely empty
76        depthFrameColor = np.zeros((frameDepth.shape[0], frameDepth.shape[1], 3), dtype=np.uint8)
77    except Exception as e:
78        raise e
79    return depthFrameColor
80
81
82rgbWeight = 0.4
83depthWeight = 0.6
84
85
86def updateBlendWeights(percentRgb):
87    """
88    Update the rgb and depth weights used to blend depth/rgb image
89
90    @param[in] percent_rgb The rgb weight expressed as a percentage (0..100)
91    """
92    global depthWeight
93    global rgbWeight
94    rgbWeight = float(percentRgb) / 100.0
95    depthWeight = 1.0 - rgbWeight
96
97
98# Connect to device and start pipeline
99with pipeline:
100    pipeline.start()
101
102    # Configure windows; trackbar adjusts blending ratio of rgb/depth
103    windowName = "rgb-depth"
104
105    # Set the window to be resizable and the initial size
106    cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
107    cv2.resizeWindow(windowName, 1280, 720)
108    cv2.createTrackbar(
109        "RGB Weight %",
110        windowName,
111        int(rgbWeight * 100),
112        100,
113        updateBlendWeights,
114    )
115    fpsCounter = FPSCounter()
116    while True:
117        messageGroup = queue.get()
118        fpsCounter.tick()
119        assert isinstance(messageGroup, dai.MessageGroup)
120        frameRgb = messageGroup["rgb"]
121        assert isinstance(frameRgb, dai.ImgFrame)
122        frameDepth = messageGroup["depth_aligned"]
123        assert isinstance(frameDepth, dai.ImgFrame)
124
125        # Blend when both received
126        if frameDepth is not None:
127            cvFrame = frameRgb.getCvFrame()
128            # Colorize the aligned depth
129            alignedDepthColorized = colorizeDepth(frameDepth.getFrame())
130            # Resize depth to match the rgb frame
131            cv2.imshow("Depth aligned", alignedDepthColorized)
132
133            if len(cvFrame.shape) == 2:
134                cvFrameUndistorted = cv2.cvtColor(cvFrame, cv2.COLOR_GRAY2BGR)
135            # print("RGB Size:", cvFrame.shape)
136            # print("Depth Size:", alignedDepthColorized.shape)
137            blended = cv2.addWeighted(
138                cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0
139            )
140            cv2.putText(
141                blended,
142                f"FPS: {fpsCounter.getFps():.2f}",
143                (10, 30),
144                cv2.FONT_HERSHEY_SIMPLEX,
145                1,
146                (255, 255, 255),
147                2,
148            )
149            cv2.imshow(windowName, blended)
150
151        key = cv2.waitKey(1)
152        if key == ord("q"):
153            break

1#include <algorithm>
2#include <chrono>
3#include <cmath>
4#include <deque>
5#include <opencv2/opencv.hpp>
6#include <optional>
7#include <string>
8#include <vector>
9
10#include "depthai/capabilities/ImgFrameCapability.hpp"
11#include "depthai/depthai.hpp"
12
13// Define constants from the Python script
14constexpr float FPS = 10.0f;
15const dai::CameraBoardSocket RGB_SOCKET = dai::CameraBoardSocket::CAM_A;
16const dai::CameraBoardSocket LEFT_SOCKET = dai::CameraBoardSocket::CAM_B;
17const dai::CameraBoardSocket RIGHT_SOCKET = dai::CameraBoardSocket::CAM_C;
18
19// FPS Counter class to calculate and display frames per second
20class FPSCounter {
21   public:
22    void tick() {
23        auto now = std::chrono::steady_clock::now();
24        frameTimes.push_back(now);
25        // Keep the last 10 timestamps, same as the Python version
26        if(frameTimes.size() > 10) {
27            frameTimes.pop_front();
28        }
29    }
30
31    double getFps() const {
32        if(frameTimes.size() <= 1) {
33            return 0.0;
34        }
35        auto duration = std::chrono::duration_cast<std::chrono::duration<double>>(frameTimes.back() - frameTimes.front()).count();
36        return (static_cast<double>(frameTimes.size()) - 1.0) / duration;
37    }
38
39   private:
40    std::deque<std::chrono::steady_clock::time_point> frameTimes;
41};
42
43// Function to colorize a depth frame for visualization
44cv::Mat colorizeDepth(const cv::Mat& frameDepth) {
45    if(frameDepth.empty() || frameDepth.channels() != 1) {
46        return cv::Mat::zeros(frameDepth.size(), CV_8UC3);
47    }
48
49    cv::Mat depth32f;
50    frameDepth.convertTo(depth32f, CV_32F);
51
52    const cv::Mat nonZeroMask = depth32f != 0.0f;
53    const int nz = cv::countNonZero(nonZeroMask);
54    if(nz == 0) {
55        return cv::Mat::zeros(frameDepth.size(), CV_8UC3);
56    }
57
58    // Extract non-zero depth values to calculate percentiles
59    std::vector<float> values;
60    values.reserve(nz);
61    for(int r = 0; r < depth32f.rows; ++r) {
62        const float* d = depth32f.ptr<float>(r);
63        const uchar* m = nonZeroMask.ptr<uchar>(r);
64        for(int c = 0; c < depth32f.cols; ++c) {
65            if(m[c]) {
66                values.push_back(d[c]);
67            }
68        }
69    }
70
71    std::sort(values.begin(), values.end());
72
73    // Lambda to calculate percentile
74    auto pct = [&](double p) {
75        if(values.empty()) return 0.0f;
76        size_t idx = static_cast<size_t>(std::round((p / 100.0) * (values.size() - 1)));
77        return values[idx];
78    };
79
80    const float minDepth = pct(3.0);
81    const float maxDepth = pct(95.0);
82
83    // Apply logarithmic scaling
84    cv::Mat logDepth;
85    depth32f.copyTo(logDepth);
86    logDepth.setTo(minDepth, ~nonZeroMask);  // Replace zeros to avoid log(0)
87    cv::log(logDepth, logDepth);
88
89    const float logMinDepth = std::log(minDepth);
90    const float logMaxDepth = std::log(maxDepth);
91
92    // Clip and linearly scale to the [0, 255] range
93    cv::min(logDepth, logMaxDepth, logDepth);
94    cv::max(logDepth, logMinDepth, logDepth);
95    if(logMaxDepth > logMinDepth) {
96        logDepth = (logDepth - logMinDepth) * (255.0f / (logMaxDepth - logMinDepth));
97    }
98
99    cv::Mat depth8U;
100    logDepth.convertTo(depth8U, CV_8U);
101
102    // Apply color map and set invalid pixels to black
103    cv::Mat depthFrameColor;
104    cv::applyColorMap(depth8U, depthFrameColor, cv::COLORMAP_JET);
105    depthFrameColor.setTo(cv::Scalar::all(0), ~nonZeroMask);
106
107    return depthFrameColor;
108}
109
110// Global variables for blending weights, controlled by the trackbar
111float rgbWeight = 0.4f;
112float depthWeight = 0.6f;
113
114// Callback function for the OpenCV trackbar
115void updateBlendWeights(int percentRgb, void*) {
116    rgbWeight = static_cast<float>(percentRgb) / 100.0f;
117    depthWeight = 1.0f - rgbWeight;
118}
119
120int main() {
121    // Create the DepthAI pipeline
122    dai::Pipeline pipeline;
123
124    // --- Define pipeline nodes ---
125    auto camRgb = pipeline.create<dai::node::Camera>();
126    camRgb->build(RGB_SOCKET);
127
128    auto left = pipeline.create<dai::node::Camera>();
129    left->build(LEFT_SOCKET);
130
131    auto right = pipeline.create<dai::node::Camera>();
132    right->build(RIGHT_SOCKET);
133
134    auto stereo = pipeline.create<dai::node::NeuralDepth>();
135    auto sync = pipeline.create<dai::node::Sync>();
136    auto align = pipeline.create<dai::node::ImageAlign>();
137
138    // --- Configure nodes ---
139    sync->setSyncThreshold(std::chrono::milliseconds(static_cast<long long>(1000.0 / (2.0 * FPS))));
140
141    // auto* rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), std::nullopt, std::nullopt, FPS, true);
142    auto* rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), std::nullopt, dai::ImgResizeMode::CROP, FPS, true);
143    auto* leftOut = left->requestFullResolutionOutput(std::nullopt, FPS);
144    auto* rightOut = right->requestFullResolutionOutput(std::nullopt, FPS);
145
146    stereo->build(*leftOut, *rightOut, dai::DeviceModelZoo::NEURAL_DEPTH_LARGE);
147
148    // --- Link pipeline nodes ---
149    stereo->depth.link(align->input);
150    rgbOut->link(align->inputAlignTo);
151    rgbOut->link(sync->inputs["rgb"]);
152    align->outputAligned.link(sync->inputs["depth_aligned"]);
153
154    // Create an output queue for the synchronized frames
155    auto queue = sync->out.createOutputQueue();
156
157    // Start the pipeline
158    pipeline.start();
159
160    // --- Setup OpenCV windows and trackbar ---
161    const std::string windowName = "rgb-depth";
162    cv::namedWindow(windowName, cv::WINDOW_NORMAL);
163    cv::resizeWindow(windowName, 1280, 720);
164    cv::createTrackbar("RGB Weight %", windowName, nullptr, 100, updateBlendWeights);
165    cv::setTrackbarPos("RGB Weight %", windowName, static_cast<int>(rgbWeight * 100));
166
167    FPSCounter fpsCounter;
168
169    while(true) {
170        // Get the synchronized message group from the queue
171        auto messageGroup = queue->get<dai::MessageGroup>();
172        fpsCounter.tick();
173
174        auto frameRgb = messageGroup->get<dai::ImgFrame>("rgb");
175        auto frameDepth = messageGroup->get<dai::ImgFrame>("depth_aligned");
176
177        if(frameRgb && frameDepth) {
178            cv::Mat cvFrame = frameRgb->getCvFrame();
179
180            // Colorize the aligned depth frame for visualization
181            cv::Mat alignedDepthColorized = colorizeDepth(frameDepth->getFrame());
182            cv::imshow("Depth aligned", alignedDepthColorized);
183
184            // Blend the RGB and colorized depth frames
185            cv::Mat blended;
186            cv::addWeighted(cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0, blended);
187
188            // Add FPS text to the blended image
189            char fpsStr[20];
190            snprintf(fpsStr, sizeof(fpsStr), "FPS: %.2f", fpsCounter.getFps());
191            cv::putText(blended, fpsStr, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 255, 255), 2);
192
193            // Show the final blended result
194            cv::imshow(windowName, blended);
195        }
196
197        // Check for 'q' key press to exit
198        int key = cv::waitKey(1);
199        if(key == 'q' || key == 27) {  // 'q' or ESC
200            break;
201        }
202    }
203
204    return 0;
205}