Sync Capture

SyncCapture.py

import cv2
import depthai as dai
from pathlib import Path
import time
from datetime import datetime, timedelta
import argparse

parser = argparse.ArgumentParser()
parser.add_argument(
    "-ls",
    "--leftSocket",
    type=str,
    default="left",
    help="choose the socket left, right or rgb depending on where the left camera is connected to",
)

parser.add_argument(
    "-rs",
    "--rightSocket",
    type=str,
    default="right",
    help="choose the socket left, right or rgb depending on where the left camera is connected to",
)

parser.add_argument(
    "-m",
    "--mode",
    type=str,
    default="mono",
    help="Default set to \"mono\". Change it to \"color\" to use color cameras for capture",
)

args = parser.parse_args()

camSocket = {'right': dai.CameraBoardSocket.RIGHT,
        'left': dai.CameraBoardSocket.LEFT,
         'rgb': dai.CameraBoardSocket.RGB}

class HostSync:
    def __init__(self, arraySize):
        self.arrays = {}
        self.arraySize = arraySize
    def add_msg(self, name, data, ts):
        if not name in self.arrays:
            self.arrays[name] = []
        # Add msg to array
        self.arrays[name].append({'data': data, 'timestamp': ts})
        # Try finding synced msgs
        synced = {}
        for name, arr in self.arrays.items():
            for i, obj in enumerate(arr):
                time_diff = abs(obj['timestamp'] - ts)
                # 20ms since we add rgb/depth frames at 30FPS => 33ms. If
                # time difference is below 20ms, it's considered as synced
                if time_diff < timedelta(milliseconds=20):
                    synced[name] = obj['data']
                    # print(f"{name}: {i}/{len(arr)}")
                    break
        if len(synced) == self.arraySize:
            def remove(t1, t2):
                return timedelta(milliseconds=500) < abs(t1 - t2)
            # Remove old msgs
            for name, arr in self.arrays.items():
                for i, obj in enumerate(arr):
                    if remove(obj['timestamp'], ts):
                        arr.remove(obj)
                    else: break
            return synced
        return False


mode = args.mode
camType = dai.node.MonoCamera
camRes = dai.MonoCameraProperties.SensorResolution.THE_800_P
if mode == 'color':
    camType = dai.node.ColorCamera
    camRes = dai.ColorCameraProperties.SensorResolution.THE_12_MP
folderName = 'dataset_' + time.strftime("%Y%m%d_%H%M%S")
datasetPath = Path(__file__).parent / folderName
datasetPath.resolve()
# Create pipeline
pipeline = dai.Pipeline()

# Define sources and outputs
camLeft = pipeline.create (camType)
camRight = pipeline.create(camType)
xoutLeft = pipeline.create(dai.node.XLinkOut)
xoutRight = pipeline.create(dai.node.XLinkOut)

xoutLeft.setStreamName('left')
xoutRight.setStreamName('right')
xoutLeft.input.setBlocking(False)
xoutLeft.input.setQueueSize(1)

xoutRight.input.setBlocking(False)
xoutRight.input.setQueueSize(1)

# Properties
camLeft.setBoardSocket(camSocket[args.leftSocket])
camLeft.setResolution(camRes)
camRight.setBoardSocket(camSocket[args.rightSocket])
camRight.setResolution(camRes)

# Linking
if mode == 'color':
    camRight.video.link(xoutRight.input)
    camLeft.video.link(xoutLeft.input)
else:
    camRight.out.link(xoutRight.input)
    camLeft.out.link(xoutLeft.input)


devInfoList = dai.Device.getAllAvailableDevices()
devices = []
devNum = 0

outputQueues = {}
for i, devInfo in enumerate(devInfoList):
    devices.append(dai.Device(pipeline, devInfo))
    qLeft = devices[i].getOutputQueue(name="left", maxSize=4, blocking=False)
    qRight = devices[i].getOutputQueue(name="right", maxSize=4, blocking=False)

    outputQueues[devInfo.getMxId() + '_left' ]  = qLeft
    outputQueues[devInfo.getMxId() + '_right' ] = qRight

sync = HostSync(len(outputQueues))

setCapture = False
captureID = 1

while True:
    msgs = False
    for key in outputQueues.keys():
        if outputQueues[key].has():
            daiImage = outputQueues[key].get()             
            msgs = msgs or sync.add_msg(key, daiImage, daiImage.getTimestamp())
            cv2.imshow(key, daiImage.getCvFrame())

    keyPressed = cv2.waitKey(1)
    if keyPressed == ord('q'):
            break
    elif keyPressed == ord(' '):
        setCapture = True
    
    if msgs and setCapture:
        setCapture = False
        for key in outputQueues.keys():
            image = msgs[key].getCvFrame()
            currDir = datasetPath / key
            currDir.resolve()
            currDir.mkdir(parents=True, exist_ok=True)
            currFileName = str(currDir) + "/capture_" + str(captureID) + '.png'
            print('Capturing images: Count -> {}'.format(captureID))
            cv2.imwrite(currFileName, image)
        captureID += 1