realsense-rbt.py

 
import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple
import numpy as np
import time
import faulthandler
faulthandler.enable()
 
from visp.core import CameraParameters, HomogeneousMatrix
from visp.core import Color, Display, ImageConvert
from visp.core import ImageGray, ImageUInt16, ImageRGBa, ImageFloat
from visp.io import ImageIo
from visp.rbt import RBTracker, RBFeatureTracker, RBFeatureTrackerInput
from visp.python.display_utils import get_display
from visp.python.rbt import PythonRBExtensions
 
import pyrealsense2 as rs
import matplotlib.pyplot as plt
 
class PyBaseFeatureTracker(RBFeatureTracker):
  '''
  The base structure that a class should have to implement a render based feature tracker
  '''
  def __init__(self):
    RBFeatureTracker.__init__(self)
 
  def requiresRGB(self) -> bool:
    return False
  def requiresDepth(self) -> bool:
    return False
  def requiresSilhouetteCandidates(self) -> bool:
    return False
 
  def onTrackingIterStart(self, frame: RBFeatureTrackerInput, cMo: HomogeneousMatrix):
    self.cov.resize(6, 6)
    self.LTL.resize(6, 6)
    self.LTR.resize(6)
    self.numFeatures = 0
 
  def extractFeatures(self, frame: RBFeatureTrackerInput, previousFrame: RBFeatureTrackerInput, cMo: HomogeneousMatrix):
    pass
 
  def trackFeatures(self, frame: RBFeatureTrackerInput, previousFrame: RBFeatureTrackerInput, cMo: HomogeneousMatrix):
    pass
 
  def initVVS(self, frame: RBFeatureTrackerInput, previousFrame: RBFeatureTrackerInput, cMo: HomogeneousMatrix):
    print('INITVVS Was called')
    pass
 
  def computeVVSIter(self, frame: RBFeatureTrackerInput, cMo: HomogeneousMatrix, iteration: int):
    pass
 
  def onTrackingIterEnd(self, cMo: HomogeneousMatrix):
    pass
 
  def display(self, cam: CameraParameters, I: ImageGray, IRGB: ImageRGBa, I_depth: ImageGray):
    pass
 
@dataclass
class FrameData:
  I: ImageGray
  IRGB: ImageRGBa
  I_depth: Optional[ImageFloat]
 
 
def read_data(depth_scale: Optional[float], IRGB: ImageRGBa, I: ImageGray, pipe: rs.pipeline):
  use_depth = depth_scale is not None
  iteration = 1
  align_to = rs.align(rs.stream.color)
  while True:
    frames = pipe.wait_for_frames()
    frames = align_to.process(frames)
    I_np = np.asanyarray(frames.get_color_frame().as_frame().get_data())
    I_np = np.concatenate((I_np, np.ones_like(I_np[..., 0:1], dtype=np.uint8)), axis=-1)
    IRGB.resize(I_np.shape[0], I_np.shape[1])
    I_rgba_ref = IRGB.numpy()
    I_rgba_ref[...] = I_np
    ImageConvert.convert(IRGB, I, 0)
    I_depth_float = None
    if use_depth:
      I_depth_raw = np.asanyarray(frames.get_depth_frame().as_frame().get_data())
      I_depth_float = I_depth_raw.astype(np.float32) * depth_scale
    iteration += 1
    yield FrameData(I, IRGB, ImageFloat(I_depth_float))
 
 
def cam_from_rs_profile(profile) -> Tuple[CameraParameters, int, int]:
  intr = profile.as_video_stream_profile().get_intrinsics() # Downcast to video_stream_profile and fetch intrinsics
  return CameraParameters(intr.fx, intr.fy, intr.ppx, intr.ppy), intr.height, intr.width
 
if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--tracker', type=str, required=True,
                      help='Path to the json file containing the tracker configuration.')
  parser.add_argument('--model', type=str, required=False,
                      help='Path to the .obj/.bam file describing the CAD model.')
 
  args = parser.parse_args()
  tracker_path: str = args.tracker
  assert Path(tracker_path).exists(), 'Tracker file not found'
  model_path = args.model
  if model_path is not None:
    assert Path(model_path).exists(), '3D CAD model file not found'
 
  # Initialize realsense2
  pipe = rs.pipeline()
  config = rs.config()
  config.enable_stream(rs.stream.depth, 848, 480, rs.format.z16, 60)
  config.enable_stream(rs.stream.color, 848, 480, rs.format.rgb8, 60)
 
  cfg = pipe.start(config)
  depth_scale = cfg.get_device().first_depth_sensor().get_depth_scale()
 
 
  tracker = RBTracker()
 
  tracker.loadConfigurationFile(tracker_path)
  extensions = PythonRBExtensions()
  extensions.parse_python_extensions(tracker, Path(tracker_path))
  if model_path is not None:
    tracker.setModelPath(model_path)
 
  custom_feature = PyBaseFeatureTracker()
  tracker.addTracker(custom_feature)
 
  cam_color, color_height, color_width = cam_from_rs_profile(cfg.get_stream(rs.stream.color))
 
  tracker.setCameraParameters(cam_color, color_height, color_width)
 
  # Camera intrinsics
 
  print('Color intrinsics:', cam_color)
  I = ImageGray()
  IRGB = ImageRGBa()
  I_depth_display = ImageGray()
  data_generator = read_data(depth_scale, IRGB, I, pipe)
  frame_data = next(data_generator) # Get first frame for init
 
  # Initialize displays
  dI = get_display()
  dI.init(I, 0, 0, 'Color image')
 
  dRGB = get_display()
  dRGB.init(IRGB, I.getWidth(), 0, 'Color image')
 
  I_depth = ImageGray()
  dDepth = get_display()
 
  ImageConvert.createDepthHistogram(frame_data.I_depth, I_depth)
  dDepth.init(I_depth,  I.getWidth() * 2, 0, 'Depth')
 
  for frame in data_generator:
    Display.display(I)
    Display.displayText(I, 50, 0, 'Click to initialize tracking', Color.red)
    Display.flush(I)
    Display.display(IRGB)
    Display.flush(IRGB)
    event = Display.getClick(I, blocking=False)
    if event:
      break
  tracker.startTracking()
  tracker.initClick(I, tracker_path.replace('.json', '.init'), True)
  start_time =  time.time()
  for frame_data in data_generator:
    if frame_data.I_depth is not None:
      I_depth_np = I_depth.numpy()
      I_depth_np[...] = ((np.minimum(frame_data.I_depth, 0.5) / 0.5) * 255.0).astype(np.uint8)
 
    displayed = [I, IRGB, I_depth]
 
    for display_image in displayed:
      Display.display(display_image)
    Display.displayText(I, 50, 0, 'Click to stop tracking', Color.red)
 
    # if args.disable_depth:
    #   tracker.track(I=I, IRGB=IRGB)
    # else:
    t1 = time.time()
    tracker.track(I=frame.I, IRGB=frame_data.IRGB, depth=frame_data.I_depth)
    tracking_time = np.round((time.time() - t1) * 1000.0, 2)
 
    Display.displayText(I, 60, 0, f'Tracking time: {tracking_time}', Color.red)
    cMo = HomogeneousMatrix()
    tracker.getPose(cMo)
 
    tracker.display(I, IRGB, I_depth)
    Display.displayFrame(I, cMo, cam_color, 0.05, Color.none, 2)
 
 
    for display_image in displayed:
      Display.flush(display_image)
 
    event = Display.getClick(I, blocking=False)
    if event:
      break
  end_time = time.time()
  print(f'total time = {end_time - start_time}s')