Feature/changes for multirobot (#18)

* start setting up vlc msgs

* (WIP) message conversions

* clean up and unittests

* (WIP) multi-robot extension ROS node

* (WIP) multi-robot extension ROS node

* set up tf and exchanging robot info

* set up node and launch file

* initial minor bug fixes

* more debugging

* ignore ouroboros-ros tests

* update has_image check

* allow inter-robot lcs to bypass time filter

* minor bug fixes and debugging prints

* fix self->other loop closures and fix embedding bugs

* finish self->other loop closures

* add depth field to vlc image

* fix bug from keypoint depths

* update how we get keypoint depths and add options for query and match

* working multi-robot

* fix unittests for query

* fix depth getting

* remove multi-robot files

* remove ouroboros msgs and any multi-robot changes

* wrap body_T_cam get in loop

* remove msg conversion tests for now

Author: yunzc

extra/ouroboros_ros/src/ouroboros_ros/utils.py

@@ -1,21 +1,81 @@
 import numpy as np
 import rospy
 import tf2_ros
+from pose_graph_tools_msgs.msg import PoseGraphEdge
+from scipy.spatial.transform import Rotation as R
 
-from ouroboros import VlcPose
+import ouroboros as ob
 
 
-def get_tf_as_pose(tf_buffer, fixed_frame, body_frame, time=None):
+def build_robot_lc_message(
+    key_from_ns,
+    key_to_ns,
+    robot_id,
+    from_T_to,
+    pose_cov,
+    body_T_cam,
+):
+    return build_lc_message(
+        key_from_ns,
+        key_to_ns,
+        robot_id,
+        robot_id,
+        from_T_to,
+        pose_cov,
+        body_T_cam,
+        body_T_cam,
+    )
+
+
+def build_lc_message(
+    key_from_ns,
+    key_to_ns,
+    robot_from,
+    robot_to,
+    from_T_to,
+    pose_cov,
+    robot_from_T_cam,
+    robot_to_T_cam,
+):
+    bodyfrom_T_bodyto = robot_from_T_cam @ from_T_to @ ob.invert_pose(robot_to_T_cam)
+    relative_position = bodyfrom_T_bodyto[:3, 3]
+    relative_orientation = R.from_matrix(bodyfrom_T_bodyto[:3, :3])
+
+    lc_edge = PoseGraphEdge()
+    lc_edge.header.stamp = rospy.Time.now()
+    lc_edge.key_from = key_from_ns
+    lc_edge.key_to = key_to_ns
+    lc_edge.robot_from = robot_from
+    lc_edge.robot_to = robot_to
+    lc_edge.type = PoseGraphEdge.LOOPCLOSE
+    lc_edge.pose.position.x = relative_position[0]
+    lc_edge.pose.position.y = relative_position[1]
+    lc_edge.pose.position.z = relative_position[2]
+    q = relative_orientation.as_quat()
+    lc_edge.pose.orientation.x = q[0]
+    lc_edge.pose.orientation.y = q[1]
+    lc_edge.pose.orientation.z = q[2]
+    lc_edge.pose.orientation.w = q[3]
+
+    lc_edge.covariance = pose_cov.flatten()
+    return lc_edge
+
+
+def get_tf_as_pose(tf_buffer, fixed_frame, body_frame, time=None, timeout=1.0):
     if time is None:
         time = rospy.Time()
     try:
-        trans = tf_buffer.lookup_transform(fixed_frame, body_frame, time)
+        trans = tf_buffer.lookup_transform(
+            fixed_frame, body_frame, time, rospy.Duration(timeout)
+        )
     except (
         tf2_ros.LookupException,
         tf2_ros.ConnectivityException,
         tf2_ros.ExtrapolationException,
-    ):
-        rospy.logwarn(" Could not transform %s from %s ", fixed_frame, body_frame)
+    ) as e:
+        rospy.logwarn(
+            " Could not transform %s from %s: %s ", fixed_frame, body_frame, str(e)
+        )
         return
 
     current_pos = np.array(
@@ -36,5 +96,14 @@ def get_tf_as_pose(tf_buffer, fixed_frame, body_frame, time=None):
     )
 
     time_ns = trans.header.stamp.to_nsec()
-    vlc_pose = VlcPose(time_ns=time_ns, position=current_pos, rotation=current_rot)
+    vlc_pose = ob.VlcPose(time_ns=time_ns, position=current_pos, rotation=current_rot)
     return vlc_pose
+
+
+def parse_camera_info(info_msg):
+    K = np.array(info_msg.K).reshape((3, 3))
+    fx = K[0, 0]
+    fy = K[1, 1]
+    cx = K[0, 2]
+    cy = K[1, 2]
+    return ob.PinholeCamera(fx=fx, fy=fy, cx=cx, cy=cy)

extra/ouroboros_ros/src/ouroboros_ros/vlc_server_ros.py

@@ -8,13 +8,16 @@
 import numpy as np
 import rospy
 import tf2_ros
