Depth Align

View as Markdown

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

1#include <atomic>
2#include <chrono>
3#include <cmath>
4#include <csignal>
5#include <deque>
6#include <iostream>
7#include <opencv2/opencv.hpp>
8#include <optional>
9#include <string>
10#include <vector>
11
12#include "depthai/depthai.hpp"
13
14std::atomic<bool> quitEvent(false);
15
16void signalHandler(int) {
17    quitEvent = true;
18}
19
20constexpr float FPS = 25.0f;
21
22const dai::CameraBoardSocket RGB_SOCKET = dai::CameraBoardSocket::CAM_A;
23const dai::CameraBoardSocket LEFT_SOCKET = dai::CameraBoardSocket::CAM_B;
24const dai::CameraBoardSocket RIGHT_SOCKET = dai::CameraBoardSocket::CAM_C;
25
26// FPS Counter class
27class FPSCounter {
28   public:
29    void tick() {
30        auto now = std::chrono::steady_clock::now();
31        frameTimes.push_back(now);
32        while(frameTimes.size() > 10) {
33            frameTimes.pop_front();
34        }
35    }
36
37    float getFps() {
38        if(frameTimes.size() <= 1) return 0.0f;
39        auto duration = std::chrono::duration_cast<std::chrono::duration<float>>(frameTimes.back() - frameTimes.front()).count();
40        return (frameTimes.size() - 1) / duration;
41    }
42
43   private:
44    std::deque<std::chrono::steady_clock::time_point> frameTimes;
45};
46
47// Depth colorization function from detection_network_remap.cpp
48cv::Mat colorizeDepth(cv::Mat frameDepth) {
49    try {
50        // Early exit if no valid pixels
51        if(cv::countNonZero(frameDepth) == 0) {
52            return cv::Mat::zeros(frameDepth.rows, frameDepth.cols, CV_8UC3);
53        }
54
55        // Convert to float once
56        cv::Mat frameDepthFloat;
57        frameDepth.convertTo(frameDepthFloat, CV_32F);
58
59        double minVal, maxVal;
60        cv::minMaxLoc(frameDepthFloat, &minVal, &maxVal, nullptr, nullptr, frameDepthFloat > 0);
61
62        // Take log in-place
63        cv::log(frameDepthFloat, frameDepthFloat);
64        float logMinDepth = std::log(minVal);
65        float logMaxDepth = std::log(maxVal);
66
67        frameDepthFloat = (frameDepthFloat - logMinDepth) * (255.0f / (logMaxDepth - logMinDepth));
68
69        cv::Mat normalizedDepth;
70        frameDepthFloat.convertTo(normalizedDepth, CV_8UC1);
71
72        cv::Mat depthFrameColor;
73        cv::applyColorMap(normalizedDepth, depthFrameColor, cv::COLORMAP_JET);
74
75        // Mask invalid pixels
76        depthFrameColor.setTo(0, frameDepth == 0);
77
78        return depthFrameColor;
79
80    } catch(const std::exception& e) {
81        std::cerr << "Error in colorizeDepth: " << e.what() << std::endl;
82        return cv::Mat::zeros(frameDepth.rows, frameDepth.cols, CV_8UC3);
83    }
84}
85
86// Global blend weights
87float rgbWeight = 0.4f;
88float depthWeight = 0.6f;
89
90// Trackbar callback
91void updateBlendWeights(int percentRgb, void*) {
92    rgbWeight = static_cast<float>(percentRgb) / 100.0f;
93    depthWeight = 1.0f - rgbWeight;
94}
95
96int main() {
97    signal(SIGTERM, signalHandler);
98    signal(SIGINT, signalHandler);
99
100    dai::Pipeline pipeline;
101
102    // Create and configure nodes
103    auto camRgb = pipeline.create<dai::node::Camera>();
104    camRgb->build(RGB_SOCKET);
105    auto left = pipeline.create<dai::node::Camera>();
106    left->build(LEFT_SOCKET);
107    auto right = pipeline.create<dai::node::Camera>();
108    right->build(RIGHT_SOCKET);
109    auto stereo = pipeline.create<dai::node::StereoDepth>();
110    auto sync = pipeline.create<dai::node::Sync>();
111
112    // Check if platform is RVC4 and create ImageAlign node if needed
113    auto platform = pipeline.getDefaultDevice()->getPlatform();
114    std::shared_ptr<dai::node::ImageAlign> align;
115    if(platform == dai::Platform::RVC4) {
116        align = pipeline.create<dai::node::ImageAlign>();
117    }
118
119    stereo->setExtendedDisparity(true);
120    sync->setSyncThreshold(std::chrono::duration<int64_t, std::nano>(static_cast<int64_t>(1e9 / (2.0 * FPS))));
121
122    // Configure outputs
123    auto rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), dai::ImgFrame::Type::NV12, dai::ImgResizeMode::CROP, FPS, true);
124    auto leftOut = left->requestOutput(std::make_pair(640, 400), std::nullopt, dai::ImgResizeMode::CROP, FPS);
125    auto rightOut = right->requestOutput(std::make_pair(640, 400), std::nullopt, dai::ImgResizeMode::CROP, FPS);
126
127    // Link nodes
128    rgbOut->link(sync->inputs["rgb"]);
129    leftOut->link(stereo->left);
130    rightOut->link(stereo->right);
131
132    if(platform == dai::Platform::RVC4) {
133        stereo->depth.link(align->input);
134        rgbOut->link(align->inputAlignTo);
135        align->outputAligned.link(sync->inputs["depth_aligned"]);
136    } else {
137        stereo->depth.link(sync->inputs["depth_aligned"]);
138        rgbOut->link(stereo->inputAlignTo);
139    }
140
141    // Create output queue
142    auto queue = sync->out.createOutputQueue();
143
144    // Create and configure windows
145    const std::string windowName = "rgb-depth";
146    cv::namedWindow(windowName, cv::WINDOW_NORMAL);
147    cv::resizeWindow(windowName, 1280, 720);
148    cv::createTrackbar("RGB Weight %", windowName, nullptr, 100, updateBlendWeights);
149    cv::setTrackbarPos("RGB Weight %", windowName, static_cast<int>(rgbWeight * 100));
150
151    FPSCounter fpsCounter;
152
153    // Start pipeline
154    pipeline.start();
155
156    while(pipeline.isRunning() && !quitEvent) {
157        auto messageGroup = queue->get<dai::MessageGroup>();
158        fpsCounter.tick();
159
160        auto frameRgb = messageGroup->get<dai::ImgFrame>("rgb");
161        auto frameDepth = messageGroup->get<dai::ImgFrame>("depth_aligned");
162
163        if(frameDepth != nullptr) {
164            cv::Mat cvFrame = frameRgb->getCvFrame();
165
166            // Colorize depth
167            cv::Mat alignedDepthColorized = colorizeDepth(frameDepth->getFrame());
168            cv::imshow("Depth aligned", alignedDepthColorized);
169
170            // Convert grayscale to BGR if needed
171            if(cvFrame.channels() == 1) {
172                cv::cvtColor(cvFrame, cvFrame, cv::COLOR_GRAY2BGR);
173            }
174
175            // Blend frames
176            cv::Mat blended;
177            cv::addWeighted(cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0, blended);
178
179            // Add FPS text
180            cv::putText(blended, "FPS: " + std::to_string(fpsCounter.getFps()), cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 255, 255), 2);
181
182            cv::imshow(windowName, blended);
183        }
184
185        if(cv::waitKey(1) == 'q') {
186            break;
187        }
188    }
189
190    pipeline.stop();
191    pipeline.wait();
192
193    return 0;
194}