XFeatVisualOdometry.py

import numpy as np
from enum import Enum
import torch
 
from visp.core import Matrix
from visp.core import CameraParameters, HomogeneousMatrix, PixelMeterConversion, ArrayInt2D
from visp.core import Color, Display
from visp.core import ImageGray, ImageRGBa
from visp.rbt import RBFeatureTrackerInput
from visp.rbt import RBVisualOdometry, RBVisualOdometryUtils, LevenbergMarquardtParameters
 
from visp.python.rbt import TrackedDescriptorMap
from visp.python.rbt.xfeat import XFeatTrackingBackend
 
 
class XFeatVisualOdometry(RBVisualOdometry):
  """
  Visual odometry implementation based on XFeat keypoints.
  The minimised error is either the reprojection or the 3D error (in meters)
  Descriptors are associated to 3D points that are stored in a map.
  This map is updated every iteration, with mismatched points being removed and novel keypoints being added.
 
  Every iteration, keypoints in the current RGB frame that are matched with points in the map are used to
  update the camera pose by minimizing the chosen error criterion.
 
  This implementation requires an RGB-D camera to compute the actual 3D positions of newly detected keypoints.
 
  The parameters are the following:
  - "numPoints" Number of points to store in the map
  - "reprojectionThreshold": Distance (in pixels) between observed keypoints and model keypoints above which a map point is considered an outlier and is removed.
  - "minDistNewPoint": Minimum distance (in meters) that a new keypoints should have to all other keypoints to be added to the map. This can help ensure that there is no overlapping keypoints in a region.
  - "maxDepthErrorVisible": Maximum depth error between 3D map point and observed depth value (from the camera) above which a point is considered as not visible (occluded)
 
  - "gain": The optimizer gain
  - "maxNumIters": Maximum number of optimizer iterations
  - "muInit": Initial value of the Levenberg-Marquardt regularization parameter
  - "muIterFactor": LM regularization scaling factor applied every iteration
  - "minImprovementFactor": Improvement factor (expressed as a ratio of the previous to current error criterion values) below which optimization is stopped.
 
 
  An example JSON representation of the available settings is:
  {
    "type": "xfeat",
    "numPoints": 8192,
    "reprojectionThreshold": 10.0,
    "minDistNewPoint": 0.0,
    "maxDepthErrorVisible": 0.02,
 
    "maxDepthErrorCandidate": 0.0,
 
    "gain": 0.5,
    "maxNumIters": 10,
    "muInit": 0.0,
    "muIterFactor": 0.0,
    "minImprovementFactor": 0.001
  }
  """
 
  class DisplayType(Enum):
    """ Feature display type for visual odometry.
    Changes the amount of displayed information
    """
    SIMPLE = 'simple'
    SIMPLE_MODEL_AND_PROJ = 'simpleModelAndProj'
    ERROR = 'error'
    WEIGHT_AND_ERROR = 'weightAndError'
 
  def __init__(self, backend: XFeatTrackingBackend):
    RBVisualOdometry.__init__(self)
    self.backend = backend
    self.environment_map = TrackedDescriptorMap(num_points=1024,
                                      reprojection_threshold=10.0,
                                      min_dist_new_point=0,
                                      max_depth_error_visible=2e-2)
 
    self.display_type = XFeatVisualOdometry.DisplayType.SIMPLE
    self.use_3d = False
    self.current_representation = None
    # XFeat params
    self.idx_curr_env_matched, self.idx_environment_map = None, None
    self.iter = 0
 
    # Odometry parameters
    self.cTw = HomogeneousMatrix()
    self.cnTc = None
 
    # Initial optimization pass: Minimize reprojection error, wrt to cTw
    self.optim_params = LevenbergMarquardtParameters()
    self.optim_params.gain = 0.5
    self.optim_params.maxNumIters = 10
    self.optim_params.muInit = 0.0
    self.optim_params.muIterFactor = 0.0
    self.optim_params.minImprovementFactor = 1e-3
 
 
  def compute_weights(self, obj_descriptors: torch.Tensor):
    Id = torch.eye(obj_descriptors.size(), device=obj_descriptors.device)
    cos = obj_descriptors @ obj_descriptors.t()
    w = 1.0 - (torch.max(cos - Id, dim=-1))
    return w
 
  def load_settings(self, d):
    """Update the VO parameters from a dictionary.
 
    Args:
        d (_type_): The dictionary containing the settings
    """
    self.environment_map.parse_settings(d)
    self.use_3d = d.get('use3d', self.use_3d)
    self.optim_params.gain = d.get('gain', self.optim_params.gain)
    self.optim_params.maxNumIters = d.get('maxNumIters', self.optim_params.maxNumIters)
    self.optim_params.muInit = d.get('muInit', self.optim_params.muInit)
    self.optim_params.muIterFactor = d.get('muIterFactor',self.optim_params.muIterFactor)
    self.optim_params.minImprovementFactor = d.get('minImprovementFactor', self.optim_params.minImprovementFactor)
    self.display_type = d.get('displayType', self.display_type)
 
 
  def compute(self, frame: RBFeatureTrackerInput, previousFrame: RBFeatureTrackerInput) -> None:
 
    if self.current_representation is not None:
      removed_indices = self.environment_map.update(frame, self.cTw, None, frame.depth, self.idx_curr_env_matched, self.idx_environment_map,
                                                     self.current_representation.keypoints, self.current_representation.descriptors)
 
    self.backend.process_frame(frame, self.iter)
    self.current_representation = self.backend.get_current_environment_data()
 
    self.idx_curr_env_matched, self.idx_environment_map = None, None
    with torch.no_grad():
      if self.environment_map.has_points() and self.current_representation is not None:
        h, w = frame.I.getRows(), frame.I.getCols()
 
        visible_indices = np.ascontiguousarray(self.environment_map.point_map.getVisiblePoints(h, w, frame.cam, self.cTw, frame.depth))
        # visible_indices = self.environment_map.get_visible_points(self.cTw, frame)
        if len(visible_indices) > 0:
          visible_env_descriptors = self.environment_map.descriptors[visible_indices]
          self.idx_curr_env_matched, self.idx_environment_map = self.backend.match(self.current_representation.descriptors, visible_env_descriptors)
          self.idx_curr_env_matched = self.idx_curr_env_matched.cpu().numpy()
          self.idx_environment_map = visible_indices[self.idx_environment_map.cpu().numpy()]
 
    self.numFeatures = 0
    self.iter += 1
 
    if self.idx_curr_env_matched is not None:
 
      if self.use_3d:
        valid_indices = []
        self.Zs = []
        for i, kp in enumerate(self.current_representation.keypoints[self.idx_curr_env_matched]):
          Z = frame.depth[int(kp[1]), int(kp[0])]
          if Z > 0.0:
            valid_indices.append(i)
            self.Zs.append(Z)
        self.Zs = np.ascontiguousarray(self.Zs)
        self.idx_curr_env_matched = self.idx_curr_env_matched[valid_indices]
        self.idx_environment_map = self.idx_environment_map[valid_indices]
 
      self.numFeatures = len(self.idx_curr_env_matched) * (2 if not self.use_3d else 3)
 
    self.optimize(frame)
 
  def optimize(self, frame: RBFeatureTrackerInput):
    self.previous_cTw = HomogeneousMatrix(self.cTw) # Calling constructor is important, we want to keep the value, not the reference as current cTw will be modified!
 
    if self.numFeatures == 0:
      self.cnTc = HomogeneousMatrix()
      return
 
    if self.idx_curr_env_matched is None or len(self.idx_curr_env_matched) == 0:
      self.cnTc = HomogeneousMatrix()
      return
 
    current_px_matched_env = self.current_representation.keypoints[self.idx_curr_env_matched]
    xoe, yoe = PixelMeterConversion.convertPoints(frame.cam, current_px_matched_env[:, 0], current_px_matched_env[:, 1])
    mapX = Matrix()
    indices = ArrayInt2D.view(np.ascontiguousarray(self.idx_environment_map[:, None].astype(np.int32)))
    self.environment_map.point_map.getPoints(indices, mapX)
    c = HomogeneousMatrix(self.cTw)
    if self.use_3d:
      current_xyz_env = Matrix.view(np.stack((xoe * self.Zs, yoe * self.Zs, self.Zs), axis=-1))
      RBVisualOdometryUtils.levenbergMarquardtKeypoints3D(mapX, current_xyz_env, self.optim_params, c)
    else:
      current_xy_env = Matrix.view(np.stack((xoe, yoe), axis=-1))
      RBVisualOdometryUtils.levenbergMarquardtKeypoints2D(mapX, current_xy_env, self.optim_params, c)
 
    self.cTw = c
    self.cnTc = self.cTw * self.previous_cTw.inverse()
 
  def reset(self):
    self.environment_map.reset()
    self.iter = 0
    self.current_representation = None
    self.previous_cTw = HomogeneousMatrix()
    self.cTw = HomogeneousMatrix()
    self.cnTc = HomogeneousMatrix()
    self.idx_curr_env_matched = None
    self.idx_environment_map = None
 
  def getCameraMotion(self) -> HomogeneousMatrix:
    return self.cnTc
 
  def getCameraPose(self) -> HomogeneousMatrix:
    return self.cTw
 
  def display(self, cam: CameraParameters, _I: ImageGray, IRGB: ImageRGBa, _I_depth: ImageGray):
    """Display the features used for visual odometry.
    XFeat keypoints are displayed as small disks in the RGB image.
 
    Depending on the display type, the behavior is different:
    - DisplayType.SIMPLE: Display the current image keypoints that were matched with the environment map
    - DisplayType.SIMPLE_MODEL_AND_PROJ: Display current image keypoints along with the map points that they were matched with.
    - DisplayType.WEIGHT_AND_ERROR: Same as the SIMPLE value, except that the color of the keypoints display the error
    (on the red component) and their attributed weight in the optimization (in blue)
 
    Args:
        cam (CameraParameters): Camera intrinsics
        _I (ImageGray): Grayscale image
        IRGB (ImageRGBa): RGB image where to display the points
        _I_depth (ImageGray): Depth image
 
    Raises:
        RuntimeError: if the display type is invalid
    """
    if self.idx_curr_env_matched is None or self.idx_environment_map is None:
      return
 
    ps = self.current_representation.keypoints[self.idx_curr_env_matched]
    cX, xs, uvs = Matrix(), Matrix(), Matrix()
    self.environment_map.point_map.project(cam, ArrayInt2D.view(np.ascontiguousarray(self.idx_environment_map[:, None]).astype(np.int32)), self.cTw, cX, xs, uvs)
    uvs = uvs.numpy().copy()
    error = np.linalg.norm(uvs - ps, axis=1)
    threshold = self.environment_map.point_map.getOutlierReprojectionErrorThreshold()
    if len(error) == 0:
      return
    max_error_display = np.maximum(np.max(error), self.environment_map.point_map.getOutlierReprojectionErrorThreshold())
    ps = np.rint(ps).astype(np.int32)
    uvs = np.rint(uvs).astype(np.int32)
 
    if self.display_type == XFeatVisualOdometry.DisplayType.SIMPLE:
      Display.displayCrosses(IRGB, ps[:, 1], ps[:, 0], 8, Color.blue, 1)
 
    elif self.display_type == XFeatVisualOdometry.DisplayType.SIMPLE_MODEL_AND_PROJ:
      Display.displayCrosses(IRGB, ps[:, 1], ps[:, 0], 8, Color.blue, 1)
      Display.displayCrosses(IRGB, uvs[:, 1], np.rint(uvs[:, 0]).astype(np.int32), 8, Color.red, 1)
 
    elif self.display_type == XFeatVisualOdometry.DisplayType.WEIGHT_AND_ERROR:
      c = Color()
      for p in range(len(error)):
        e = np.minimum(error[p], threshold)  / threshold
        c.setColor(
          np.rint((e) * 255).astype(np.uint8),
          0,
          np.rint((1.0 - e) * 255).astype(np.uint8),
          255
        )
 
        Display.displayCircleStatic(IRGB, ps[p, 1], ps[p, 0], 2, c, True, 1)
    else:
      raise RuntimeError('Display type not implemented')