-from pose_graph_tools_msgs.msg import PoseGraph, PoseGraphEdge
-from scipy.spatial.transform import Rotation as R
+from pose_graph_tools_msgs.msg import PoseGraph
 from sensor_msgs.msg import CameraInfo, Image
 
 import ouroboros as ob
 from ouroboros.utils.plotting_utils import plt_fast_pause
-from ouroboros_ros.utils import get_tf_as_pose
+from ouroboros_ros.utils import (
+    build_robot_lc_message,
+    get_tf_as_pose,
+    parse_camera_info,
+)
 
 
 def update_plot(line, pts, images_to_pose):
@@ -56,38 +59,6 @@ def get_query_and_est_match_position(lc, image_to_pose, body_T_cam):
     return query_position, est_match_position
 
 
-def build_lc_message(
-    key_from_ns,
-    key_to_ns,
-    robot_id,
-    from_T_to,
-    pose_cov,
-    body_T_cam,
-):
-    bodyfrom_T_bodyto = body_T_cam @ from_T_to @ ob.invert_pose(body_T_cam)
-    relative_position = bodyfrom_T_bodyto[:3, 3]
-    relative_orientation = R.from_matrix(bodyfrom_T_bodyto[:3, :3])
-
-    lc_edge = PoseGraphEdge()
-    lc_edge.header.stamp = rospy.Time.now()
-    lc_edge.key_from = key_from_ns
-    lc_edge.key_to = key_to_ns
-    lc_edge.robot_from = 0
-    lc_edge.robot_to = 0
-    lc_edge.type = PoseGraphEdge.LOOPCLOSE
-    lc_edge.pose.position.x = relative_position[0]
-    lc_edge.pose.position.y = relative_position[1]
-    lc_edge.pose.position.z = relative_position[2]
-    q = relative_orientation.as_quat()
-    lc_edge.pose.orientation.x = q[0]
-    lc_edge.pose.orientation.y = q[1]
-    lc_edge.pose.orientation.z = q[2]
-    lc_edge.pose.orientation.w = q[3]
-
-    lc_edge.covariance = pose_cov.flatten()
-    return lc_edge
-
-
 class VlcServerRos:
     def __init__(self):
         self.tf_buffer = tf2_ros.Buffer()
@@ -111,13 +82,13 @@ def __init__(self):
         server_config = ob.VlcServerConfig.load(config_path)
         self.vlc_server = ob.VlcServer(
             server_config,
-            robot_id=0,
+            robot_id=self.robot_id,
         )
 
-        camera_config = self.get_camera_config_ros()
-        print(f"camera config: {camera_config}")
+        self.camera_config = self.get_camera_config_ros()
+        print(f"camera config: {self.camera_config}")
         self.session_id = self.vlc_server.register_camera(
-            0, camera_config, rospy.Time.now().to_nsec()
+            self.robot_id, self.camera_config, rospy.Time.now().to_nsec()
         )
 
         self.loop_rate = rospy.Rate(10)
@@ -145,6 +116,8 @@ def __init__(self):
         )
         self.image_depth_sub.registerCallback(self.image_depth_callback)
 
+        self.body_T_cam = self.get_body_T_cam_ros()
+
     def get_camera_config_ros(self):
         rate = rospy.Rate(10)
         while not rospy.is_shutdown():
@@ -154,14 +127,29 @@ def get_camera_config_ros(self):
                 rospy.logerr("Timed out waiting for camera info")
                 rate.sleep()
                 continue
