|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import cv2 |
| 4 | +import numpy as np |
| 5 | +import depthai as dai |
| 6 | +from time import sleep |
| 7 | +import datetime |
| 8 | +import argparse |
| 9 | +from pathlib import Path |
| 10 | + |
| 11 | +datasetDefault = str((Path(__file__).parent / Path('models/dataset')).resolve().absolute()) |
| 12 | +parser = argparse.ArgumentParser() |
| 13 | +parser.add_argument('-dataset', nargs='?', help="Path to recorded frames", default=datasetDefault) |
| 14 | +args = parser.parse_args() |
| 15 | + |
| 16 | +if not Path(datasetDefault).exists(): |
| 17 | + import sys |
| 18 | + raise FileNotFoundError(f'Required file/s not found, please run "{sys.executable} install_requirements.py"') |
| 19 | + |
| 20 | + |
| 21 | +# StereoDepth config options. |
| 22 | +out_depth = False # Disparity by default |
| 23 | +out_rectified = True # Output and display rectified streams |
| 24 | +lrcheck = True # Better handling for occlusions |
| 25 | +extended = False # Closer-in minimum depth, disparity range is doubled |
| 26 | +subpixel = True # Better accuracy for longer distance, fractional disparity 32-levels |
| 27 | +median = dai.StereoDepthProperties.MedianFilter.KERNEL_7x7 |
| 28 | + |
| 29 | +# Sanitize some incompatible options |
| 30 | +if lrcheck or extended or subpixel: |
| 31 | + median = dai.StereoDepthProperties.MedianFilter.MEDIAN_OFF |
| 32 | + |
| 33 | +print("StereoDepth config options: ") |
| 34 | +print("Left-Right check: ", lrcheck) |
| 35 | +print("Extended disparity: ", extended) |
| 36 | +print("Subpixel: ", subpixel) |
| 37 | +print("Median filtering: ", median) |
| 38 | + |
| 39 | +right_intrinsic = [[860.0, 0.0, 640.0], [0.0, 860.0, 360.0], [0.0, 0.0, 1.0]] |
| 40 | + |
| 41 | + |
| 42 | +def create_stereo_depth_pipeline(): |
| 43 | + print("Creating Stereo Depth pipeline: ", end='') |
| 44 | + |
| 45 | + print("XLINK IN -> STEREO -> XLINK OUT") |
| 46 | + pipeline = dai.Pipeline() |
| 47 | + |
| 48 | + camLeft = pipeline.createXLinkIn() |
| 49 | + camRight = pipeline.createXLinkIn() |
| 50 | + stereo = pipeline.createStereoDepth() |
| 51 | + xoutLeft = pipeline.createXLinkOut() |
| 52 | + xoutRight = pipeline.createXLinkOut() |
| 53 | + xoutDepth = pipeline.createXLinkOut() |
| 54 | + xoutDisparity = pipeline.createXLinkOut() |
| 55 | + xoutRectifLeft = pipeline.createXLinkOut() |
| 56 | + xoutRectifRight = pipeline.createXLinkOut() |
| 57 | + |
| 58 | + camLeft.setStreamName('in_left') |
| 59 | + camRight.setStreamName('in_right') |
| 60 | + |
| 61 | + stereo.setOutputDepth(out_depth) |
| 62 | + stereo.setOutputRectified(out_rectified) |
| 63 | + stereo.setConfidenceThreshold(200) |
| 64 | + stereo.setRectifyEdgeFillColor(0) # Black, to better see the cutout |
| 65 | + stereo.setMedianFilter(median) # KERNEL_7x7 default |
| 66 | + stereo.setLeftRightCheck(lrcheck) |
| 67 | + stereo.setExtendedDisparity(extended) |
| 68 | + stereo.setSubpixel(subpixel) |
| 69 | + |
| 70 | + stereo.setEmptyCalibration() # Set if the input frames are already rectified |
| 71 | + stereo.setInputResolution(1280, 720) |
| 72 | + |
| 73 | + xoutLeft.setStreamName('left') |
| 74 | + xoutRight.setStreamName('right') |
| 75 | + xoutDepth.setStreamName('depth') |
| 76 | + xoutDisparity.setStreamName('disparity') |
| 77 | + xoutRectifLeft.setStreamName('rectified_left') |
| 78 | + xoutRectifRight.setStreamName('rectified_right') |
| 79 | + |
| 80 | + camLeft.out.link(stereo.left) |
| 81 | + camRight.out.link(stereo.right) |
| 82 | + stereo.syncedLeft.link(xoutLeft.input) |
| 83 | + stereo.syncedRight.link(xoutRight.input) |
| 84 | + stereo.depth.link(xoutDepth.input) |
| 85 | + stereo.disparity.link(xoutDisparity.input) |
| 86 | + stereo.rectifiedLeft.link(xoutRectifLeft.input) |
| 87 | + stereo.rectifiedRight.link(xoutRectifRight.input) |
| 88 | + |
| 89 | + streams = ['left', 'right'] |
| 90 | + if out_rectified: |
| 91 | + streams.extend(['rectified_left', 'rectified_right']) |
| 92 | + streams.extend(['disparity', 'depth']) |
| 93 | + |
| 94 | + return pipeline, streams |
| 95 | + |
| 96 | + |
| 97 | +def convert_to_cv2_frame(name, image): |
| 98 | + baseline = 75 #mm |
| 99 | + focal = right_intrinsic[0][0] |
| 100 | + max_disp = 96 |
| 101 | + disp_type = np.uint8 |
| 102 | + disp_levels = 1 |
| 103 | + if (extended): |
| 104 | + max_disp *= 2 |
| 105 | + if (subpixel): |
| 106 | + max_disp *= 32 |
| 107 | + disp_type = np.uint16 |
| 108 | + disp_levels = 32 |
| 109 | + |
| 110 | + data, w, h = image.getData(), image.getWidth(), image.getHeight() |
| 111 | + if name == 'depth': |
| 112 | + # this contains FP16 with (lrcheck or extended or subpixel) |
| 113 | + frame = np.array(data).astype(np.uint8).view(np.uint16).reshape((h, w)) |
| 114 | + elif name == 'disparity': |
| 115 | + disp = np.array(data).astype(np.uint8).view(disp_type).reshape((h, w)) |
| 116 | + |
| 117 | + # Compute depth from disparity |
| 118 | + with np.errstate(divide='ignore'): |
| 119 | + depth = (disp_levels * baseline * focal / disp).astype(np.uint16) |
| 120 | + |
| 121 | + if 1: # Optionally, extend disparity range to better visualize it |
| 122 | + frame = (disp * 255. / max_disp).astype(np.uint8) |
| 123 | + |
| 124 | + if 1: # Optionally, apply a color map |
| 125 | + frame = cv2.applyColorMap(frame, cv2.COLORMAP_HOT) |
| 126 | + |
| 127 | + else: # mono streams / single channel |
| 128 | + frame = np.array(data).reshape((h, w)).astype(np.uint8) |
| 129 | + if name.startswith('rectified_'): |
| 130 | + frame = cv2.flip(frame, 1) |
| 131 | + if name == 'rectified_right': |
| 132 | + last_rectif_right = frame |
| 133 | + return frame |
| 134 | + |
| 135 | + |
| 136 | +pipeline, streams = create_stereo_depth_pipeline() |
| 137 | + |
| 138 | +print("Creating DepthAI device") |
| 139 | +with dai.Device(pipeline) as device: |
| 140 | + print("Starting pipeline") |
| 141 | + device.startPipeline() |
| 142 | + |
| 143 | + inStreams = ['in_right', 'in_left'] |
| 144 | + inStreamsCameraID = [dai.CameraBoardSocket.RIGHT, dai.CameraBoardSocket.LEFT] |
| 145 | + in_q_list = [] |
| 146 | + for s in inStreams: |
| 147 | + q = device.getInputQueue(s) |
| 148 | + in_q_list.append(q) |
| 149 | + |
| 150 | + # Create a receive queue for each stream |
| 151 | + q_list = [] |
| 152 | + for s in streams: |
| 153 | + q = device.getOutputQueue(s, 8, blocking=False) |
| 154 | + q_list.append(q) |
| 155 | + |
| 156 | + # Need to set a timestamp for input frames, for the sync stage in Stereo node |
| 157 | + timestamp_ms = 0 |
| 158 | + index = 0 |
| 159 | + while True: |
| 160 | + # Handle input streams, if any |
| 161 | + if in_q_list: |
| 162 | + dataset_size = 2 # Number of image pairs |
| 163 | + frame_interval_ms = 500 |
| 164 | + for i, q in enumerate(in_q_list): |
| 165 | + path = args.dataset + '/' + str(index) + '/' + q.getName() + '.png' |
| 166 | + data = cv2.imread(path, cv2.IMREAD_GRAYSCALE).reshape(720*1280) |
| 167 | + tstamp = datetime.timedelta(seconds = timestamp_ms // 1000, |
| 168 | + milliseconds = timestamp_ms % 1000) |
| 169 | + img = dai.ImgFrame() |
| 170 | + img.setData(data) |
| 171 | + img.setTimestamp(tstamp) |
| 172 | + img.setInstanceNum(inStreamsCameraID[i]) |
| 173 | + img.setWidth(1280) |
| 174 | + img.setHeight(720) |
| 175 | + q.send(img) |
| 176 | + if timestamp_ms == 0: # Send twice for first iteration |
| 177 | + q.send(img) |
| 178 | + print("Sent frame: {:25s}".format(path), 'timestamp_ms:', timestamp_ms) |
| 179 | + timestamp_ms += frame_interval_ms |
| 180 | + index = (index + 1) % dataset_size |
| 181 | + sleep(frame_interval_ms / 1000) |
| 182 | + # Handle output streams |
| 183 | + for q in q_list: |
| 184 | + if q.getName() in ['left', 'right', 'depth']: continue |
| 185 | + frame = convert_to_cv2_frame(q.getName(), q.get()) |
| 186 | + cv2.imshow(q.getName(), frame) |
| 187 | + if cv2.waitKey(1) == ord('q'): |
| 188 | + break |
0 commit comments