diff --git a/modules/python/examples/realsense_mbt.py b/modules/python/examples/realsense_mbt.py new file mode 100644 index 0000000000..53b6013aa3 --- /dev/null +++ b/modules/python/examples/realsense_mbt.py @@ -0,0 +1,228 @@ +import argparse +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional +import numpy as np +import time +import faulthandler +faulthandler.enable() + + +from visp.core import XmlParserCamera, CameraParameters, ColVector, HomogeneousMatrix, Display, ImageConvert +from visp.core import ImageGray, ImageUInt16 +from visp.io import ImageIo +from visp.mbt import MbGenericTracker, MbTracker +from visp.gui import DisplayOpenCV +from visp.core import Color +from visp.core import PixelMeterConversion + +import pyrealsense2 as rs + + +try: + import cv2 +except: + print('Could not import opencv python! make sure that it is installed as it is required') + import sys + sys.exit(1) + +import matplotlib.pyplot as plt + + +class MBTModelData: + def __init__(self, data_root: Path): + model_path = data_root / 'model' / 'teabox' + assert model_path.exists() + self.config_color = model_path / 'teabox_color.xml' + self.config_depth = model_path / 'teabox_depth.xml' + self.cad_file = model_path / 'teabox.cao' + self.init_file = model_path / 'teabox.init' + + +class MBTConfig: + def __init__(self, data_root: Path): + data_path = data_root / 'data' / 'teabox' + assert data_path.exists() + self.color_camera_name = 'Camera_L' + self.depth_camera_name = 'Camera_R' + + self.color_intrinsics_file = data_path / f'{self.color_camera_name}.xml' + self.depth_intrinsics_file = data_path / f'{self.depth_camera_name}.xml' + + self.color_images_dir = data_path / 'color' + self.depth_images_dir = data_path / 'depth' + self.ground_truth_dir = data_path / 'ground-truth' + + + self.depth_intrinsics_file = data_path / f'{self.depth_camera_name}.xml' + + self.extrinsic_file = str(data_path / 'depth_M_color.txt') + # self.ground_truth = str(data_root / 'data' / 'depth_M_color.txt') + +@dataclass +class FrameData: + I: ImageGray + I_depth: Optional[ImageUInt16] + point_cloud: Optional[np.ndarray] + cMo_ground_truth: HomogeneousMatrix + + + + +def read_data(exp_config: MBTConfig, cam_depth: CameraParameters | None, I: ImageGray): + pipe = rs.pipeline() + profile = pipe.start() + + use_depth = cam_depth is not None + iteration = 1 + while True: + frames = pipe.wait_for_frames() + print(frames) + I_depth_raw = np.asanyarray(frames.get_depth_frame().as_frame().get_data()) + + print(I_depth_raw.shape) + I = np.asanyarray(frames.get_color_frame().as_frame().get_data()) + print(I.shape) + point_cloud = None + if use_depth: + t = time.time() + depth_filepath = exp_config.depth_images_dir / depth_format.format(iteration) + if not depth_filepath.exists(): + print(f'Could not find image {depth_filepath}') + return + I_depth_np = cv2.imread(str(depth_filepath), cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH) + I_depth_np = I_depth_np[..., 0] + print(f'\tDepth load took {(time.time() - t) * 1000}ms') + I_depth_raw = ImageUInt16(I_depth_np * 32767.5) + if I_depth_np.size == 0: + print('Could not successfully read the depth image') + return + t = time.time() + # point_cloud = np.empty((*I_depth_np.shape, 3), dtype=np.float64) + Z = I_depth_np.copy() + Z[Z > 2] = 0.0 # Clamping values that are too high + + vs, us = np.meshgrid(range(I_depth_np.shape[0]), range(I_depth_np.shape[1]), indexing='ij') + xs, ys = PixelMeterConversion.convertPoints(cam_depth, us, vs) + point_cloud = np.stack((xs * Z, ys * Z, Z), axis=-1) + + print(f'\tPoint_cloud took {(time.time() - t) * 1000}ms') + + cMo_ground_truth = HomogeneousMatrix() + ground_truth_file = exp_config.ground_truth_dir / (exp_config.color_camera_name + '_{:04d}.txt'.format(iteration)) + cMo_ground_truth.load(str(ground_truth_file)) + iteration += 1 + end_parse_time = time.time() + print(f'Data parsing took: {(end_parse_time - start_parse_time) * 1000}ms') + yield FrameData(I, I_depth_raw, point_cloud, cMo_ground_truth) + + + +def parse_camera_file(exp_config: MBTConfig, is_color: bool) -> CameraParameters: + cam = CameraParameters() + xml_parser = XmlParserCamera() + if is_color: + camera_name, file_path = exp_config.color_camera_name, exp_config.color_intrinsics_file + else: + camera_name, file_path = exp_config.depth_camera_name, exp_config.depth_intrinsics_file + parse_res = xml_parser.parse(cam, str(file_path), camera_name, + CameraParameters.perspectiveProjWithoutDistortion, 0, 0, True) + assert parse_res == XmlParserCamera.SEQUENCE_OK # Check result + return cam + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--data-root', type=str, required=True, + help='Path to the folder containing all the data for the synthetic MBT example') + parser.add_argument('--display-ground-truth', action='store_true') + parser.add_argument('--disable-klt', action='store_true') + parser.add_argument('--disable-depth', action='store_true') + parser.add_argument('--step-by-step', action='store_true') + parser.add_argument('--init-ground-truth', action='store_true') + + + args = parser.parse_args() + data_root = Path(args.data_root) + + mbt_model = MBTModelData(data_root) + exp_config = MBTConfig(data_root) + + assert data_root.exists() and data_root.is_dir() + + rgb_tracker: int = MbGenericTracker.EDGE_TRACKER | (MbGenericTracker.KLT_TRACKER if not args.disable_klt else 0) + tracker_types: List[int] = [rgb_tracker] + if not args.disable_depth: + depth_tracker = MbGenericTracker.DEPTH_DENSE_TRACKER + tracker_types.append(depth_tracker) + + tracker = MbGenericTracker(tracker_types) + + if args.disable_depth: + tracker.loadConfigFile(str(mbt_model.config_color)) + else: + tracker.loadConfigFile(str(mbt_model.config_color), str(mbt_model.config_depth)) + tracker.loadModel(str(mbt_model.cad_file)) + + # Camera intrinsics + cam_color = parse_camera_file(exp_config, True) + cam_depth = parse_camera_file(exp_config, False) if not args.disable_depth else None + + tracker.setCameraParameters(*((cam_color,) if args.disable_depth else (cam_color, cam_depth))) + tracker.setDisplayFeatures(True) + + print('Color intrinsics:', cam_color) + print('Depth intrinsics:', cam_depth) + I = ImageGray() + data_generator = read_data(exp_config, cam_depth, I) + frame_data = next(data_generator) # Get first frame for init + + depth_M_color = HomogeneousMatrix() + if not args.disable_depth: + depth_M_color.load(exp_config.extrinsic_file) + tracker.setCameraTransformationMatrix('Camera2', depth_M_color) + + # Initialize displays + dI = DisplayOpenCV() + dI.init(I, 0, 0, 'Color image') + + I_depth = None if args.disable_depth else ImageGray() + dDepth = DisplayOpenCV() + if not args.disable_depth: + ImageConvert.createDepthHistogram(frame_data.I_depth, I_depth) + dDepth.init(I_depth, I.getWidth(), 0, 'Depth') + + if args.init_ground_truth: + tracker.initFromPose(I, frame_data.cMo_ground_truth) + else: + tracker.initClick(I, str(mbt_model.init_file)) + start_time = time.time() + for frame_data in data_generator: + if frame_data.I_depth is not None: + ImageConvert.createDepthHistogram(frame_data.I_depth, I_depth) + + Display.display(I) + if not args.disable_depth: + Display.display(I_depth) + + if args.disable_depth: + tracker.track(I=I) + else: + pc = frame_data.point_cloud + image_dict = { + 'Camera1': I + } + t = time.time() + tracker.track(image_dict, {'Camera2': pc}) + print(f'Tracking took {(time.time() - t) * 1000}ms') + cMo = HomogeneousMatrix() + tracker.getPose(cMo) + + Display.displayFrame(I, cMo, cam_color, 0.05, Color.none, 2) + tracker.display(I, cMo, cam_color, Color.red, 2) + Display.flush(I) + if not args.disable_depth: + Display.flush(I_depth) + if args.step_by_step: + Display.getKeyboardEvent(I, blocking=True) + end_time = time.time() + print(f'total time = {end_time - start_time}s') diff --git a/modules/python/examples/synthetic_data_mbt.py b/modules/python/examples/synthetic_data_mbt.py index d140b79d89..e1fbcb506e 100644 --- a/modules/python/examples/synthetic_data_mbt.py +++ b/modules/python/examples/synthetic_data_mbt.py @@ -224,4 +224,4 @@ def parse_camera_file(exp_config: MBTConfig, is_color: bool) -> CameraParameters if args.step_by_step: Display.getKeyboardEvent(I, blocking=True) end_time = time.time() - print(f'total time = {end_time - start_time}') + print(f'total time = {end_time - start_time}s')