神经深度对齐

1#!/usr/bin/env python3
2
3import cv2
4import depthai as dai
5import time
6from datetime import timedelta
7FPS = 25
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.NeuralDepth)
36sync = pipeline.create(dai.node.Sync)
37align = pipeline.create(dai.node.ImageAlign)
38
39sync.setSyncThreshold(timedelta(seconds=1/(2*FPS)))
40
41rgbOut = camRgb.requestOutput(size = (1280, 960), fps = FPS, enableUndistortion=True)
42leftOut = left.requestFullResolutionOutput(fps = FPS)
43rightOut = right.requestFullResolutionOutput(fps = FPS)
44stereo.build(leftOut, rightOut, dai.DeviceModelZoo.NEURAL_DEPTH_LARGE)
45# Linking
46stereo.depth.link(align.input)
47rgbOut.link(align.inputAlignTo)
48rgbOut.link(sync.inputs["rgb"])
49align.outputAligned.link(sync.inputs["depth_aligned"])
50
51queue = sync.out.createOutputQueue()
52
53rgbWeight = 0.4
54depthWeight = 0.6
55
56
57def updateBlendWeights(percentRgb):
58    """
59    Update the rgb and depth weights used to blend depth/rgb image
60
61    @param[in] percent_rgb The rgb weight expressed as a percentage (0..100)
62    """
63    global depthWeight
64    global rgbWeight
65    rgbWeight = float(percentRgb) / 100.0
66    depthWeight = 1.0 - rgbWeight
67
68
69# Connect to device and start pipeline
70with pipeline:
71    pipeline.start()
72
73    # Configure windows; trackbar adjusts blending ratio of rgb/depth
74    windowName = "rgb-depth"
75
76    # Set the window to be resizable and the initial size
77    cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
78    cv2.resizeWindow(windowName, 1280, 720)
79    cv2.createTrackbar(
80        "RGB Weight %",
81        windowName,
82        int(rgbWeight * 100),
83        100,
84        updateBlendWeights,
85    )
86    fpsCounter = FPSCounter()
87    while True:
88        messageGroup = queue.get()
89        fpsCounter.tick()
90        assert isinstance(messageGroup, dai.MessageGroup)
91        frameRgb = messageGroup["rgb"]
92        assert isinstance(frameRgb, dai.ImgFrame)
93        frameDepth = messageGroup["depth_aligned"]
94        assert isinstance(frameDepth, dai.ImgFrame)
95
96        # Blend when both received
97        if frameDepth is not None:
98            cvFrame = frameRgb.getCvFrame()
99            # Colorize the aligned depth
100            alignedDepthColorized = dai.utility.colorizeDepthFrame(frameDepth).getCvFrame()
101            # Resize depth to match the rgb frame
102            cv2.imshow("Depth aligned", alignedDepthColorized)
103
104            if len(cvFrame.shape) == 2:
105                cvFrameUndistorted = cv2.cvtColor(cvFrame, cv2.COLOR_GRAY2BGR)
106            # print("RGB Size:", cvFrame.shape)
107            # print("Depth Size:", alignedDepthColorized.shape)
108            blended = cv2.addWeighted(
109                cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0
110            )
111            cv2.putText(
112                blended,
113                f"FPS: {fpsCounter.getFps():.2f}",
114                (10, 30),
115                cv2.FONT_HERSHEY_SIMPLEX,
116                1,
117                (255, 255, 255),
118                2,
119            )
120            cv2.imshow(windowName, blended)
121
122        key = cv2.waitKey(1)
123        if key == ord("q"):
124            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 = 25.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// Global variables for blending weights, controlled by the trackbar
44float rgbWeight = 0.4f;
45float depthWeight = 0.6f;
46
47// Callback function for the OpenCV trackbar
48void updateBlendWeights(int percentRgb, void*) {
49    rgbWeight = static_cast<float>(percentRgb) / 100.0f;
50    depthWeight = 1.0f - rgbWeight;
51}
52
53int main() {
54    // Create the DepthAI pipeline
55    dai::Pipeline pipeline;
56
57    // --- Define pipeline nodes ---
58    auto camRgb = pipeline.create<dai::node::Camera>();
59    camRgb->build(RGB_SOCKET);
60
61    auto left = pipeline.create<dai::node::Camera>();
62    left->build(LEFT_SOCKET);
63
64    auto right = pipeline.create<dai::node::Camera>();
65    right->build(RIGHT_SOCKET);
66
67    auto stereo = pipeline.create<dai::node::NeuralDepth>();
68    auto sync = pipeline.create<dai::node::Sync>();
69    auto align = pipeline.create<dai::node::ImageAlign>();
70
71    // --- Configure nodes ---
72    sync->setSyncThreshold(std::chrono::milliseconds(static_cast<long long>(1000.0 / (2.0 * FPS))));
73
74    // auto* rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), std::nullopt, std::nullopt, FPS, true);
75    auto* rgbOut = camRgb->requestOutput(std::make_pair(1280, 960), std::nullopt, dai::ImgResizeMode::CROP, FPS, true);
76    auto* leftOut = left->requestFullResolutionOutput(std::nullopt, FPS);
77    auto* rightOut = right->requestFullResolutionOutput(std::nullopt, FPS);
78
79    stereo->build(*leftOut, *rightOut, dai::DeviceModelZoo::NEURAL_DEPTH_LARGE);
80
81    // --- Link pipeline nodes ---
82    stereo->depth.link(align->input);
83    rgbOut->link(align->inputAlignTo);
84    rgbOut->link(sync->inputs["rgb"]);
85    align->outputAligned.link(sync->inputs["depth_aligned"]);
86
87    // Create an output queue for the synchronized frames
88    auto queue = sync->out.createOutputQueue();
89
90    // Start the pipeline
91    pipeline.start();
92
93    // --- Setup OpenCV windows and trackbar ---
94    const std::string windowName = "rgb-depth";
95    cv::namedWindow(windowName, cv::WINDOW_NORMAL);
96    cv::resizeWindow(windowName, 1280, 720);
97    cv::createTrackbar("RGB Weight %", windowName, nullptr, 100, updateBlendWeights);
98    cv::setTrackbarPos("RGB Weight %", windowName, static_cast<int>(rgbWeight * 100));
99
100    FPSCounter fpsCounter;
101
102    while(true) {
103        // Get the synchronized message group from the queue
104        auto messageGroup = queue->get<dai::MessageGroup>();
105        fpsCounter.tick();
106
107        auto frameRgb = messageGroup->get<dai::ImgFrame>("rgb");
108        auto frameDepth = messageGroup->get<dai::ImgFrame>("depth_aligned");
109
110        if(frameRgb && frameDepth) {
111            cv::Mat cvFrame = frameRgb->getCvFrame();
112
113            // Colorize the aligned depth frame for visualization
114            cv::Mat alignedDepthColorized = dai::utility::colorizeDepthFrame(*frameDepth).getCvFrame();
115            cv::imshow("Depth aligned", alignedDepthColorized);
116
117            // Blend the RGB and colorized depth frames
118            cv::Mat blended;
119            cv::addWeighted(cvFrame, rgbWeight, alignedDepthColorized, depthWeight, 0, blended);
120
121            // Add FPS text to the blended image
122            char fpsStr[20];
123            snprintf(fpsStr, sizeof(fpsStr), "FPS: %.2f", fpsCounter.getFps());
124            cv::putText(blended, fpsStr, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1, cv::Scalar(255, 255, 255), 2);
125
126            // Show the final blended result
127            cv::imshow(windowName, blended);
128        }
129
130        // Check for 'q' key press to exit
131        int key = cv::waitKey(1);
132        if(key == 'q' || key == 27) {  // 'q' or ESC
133            break;
134        }
135    }
136
137    return 0;
138}