深度对齐

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import time
6from datetime import timedelta
7FPS = 25.0
8
9RGB_SOCKET = dai.CameraBoardSocket.CAM_A
10LEFT_SOCKET = dai.CameraBoardSocket.CAM_B
11RIGHT_SOCKET = dai.CameraBoardSocket.CAM_C
12
13class FPSCounter:
14    def __init__(self):
15        self.frameTimes = []
16
17    def tick(self):
18        now = time.time()
19        self.frameTimes.append(now)
20        self.frameTimes = self.frameTimes[-10:]
21
22    def getFps(self):
23        if len(self.frameTimes) <= 1:
24            return 0
25        return (len(self.frameTimes) - 1) / (self.frameTimes[-1] - self.frameTimes[0])
26
27pipeline = dai.Pipeline()
28
29platform = pipeline.getDefaultDevice().getPlatform()
30
31# Define sources and outputs
32camRgb = pipeline.create(dai.node.Camera).build(RGB_SOCKET)
33left = pipeline.create(dai.node.Camera).build(LEFT_SOCKET)
34right = pipeline.create(dai.node.Camera).build(RIGHT_SOCKET)
35stereo = pipeline.create(dai.node.StereoDepth)
36sync = pipeline.create(dai.node.Sync)
37if platform == dai.Platform.RVC4:
38    align = pipeline.create(dai.node.ImageAlign)
39
40stereo.setExtendedDisparity(True)
41sync.setSyncThreshold(timedelta(seconds=1/(2*FPS)))
42
43rgbOut = camRgb.requestOutput(size = (1280, 960), fps = FPS, enableUndistortion=True)
44leftOut = left.requestOutput(size = (640, 400), fps = FPS)
45rightOut = right.requestOutput(size = (640, 400), fps = FPS)
46
47# Linking
48rgbOut.link(sync.inputs["rgb"])
49leftOut.link(stereo.left)
50rightOut.link(stereo.right)
51if platform == dai.Platform.RVC4:
52    stereo.depth.link(align.input)
53    rgbOut.link(align.inputAlignTo)
54    align.outputAligned.link(sync.inputs["depth_aligned"])
55else:
56    stereo.depth.link(sync.inputs["depth_aligned"])
57    rgbOut.link(stereo.inputAlignTo)
58
59queue = sync.out.createOutputQueue()
60
61rgbWeight = 0.4
62depthWeight = 0.6
63
64
65def updateBlendWeights(percentRgb):
66    """
67    Update the rgb and depth weights used to blend depth/rgb image
68
69    @param[in] percent_rgb The rgb weight expressed as a percentage (0..100)
70    """
71    global depthWeight
72    global rgbWeight
73    rgbWeight = float(percentRgb) / 100.0
74    depthWeight = 1.0 - rgbWeight
75
76
77# Connect to device and start pipeline
78with pipeline:
79    pipeline.start()
80
81    # Configure windows; trackbar adjusts blending ratio of rgb/depth
82    windowName = "rgb-depth"
83
84    # Set the window to be resizable and the initial size
85    cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
86    cv2.resizeWindow(windowName, 1280, 720)
87    cv2.createTrackbar(
88        "RGB Weight %",
89        windowName,
90        int(rgbWeight * 100),
91        100,
92        updateBlendWeights,
93    )
94    fpsCounter = FPSCounter()
95    while True:
96        messageGroup = queue.get()
97        fpsCounter.tick()
98        assert isinstance(messageGroup, dai.MessageGroup)
99        frameRgb = messageGroup["rgb"]
100        assert isinstance(frameRgb, dai.ImgFrame)
101        frameDepth = messageGroup["depth_aligned"]
102        assert isinstance(frameDepth, dai.ImgFrame)
103
104        # Blend when both received
105        if frameDepth is not None:
106            cvFrame = frameRgb.getCvFrame()
107            # Colorize the aligned depth
108            alignedDepthColorized = dai.utility.colorizeDepthFrame(frameDepth).getCvFrame()
109            # Resize depth to match the rgb frame
110            cv2.imshow("Depth aligned", alignedDepthColorized)
111
112            if len(cvFrame.shape) == 2:
113                cvFrameUndistorted = cv2.cvtColor(cvFrame, cv2.COLOR_GRAY2BGR)
114            blended = cv2.addWeighted(
115                cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0
116            )
117            cv2.putText(
118                blended,
119                f"FPS: {fpsCounter.getFps():.2f}",
120                (10, 30),
121                cv2.FONT_HERSHEY_SIMPLEX,
122                1,
123                (255, 255, 255),
124                2,
125            )
126            cv2.imshow(windowName, blended)
127
128        key = cv2.waitKey(1)
129        if key == ord("q"):
130            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// Global blend weights
48float rgbWeight = 0.4f;
49float depthWeight = 0.6f;
50
51// Trackbar callback
52void updateBlendWeights(int percentRgb, void*) {
53    rgbWeight = static_cast<float>(percentRgb) / 100.0f;
54    depthWeight = 1.0f - rgbWeight;
55}
56
57int main() {
58    signal(SIGTERM, signalHandler);
59    signal(SIGINT, signalHandler);
60
61    dai::Pipeline pipeline;
62
63    // Create and configure nodes
64    auto camRgb = pipeline.create<dai::node::Camera>();
65    camRgb->build(RGB_SOCKET);
66    auto left = pipeline.create<dai::node::Camera>();
67    left->build(LEFT_SOCKET);
68    auto right = pipeline.create<dai::node::Camera>();
69    right->build(RIGHT_SOCKET);
70    auto stereo = pipeline.create<dai::node::StereoDepth>();
71    auto sync = pipeline.create<dai::node::Sync>();
72
73    // Check if platform is RVC4 and create ImageAlign node if needed
74    auto platform = pipeline.getDefaultDevice()->getPlatform();
75    std::shared_ptr<dai::node::ImageAlign> align;
76    if(platform == dai::Platform::RVC4) {
77        align = pipeline.create<dai::node::ImageAlign>();
78    }
79
80    stereo->setExtendedDisparity(true);
81    sync->setSyncThreshold(std::chrono::duration<int64_t, std::nano>(static_cast<int64_t>(1e9 / (2.0 * FPS))));
82
83    // Configure outputs
84    auto rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), dai::ImgFrame::Type::NV12, dai::ImgResizeMode::CROP, FPS, true);
85    auto leftOut = left->requestOutput(std::make_pair(640, 400), std::nullopt, dai::ImgResizeMode::CROP, FPS);
86    auto rightOut = right->requestOutput(std::make_pair(640, 400), std::nullopt, dai::ImgResizeMode::CROP, FPS);
87
88    // Link nodes
89    rgbOut->link(sync->inputs["rgb"]);
90    leftOut->link(stereo->left);
91    rightOut->link(stereo->right);
92
93    if(platform == dai::Platform::RVC4) {
94        stereo->depth.link(align->input);
95        rgbOut->link(align->inputAlignTo);
96        align->outputAligned.link(sync->inputs["depth_aligned"]);
97    } else {
98        stereo->depth.link(sync->inputs["depth_aligned"]);
99        rgbOut->link(stereo->inputAlignTo);
100    }
101
102    // Create output queue
103    auto queue = sync->out.createOutputQueue();
104
105    // Create and configure windows
106    const std::string windowName = "rgb-depth";
107    cv::namedWindow(windowName, cv::WINDOW_NORMAL);
108    cv::resizeWindow(windowName, 1280, 720);
109    cv::createTrackbar("RGB Weight %", windowName, nullptr, 100, updateBlendWeights);
110    cv::setTrackbarPos("RGB Weight %", windowName, static_cast<int>(rgbWeight * 100));
111
112    FPSCounter fpsCounter;
113
114    // Start pipeline
115    pipeline.start();
116
117    while(pipeline.isRunning() && !quitEvent) {
118        auto messageGroup = queue->get<dai::MessageGroup>();
119        fpsCounter.tick();
120
121        auto frameRgb = messageGroup->get<dai::ImgFrame>("rgb");
122        auto frameDepth = messageGroup->get<dai::ImgFrame>("depth_aligned");
123
124        if(frameDepth != nullptr) {
125            cv::Mat cvFrame = frameRgb->getCvFrame();
126
127            // Colorize depth
128            cv::Mat alignedDepthColorized = dai::utility::colorizeDepthFrame(*frameDepth).getCvFrame();
129            cv::imshow("Depth aligned", alignedDepthColorized);
130
131            // Convert grayscale to BGR if needed
132            if(cvFrame.channels() == 1) {
133                cv::cvtColor(cvFrame, cvFrame, cv::COLOR_GRAY2BGR);
134            }
135
136            // Blend frames
137            cv::Mat blended;
138            cv::addWeighted(cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0, blended);
139
140            // Add FPS text
141            cv::putText(blended, "FPS: " + std::to_string(fpsCounter.getFps()), cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 255, 255), 2);
142
143            cv::imshow(windowName, blended);
144        }
145
146        if(cv::waitKey(1) == 'q') {
147            break;
148        }
149    }
150
151    pipeline.stop();
152    pipeline.wait();
153
154    return 0;
155}