-            break
-        K = np.array(info_msg.K).reshape((3, 3))
-        fx = K[0, 0]
-        fy = K[1, 1]
-        cx = K[0, 2]
-        cy = K[1, 2]
-        pinhole = ob.PinholeCamera(fx=fx, fy=fy, cx=cx, cy=cy)
-        return pinhole
+            pinhole = parse_camera_info(info_msg)
+            return pinhole
+        return None
+
+    def get_body_T_cam_ros(self, max_retries=10, retry_delay=0.5):
+        rate = rospy.Rate(10)
+        while not rospy.is_shutdown():
+            try:
+                return get_tf_as_pose(
+                    self.tf_buffer, self.body_frame, self.camera_frame, rospy.Time.now()
+                )
+            except rospy.ROSException:
+                rospy.logerr("Failed to get body_T_cam transform")
+                rate.sleep()
+        return None
+
+    def process_new_frame(self, image, stamp_ns, hint_pose):
+        return self.vlc_server.add_and_query_frame(
+            self.session_id,
+            image,
+            stamp_ns,
+            pose_hint=hint_pose,
+        )
 
     def image_depth_callback(self, img_msg, depth_msg):
         if not (
@@ -199,17 +187,11 @@ def image_depth_callback(self, img_msg, depth_msg):
             print("Cannot get current pose, skipping frame!")
             return
 
-        body_T_cam = get_tf_as_pose(
-            self.tf_buffer, self.body_frame, self.camera_frame, img_msg.header.stamp
-        ).as_matrix()
-
         spark_image = ob.SparkImage(rgb=color_image, depth=depth_image)
-        image_uuid, loop_closures = self.vlc_server.add_and_query_frame(
-            self.session_id,
-            spark_image,
-            img_msg.header.stamp.to_nsec(),
-            pose_hint=hint_pose,
+        image_uuid, loop_closures = self.process_new_frame(
+            spark_image, img_msg.header.stamp.to_nsec(), hint_pose
         )
+
         with self.image_pose_lock:
             self.images_to_pose[image_uuid] = hint_pose
 
@@ -223,7 +205,7 @@ def image_depth_callback(self, img_msg, depth_msg):
             if self.show_plots:
                 with self.image_pose_lock:
                     query_pos, match_pos = get_query_and_est_match_position(
-                        lc, self.images_to_pose, body_T_cam
+                        lc, self.images_to_pose, self.body_T_cam.as_matrix()
                     )
                     true_match_pos = self.images_to_pose[lc.to_image_uuid].position
                 if not lc.is_metric:
@@ -239,13 +221,13 @@ def image_depth_callback(self, img_msg, depth_msg):
 
             from_time_ns, to_time_ns = self.vlc_server.get_lc_times(lc.metadata.lc_uuid)
 
-            lc_edge = build_lc_message(
+            lc_edge = build_robot_lc_message(
                 from_time_ns,
                 to_time_ns,
                 self.robot_id,
                 lc.f_T_t,
                 pose_cov_mat,
-                body_T_cam,
+                self.body_T_cam.as_matrix(),
             )
 
             pg.edges.append(lc_edge)

src/ouroboros/pose_recovery.py

@@ -74,7 +74,11 @@ def from_image(
         cls, cam: PinholeCamera, img: VlcImage, indices: Optional[np.ndarray] = None
     ):
         """Get undistorted geometry from keypoints."""
-        depths = img.get_feature_depths()
+        if img.keypoint_depths is None:
+            depths = img.get_feature_depths()
+        else:
+            depths = img.keypoint_depths
+
         keypoints = img.keypoints
         if indices is not None:
             keypoints = keypoints[indices, :]

src/ouroboros/vlc_db/vlc_db.py

@@ -41,6 +41,9 @@ def add_image(
             session_id, image_timestamp, image, pose_hint
         )
 
+    def has_image(self, image_uuid: str) -> bool:
+        return self._image_table.has_image(image_uuid)
+
     def get_image(self, image_uuid: str) -> VlcImage:
         img = self._image_table.get_image(image_uuid)
         if img is None:
@@ -94,8 +97,10 @@ def query_embeddings_max_time(
         self,
         embedding: np.ndarray,
         k: int,
+        sessions_to_time_filter: [str],
         max_time: Union[float, int, datetime],
         similarity_metric: Union[str, callable] = "ip",
+        search_sessions: Optional[List[str]] = None,
     ) -> ([VlcImage], [float]):
         """Query image embeddings to find the k closest vectors with timestamp older than max_time."""
 
@@ -105,11 +110,24 @@ def query_embeddings_max_time(
         # this to be more efficient and not iterate through the full set of
         # vectors
 
+        def session_filter(_, vlc_image, similarity):
+            if search_sessions is None:
+                return True
+            return vlc_image.metadata.session_id in search_sessions
+
         def time_filter(_, vlc_image, similarity):
-            return vlc_image.metadata.epoch_ns < max_time
+            return (
+                vlc_image.metadata.epoch_ns < max_time
+                or vlc_image.metadata.session_id not in sessions_to_time_filter
+            )
+
+        def combined_filter(_, vlc_image, similarity):
+            return time_filter(_, vlc_image, similarity) and session_filter(
+                _, vlc_image, similarity
+            )
 
         ret = self.batch_query_embeddings_filter(
-            np.array([embedding]), k, time_filter, similarity_metric
+            np.array([embedding]), k, combined_filter, similarity_metric
         )
         matches = [t[2] for t in ret[0]]
         similarities = [t[0] for t in ret[0]]
@@ -210,6 +228,10 @@ def update_keypoints(self, image_uuid: str, keypoints, descriptors=None):
         )
         return self.get_image(image_uuid)
 
+    def update_keypoint_depths(self, image_uuid: str, keypoint_depths):
+        self._image_table.update_keypoint_depths(image_uuid, keypoint_depths)
+        return self.get_image(image_uuid)
+
     def get_keypoints(self, image_uuid: str):
         return self._image_table.get_keypoints(image_uuid)
 

src/ouroboros/vlc_db/vlc_image.py

@@ -19,6 +19,7 @@ class VlcImage:
     image: SparkImage
     embedding: np.ndarray = None
     keypoints: np.ndarray = None
+    keypoint_depths: np.ndarray = None
     descriptors: np.ndarray = None
     pose_hint: VlcPose = None
 
@@ -39,7 +40,7 @@ def get_feature_depths(self):
         if self.keypoints is None:
             return None
 
-        if self.image.depth is None:
+        if self.image is None or self.image.depth is None:
             return None
 
         # NOTE(nathan) this is ugly, but: