Stereo Depth Video

1#!/usr/bin/env python3
2
3import cv2
4import numpy as np
5import depthai as dai
6import argparse
7
8parser = argparse.ArgumentParser()
9parser.add_argument(
10    "-res",
11    "--resolution",
12    type=str,
13    default="720",
14    help="Sets the resolution on mono cameras. Options: 800 | 720 | 400",
15)
16parser.add_argument(
17    "-md",
18    "--mesh_dir",
19    type=str,
20    default=None,
21    help="Output directory for mesh files. If not specified mesh files won't be saved",
22)
23parser.add_argument(
24    "-lm",
25    "--load_mesh",
26    default=False,
27    action="store_true",
28    help="Read camera intrinsics, generate mesh files and load them into the stereo node.",
29)
30parser.add_argument(
31    "-rect",
32    "--out_rectified",
33    default=False,
34    action="store_true",
35    help="Generate and display rectified streams",
36)
37parser.add_argument(
38    "-lr",
39    "--lrcheck",
40    default=False,
41    action="store_true",
42    help="Better handling for occlusions",
43)
44parser.add_argument(
45    "-e",
46    "--extended",
47    default=False,
48    action="store_true",
49    help="Closer-in minimum depth, disparity range is doubled",
50)
51parser.add_argument(
52    "-s",
53    "--subpixel",
54    default=False,
55    action="store_true",
56    help="Better accuracy for longer distance, fractional disparity 32-levels",
57)
58parser.add_argument(
59    "-m",
60    "--median",
61    type=str,
62    default="7x7",
63    help="Choose the size of median filtering. Options: OFF | 3x3 | 5x5 | 7x7 (default)",
64)
65parser.add_argument(
66    "-d",
67    "--depth",
68    default=False,
69    action="store_true",
70    help="Display depth frames",
71)
72parser.add_argument(
73    "-swlr",
74    "--swap_left_right",
75    default=False,
76    action="store_true",
77    help="Swap left right frames",
78)
79parser.add_argument(
80    "-a",
81    "--alpha",
82    type=float,
83    default=None,
84    help="Alpha scaling parameter to increase FOV",
85)
86args = parser.parse_args()
87
88RES_MAP = {
89    '800': {'w': 1280, 'h': 800, 'res': dai.MonoCameraProperties.SensorResolution.THE_800_P },
90    '720': {'w': 1280, 'h': 720, 'res': dai.MonoCameraProperties.SensorResolution.THE_720_P },
91    '400': {'w': 640, 'h': 400, 'res': dai.MonoCameraProperties.SensorResolution.THE_400_P }
92}
93if args.resolution not in RES_MAP:
94    exit("Unsupported resolution!")
95
96resolution = RES_MAP[args.resolution]
97
98meshDirectory = args.mesh_dir  # Output dir for mesh files
99generateMesh = args.load_mesh  # Load mesh files
100
101outRectified = args.out_rectified  # Output and display rectified streams
102lrcheck = args.lrcheck  # Better handling for occlusions
103extended = args.extended  # Closer-in minimum depth, disparity range is doubled
104subpixel = args.subpixel  # Better accuracy for longer distance, fractional disparity 32-levels
105depth = args.depth  # Display depth frames
106
107medianMap = {
108    "OFF": dai.StereoDepthProperties.MedianFilter.MEDIAN_OFF,
109    "3x3": dai.StereoDepthProperties.MedianFilter.KERNEL_3x3,
110    "5x5": dai.StereoDepthProperties.MedianFilter.KERNEL_5x5,
111    "7x7": dai.StereoDepthProperties.MedianFilter.KERNEL_7x7,
112}
113if args.median not in medianMap:
114    exit("Unsupported median size!")
115
116median = medianMap[args.median]
117
118print("StereoDepth config options:")
119print(f"    Resolution:  {resolution['w']}x{resolution['h']}")
120print("    Left-Right check:  ", lrcheck)
121print("    Extended disparity:", extended)
122print("    Subpixel:          ", subpixel)
123print("    Median filtering:  ", median)
124print("    Generating mesh files:  ", generateMesh)
125print("    Outputting mesh files to:  ", meshDirectory)
126
127
128def getMesh(calibData):
129    M1 = np.array(calibData.getCameraIntrinsics(dai.CameraBoardSocket.CAM_B, resolution[0], resolution[1]))
130    d1 = np.array(calibData.getDistortionCoefficients(dai.CameraBoardSocket.CAM_B))
131    R1 = np.array(calibData.getStereoLeftRectificationRotation())
132    M2 = np.array(calibData.getCameraIntrinsics(dai.CameraBoardSocket.CAM_C, resolution[0], resolution[1]))
133    d2 = np.array(calibData.getDistortionCoefficients(dai.CameraBoardSocket.CAM_C))
134    R2 = np.array(calibData.getStereoRightRectificationRotation())
135    mapXL, mapYL = cv2.initUndistortRectifyMap(M1, d1, R1, M2, resolution, cv2.CV_32FC1)
136    mapXR, mapYR = cv2.initUndistortRectifyMap(M2, d2, R2, M2, resolution, cv2.CV_32FC1)
137
138    meshCellSize = 16
139    meshLeft = []
140    meshRight = []
141
142    for y in range(mapXL.shape[0] + 1):
143        if y % meshCellSize == 0:
144            rowLeft = []
145            rowRight = []
146            for x in range(mapXL.shape[1] + 1):
147                if x % meshCellSize == 0:
148                    if y == mapXL.shape[0] and x == mapXL.shape[1]:
149                        rowLeft.append(mapYL[y - 1, x - 1])
150                        rowLeft.append(mapXL[y - 1, x - 1])
151                        rowRight.append(mapYR[y - 1, x - 1])
152                        rowRight.append(mapXR[y - 1, x - 1])
153                    elif y == mapXL.shape[0]:
154                        rowLeft.append(mapYL[y - 1, x])
155                        rowLeft.append(mapXL[y - 1, x])
156                        rowRight.append(mapYR[y - 1, x])
157                        rowRight.append(mapXR[y - 1, x])
158                    elif x == mapXL.shape[1]:
159                        rowLeft.append(mapYL[y, x - 1])
160                        rowLeft.append(mapXL[y, x - 1])
161                        rowRight.append(mapYR[y, x - 1])
162                        rowRight.append(mapXR[y, x - 1])
163                    else:
164                        rowLeft.append(mapYL[y, x])
165                        rowLeft.append(mapXL[y, x])
166                        rowRight.append(mapYR[y, x])
167                        rowRight.append(mapXR[y, x])
168            if (mapXL.shape[1] % meshCellSize) % 2 != 0:
169                rowLeft.append(0)
170                rowLeft.append(0)
171                rowRight.append(0)
172                rowRight.append(0)
173
174            meshLeft.append(rowLeft)
175            meshRight.append(rowRight)
176
177    meshLeft = np.array(meshLeft)
178    meshRight = np.array(meshRight)
179
180    return meshLeft, meshRight
181
182
183def saveMeshFiles(meshLeft, meshRight, outputPath):
184    print("Saving mesh to:", outputPath)
185    meshLeft.tofile(outputPath + "/left_mesh.calib")
186    meshRight.tofile(outputPath + "/right_mesh.calib")
187
188
189def getDisparityFrame(frame, cvColorMap):
190    maxDisp = stereo.initialConfig.getMaxDisparity()
191    disp = (frame * (255.0 / maxDisp)).astype(np.uint8)
192    disp = cv2.applyColorMap(disp, cvColorMap)
193
194    return disp
195
196device = dai.Device()
197calibData = device.readCalibration()
198print("Creating Stereo Depth pipeline")
199pipeline = dai.Pipeline()
200
201camLeft = pipeline.create(dai.node.MonoCamera)
202camRight = pipeline.create(dai.node.MonoCamera)
203stereo = pipeline.create(dai.node.StereoDepth)
204xoutLeft = pipeline.create(dai.node.XLinkOut)
205xoutRight = pipeline.create(dai.node.XLinkOut)
206xoutDisparity = pipeline.create(dai.node.XLinkOut)
207xoutDepth = pipeline.create(dai.node.XLinkOut)
208xoutRectifLeft = pipeline.create(dai.node.XLinkOut)
209xoutRectifRight = pipeline.create(dai.node.XLinkOut)
210
211if args.swap_left_right:
212    camLeft.setCamera("right")
213    camRight.setCamera("left")
214else:
215    camLeft.setCamera("left")
216    camRight.setCamera("right")
217
218for monoCam in (camLeft, camRight):  # Common config
219    monoCam.setResolution(resolution['res'])
220    # monoCam.setFps(20.0)
221
222stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
223stereo.initialConfig.setMedianFilter(median)  # KERNEL_7x7 default
224stereo.setRectifyEdgeFillColor(0)  # Black, to better see the cutout
225stereo.setLeftRightCheck(lrcheck)
226stereo.setExtendedDisparity(extended)
227stereo.setSubpixel(subpixel)
228if args.alpha is not None:
229    stereo.setAlphaScaling(args.alpha)
230    config = stereo.initialConfig.get()
231    config.postProcessing.brightnessFilter.minBrightness = 0
232    stereo.initialConfig.set(config)
233
234xoutLeft.setStreamName("left")
235xoutRight.setStreamName("right")
236xoutDisparity.setStreamName("disparity")
237xoutDepth.setStreamName("depth")
238xoutRectifLeft.setStreamName("rectifiedLeft")
239xoutRectifRight.setStreamName("rectifiedRight")
240
241camLeft.out.link(stereo.left)
242camRight.out.link(stereo.right)
243stereo.syncedLeft.link(xoutLeft.input)
244stereo.syncedRight.link(xoutRight.input)
245stereo.disparity.link(xoutDisparity.input)
246if depth:
247    stereo.depth.link(xoutDepth.input)
248if outRectified:
249    stereo.rectifiedLeft.link(xoutRectifLeft.input)
250    stereo.rectifiedRight.link(xoutRectifRight.input)
251
252streams = ["left", "right"]
253if outRectified:
254    streams.extend(["rectifiedLeft", "rectifiedRight"])
255streams.append("disparity")
256if depth:
257    streams.append("depth")
258
259cvColorMap = cv2.applyColorMap(np.arange(256, dtype=np.uint8), cv2.COLORMAP_JET)
260cvColorMap[0] = [0, 0, 0]
261print("Creating DepthAI device")
262with device:
263    device.startPipeline(pipeline)
264
265    # Create a receive queue for each stream
266    qList = [device.getOutputQueue(stream, 8, blocking=False) for stream in streams]
267
268    while True:
269        for q in qList:
270            name = q.getName()
271            frame = q.get().getCvFrame()
272            if name == "depth":
273                frame = frame.astype(np.uint16)
274            elif name == "disparity":
275                frame = getDisparityFrame(frame, cvColorMap)
276
277            cv2.imshow(name, frame)
278        if cv2.waitKey(1) == ord("q"):
279            break