Skip to content

Instantly share code, notes, and snippets.

@franva
Created April 16, 2023 14:36
Show Gist options
  • Save franva/0c57e5c002963f6a9561638f829f3f6e to your computer and use it in GitHub Desktop.
Save franva/0c57e5c002963f6a9561638f829f3f6e to your computer and use it in GitHub Desktop.
Calculate Depth on AICAm
import math
import depthai as dai
import cv2
pipeline = dai.Pipeline()
img_size = 672
def init_depth(pipeline):
# Define sources and outputs
monoLeft = pipeline.create(dai.node.MonoCamera)
monoRight = pipeline.create(dai.node.MonoCamera)
stereo = pipeline.create(dai.node.StereoDepth)
spatialLocationCalculator = pipeline.create(dai.node.SpatialLocationCalculator)
xoutSpatialData = pipeline.create(dai.node.XLinkOut)
xoutSpatialData.setStreamName("spatialData")
# Properties
monoLeft.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
monoLeft.setBoardSocket(dai.CameraBoardSocket.LEFT)
monoRight.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)
monoRight.setBoardSocket(dai.CameraBoardSocket.RIGHT)
lrcheck = True
subpixel = False
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
stereo.setLeftRightCheck(lrcheck)
stereo.setExtendedDisparity(True) if subpixel == False else stereo.setSubpixel(subpixel) # ExtendedDisparity and SubPixel for now is incompatiable.
stereo.setDepthAlign(dai.CameraBoardSocket.RGB)
stereo.setOutputSize(img_size,img_size)
# Create 10*10 ROIs grid which covers the entire image
for i in range(10):
for j in range(10):
config = dai.SpatialLocationCalculatorConfigData()
config.depthThresholds.lowerThreshold = 200
config.depthThresholds.upperThreshold = 10000
config.roi = dai.Rect(dai.Point2f(i*0.1, j * 0.1), dai.Point2f((i+1)*0.1, (j+1)*0.1))
spatialLocationCalculator.initialConfig.addROI(config)
# Linking
monoLeft.out.link(stereo.left)
monoRight.out.link(stereo.right)
stereo.depth.link(spatialLocationCalculator.inputDepth)
spatialLocationCalculator.out.link(xoutSpatialData.input)
def get_depth_data(spatialCalcQueue, nn_size):
spatialData = spatialCalcQueue.get().getSpatialLocations()
depth_matrix_10_10 = [[{} for _ in range(10)] for _ in range(10)]
index = 0
for depthData in spatialData:
roi = depthData.config.roi
roi = roi.denormalize(width=nn_size, height=nn_size)
xmin = int(roi.topLeft().x)
ymin = int(roi.topLeft().y)
xmax = int(roi.bottomRight().x)
ymax = int(roi.bottomRight().y)
coords = depthData.spatialCoordinates
distance = math.sqrt(coords.x ** 2 + coords.y ** 2 + coords.z ** 2)
row = index // 10
col = index % 10
index += 1
depth_matrix_10_10[row][col] = {
"distance": distance,
"topLeft": (xmin, ymin),
"bottomRight": (xmax, ymax)
}
return depth_matrix_10_10
def draw_depth(frame, depth_matrix):
color = (0,200,40)
font_color = (0,0,200)
fontType = cv2.FONT_HERSHEY_TRIPLEX
for i in range(10):
for j in range(10):
distance, topLeft, bottomRight = depth_matrix[i][j].values()
cv2.rectangle(frame, topLeft[::-1], bottomRight[::-1], color, thickness=2)
cv2.putText(frame, "{:.1f}".format(distance/1000), (topLeft[0],topLeft[1] + 30), fontType, 1, font_color)
return frame
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_12_MP)
cam_rgb.setIspScale(5,16)
cam_rgb.setPreviewSize(img_size,img_size)
cam_rgb.setInterleaved(False)
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName('rgb')
cam_rgb.preview.link(xout_rgb.input)
init_depth(pipeline)
with dai.Device(pipeline) as device:
q_rgb = device.getOutputQueue('rgb', maxSize=1, blocking=False)
q_spatialData = device.getOutputQueue('spatialData', maxSize=1, blocking=False)
while True:
depth_matrix = get_depth_data(q_spatialData, img_size)
raw_data_rgb = q_rgb.tryGet()
if raw_data_rgb is not None:
frame = raw_data_rgb.getCvFrame()
frame = draw_depth(frame, depth_matrix)
cv2.imshow("show", frame)
if cv2.waitKey(1) == ord('q'):
break
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment