Stereo 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 = 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)
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.log(frameDepth, where=frameDepth != 0)
69        logMinDepth = np.log(minDepth)
70        logMaxDepth = np.log(maxDepth)
71        np.nan_to_num(logDepth, copy=False, nan=logMinDepth)
72        # Clip the values to be in the 0-255 range
73        logDepth = np.clip(logDepth, logMinDepth, logMaxDepth)
74
75        # Interpolate only valid logDepth values, setting the rest based on the mask
76        depthFrameColor = np.interp(logDepth, (logMinDepth, logMaxDepth), (0, 255))
77        depthFrameColor = np.nan_to_num(depthFrameColor)
78        depthFrameColor = depthFrameColor.astype(np.uint8)
79        depthFrameColor = cv2.applyColorMap(depthFrameColor, cv2.COLORMAP_JET)
80        # Set invalid depth pixels to black
81        depthFrameColor[invalidMask] = 0
82    except IndexError:
83        # Frame is likely empty
84        depthFrameColor = np.zeros((frameDepth.shape[0], frameDepth.shape[1], 3), dtype=np.uint8)
85    except Exception as e:
86        raise e
87    return depthFrameColor
88
89
90rgbWeight = 0.4
91depthWeight = 0.6
92
93
94def updateBlendWeights(percentRgb):
95    """
96    Update the rgb and depth weights used to blend depth/rgb image
97
98    @param[in] percent_rgb The rgb weight expressed as a percentage (0..100)
99    """
100    global depthWeight
101    global rgbWeight
102    rgbWeight = float(percentRgb) / 100.0
103    depthWeight = 1.0 - rgbWeight
104
105
106# Connect to device and start pipeline
107with pipeline:
108    pipeline.start()
109
110    # Configure windows; trackbar adjusts blending ratio of rgb/depth
111    windowName = "rgb-depth"
112
113    # Set the window to be resizable and the initial size
114    cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
115    cv2.resizeWindow(windowName, 1280, 720)
116    cv2.createTrackbar(
117        "RGB Weight %",
118        windowName,
119        int(rgbWeight * 100),
120        100,
121        updateBlendWeights,
122    )
123    fpsCounter = FPSCounter()
124    while True:
125        messageGroup = queue.get()
126        fpsCounter.tick()
127        assert isinstance(messageGroup, dai.MessageGroup)
128        frameRgb = messageGroup["rgb"]
129        assert isinstance(frameRgb, dai.ImgFrame)
130        frameDepth = messageGroup["depth_aligned"]
131        assert isinstance(frameDepth, dai.ImgFrame)
132
133        # Blend when both received
134        if frameDepth is not None:
135            cvFrame = frameRgb.getCvFrame()
136            cvFrameUndistorted = cv2.undistort(
137                cvFrame,
138                np.array(frameRgb.getTransformation().getIntrinsicMatrix()),
139                np.array(frameRgb.getTransformation().getDistortionCoefficients()),
140            )
141            # Colorize the aligned depth
142            alignedDepthColorized = colorizeDepth(frameDepth.getFrame())
143            # Resize depth to match the rgb frame
144            cv2.imshow("Depth aligned", alignedDepthColorized)
145
146            if len(cvFrameUndistorted.shape) == 2:
147                cvFrameUndistorted = cv2.cvtColor(cvFrameUndistorted, cv2.COLOR_GRAY2BGR)
148            blended = cv2.addWeighted(
149                cvFrameUndistorted, rgbWeight, alignedDepthColorized, depthWeight, 0
150            )
151            cv2.putText(
152                blended,
153                f"FPS: {fpsCounter.getFps():.2f}",
154                (10, 30),
155                cv2.FONT_HERSHEY_SIMPLEX,
156                1,
157                (255, 255, 255),
158                2,
159            )
160            cv2.imshow(windowName, blended)
161
162        key = cv2.waitKey(1)
163        if key == ord("q"):
164            break