multi_static.py

import os
import os.path as osp
import time

import cv2
import pyrealsense2 as rs
import logging
import numpy as np
import json
import yaml


def enumerate_connected_devices(context):
    """
    Enumerate the connected Intel RealSense devices

    Parameters:
    -----------
    context 	   : rs.context()
                     The context created for using the realsense library

    Return:
    -----------
    connect_device : array
                     Array of (serial, product-line) tuples of devices which are connected to the PC

    """
    connect_device = []

    for d in context.devices:
        if d.get_info(rs.camera_info.name).lower() != 'platform camera':
            serial = d.get_info(rs.camera_info.serial_number)
            product_line = d.get_info(rs.camera_info.product_line)
            device_info = (serial, product_line) # (serial_number, product_line)
            connect_device.append( device_info )
    return connect_device



# No. 1 step
def main(save_path: str):

    logging.basicConfig(level=logging.INFO)

    # Realsense
    context = rs.context()
    valid_devices = enumerate_connected_devices(context)
    print('Find {} valid devices!'.format(len(valid_devices)))

    intrinsics_dict = dict()
    realsense_cameras = []
    for idx, device in enumerate(valid_devices):
        print('[{}] {}:{}'.format(idx, device[0], device[1]))
        # Configure depth and color streams
        pipeline = rs.pipeline()
        config = rs.config()
        device_serial = device[0]
        device_product_line = device[1]
        if device[1] == 'L500':
            config.enable_stream(rs.stream.depth, 1024, 768, rs.format.z16, 30)
        else:
            # D400
            config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
        config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
        # Start streaming
        config.enable_device(device_serial)
        init_profile = pipeline.start(config)
        # Set preset
        depth_sensor = init_profile.get_device().first_depth_sensor()
        preset_range = depth_sensor.get_option_range(rs.option.visual_preset)
        for i in range(int(preset_range.max) + 1):
            visulpreset = depth_sensor.get_option_value_description(rs.option.visual_preset, i)
            if visulpreset == "Short Range":
                depth_sensor.set_option(rs.option.visual_preset, i)
                print('Changed to Short Range')
        # change exposure
        sensor = init_profile.get_device().query_sensors()[1]
        sensor.set_option(rs.option.enable_auto_exposure, 1)
        intrinsics = init_profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()
        intrinsics_matrix = np.array([
            [intrinsics.fx, 0., intrinsics.ppx],
            [0., intrinsics.fy, intrinsics.ppy],
            [0., 0., 1.]
        ])
        print(f'intrinsics matrix: {intrinsics_matrix}')

        align = rs.align(rs.stream.color)  # align to color is better than to depth

        realsense_cameras.append((device[0][-2:], (pipeline, align)))
        intrinsics_dict['r' + str(device[0][-2:])] = intrinsics_matrix.tolist()

        device_save_path = os.path.join(save_path, 'r' + device[0][-2:])
        if not os.path.lexists(device_save_path):
            os.makedirs(device_save_path)
        device_save_path = os.path.join(save_path, 'r' + device[0][-2:] + '_depth')
        if not os.path.lexists(device_save_path):
            os.makedirs(device_save_path)

    intrinsics_save_path = os.path.join(save_path, 'intrinsics_all.json')
    with open(intrinsics_save_path, 'w') as f:
        json.dump(intrinsics_dict, f)
        print('Saving intrinsics matrix to {}!'.format(intrinsics_save_path))

    capture_count = 0
    try:
        while True:
            command = input('Press Enter to capture, q to exit.')
            if command == '' or command == 'c':
                # Realsense
                for cname, (pipe, align) in realsense_cameras:
                    # Wait for a coherent pair of frames: depth and color
                    frames = pipe.wait_for_frames()  # @fixme: will get multiple repeated depth images for r22 camera
                    # align
                    frames = align.process(frames)
                    depth_frame = frames.get_depth_frame()
                    color_frame = frames.get_color_frame()
                    depth_image = np.asanyarray(depth_frame.get_data())
                    color_image = np.asanyarray(color_frame.get_data())
                    fname = os.path.join(save_path, 'r'+cname, f"{capture_count}.jpg")
                    cv2.imwrite(fname, color_image)
                    fname = os.path.join(save_path,'r'+cname+'_depth', f"{capture_count}.png")
                    cv2.imwrite(fname, (depth_image / 4).astype(np.uint16))
                    print('-------', cname + 'OK')

                    # time.sleep(0.03)

                capture_count += 1

            elif command == 'q':
                raise RuntimeError('Quit!')

            else:
                print(command, 'not found')
    except Exception as e:
        print(e)

    finally:
        for cname, (pipe, align) in realsense_cameras:
            pipe.stop()
            print(cname, 'closed')
        # quit()


if __name__ == '__main__':
    # 修改这里，指定保存路径，会生成./records/rXX/ 和 ./records/rXX_depth/ 两个文件夹，分别存放彩色图和深度图，XX为设备序号后两位，如r85.
    SAVE_PATH = './records' 
    main(SAVE_PATH)