Source code for concepts.hw_interface.realsense.visualizer

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : visualizer.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 09/12/2022
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

import os
import numpy as np
import cv2
import jacinle.io as io

from datetime import datetime
from enum import Enum
from typing import TYPE_CHECKING, Optional, Tuple, List

if TYPE_CHECKING:
    from concepts.hw_interface.realsense.device import RealSenseDevice

__all__ = ['WindowEvent', 'RealSenseVisualizer', 'visualize_devices', 'run_4corner_calibration']


[docs]class WindowEvent(Enum): EXIT = "exit" SAVE = "save" NONE = "none"
[docs]class RealSenseVisualizer(object):
[docs] def __init__(self, device): self.device: 'RealSenseDevice' = device
[docs] def visualize(self, save_dir: str = "", save_image: bool = False) -> WindowEvent: """Visualize color and depth images in a cv2 window. - Terminates when 'esc' or 'q' key is pressed. - Saves an image when the 's' key is pressed or if the 'save_image' flag is specified. - Images are saved to the specified save_dir, which we assume to already exist. Returns: WindowEvent indicating status of cv2. """ color_image, depth_image = self.device.capture_images() # Form heatmap for depth and stack with color image depth_colormap = cv2.applyColorMap( cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET ) images = np.hstack((color_image, depth_colormap)) window_name = str(self) cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.imshow(window_name, images) cv2.setMouseCallback(window_name, self.register_point) key = cv2.waitKey(1) # Exit on 'esc' or 'q' if key & 0xFF == ord("q") or key == 27: cv2.destroyWindow(window_name) return WindowEvent.EXIT # Save images if 's' key is pressed if key == 115 or save_image: timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") depth_fname = os.path.join( save_dir, f"{timestamp}-{self.device.serial_number}-depth.png" ) color_fname = os.path.join( save_dir, f"{timestamp}-{self.device.serial_number}-color.png" ) cv2.imwrite(depth_fname, depth_image) cv2.imwrite(color_fname, color_image) print(f"Saved depth image for {self} to {depth_fname}") print(f"Saved color image to {self} to {color_fname}") return WindowEvent.SAVE return WindowEvent.NONE
[docs]def visualize_devices(devices: List['RealSenseDevice'], save_dir: str = "") -> None: """Visualizes all the devices in a cv2 window. Press 'q' or 'esc' on any of the windows to exit the infinite loop. Press the 's' key in a specific window to save the color and depth image to disk. You can use a similar loop interface in other places where you need a live camera feed (e.g. collecting demonstrations). """ while True: should_exit = False save = False for device in devices: window_event = RealSenseVisualizer(device).visualize(save_dir, save_image=save) if window_event == WindowEvent.SAVE: # We use this to propagate a save command across all windows save = True # Exit all windows if window_event == WindowEvent.EXIT: should_exit = True break if should_exit: print("Exit key pressed.") cv2.destroyAllWindows() break
[docs]def run_4corner_calibration(devices: List['RealSenseDevice'], save_dir: str = "") -> None: while True: should_exit = False for device in devices: window_event = RealSenseVisualizer(device).visualize(save_dir, save_image=False) # Exit all windows if window_event == WindowEvent.EXIT: should_exit = True break if should_exit: print("Exit key pressed.") cv2.destroyAllWindows() break
def _main(): import jacinle from concepts.hw_interface.realsense.device import start_pipelines, stop_pipelines parser = jacinle.JacArgumentParser() parser.add_argument('--device', choices=['l515', 'd435'], default='d435') args = parser.parse_args() # Detect devices, start pipelines, visualize, and stop pipelines devices = RealSenseDevice.find_devices(args.device) start_pipelines(devices) try: visualize_devices(devices) finally: stop_pipelines(devices) if __name__ == "__main__": _main()