Extract projected point cloud

scripts/dataset_generation/extract_binary_maps.py

@@ -168,7 +168,7 @@ def do(n, dry_run):
                     info_msg.header = msg.header
                     try:
                         wvn_ros_interface.image_callback(image_msg, point_cloud_msg, elevation_map_msg, info_msg, camera_options)
-                        print("Time diff", abs(point_cloud_msg.header.stamp.to_sec() - image_msg.header.stamp.to_sec()))
+                        # print("Time diff", abs(point_cloud_msg.header.stamp.to_sec() - image_msg.header.stamp.to_sec()))
                     except Exception as e:
                         print("Bad image_callback", e)
 

wild_visual_navigation/traversability_estimator/nodes.py

@@ -113,6 +113,8 @@ def __init__(
         pose_base_in_world: torch.tensor = torch.eye(4),
         pose_cam_in_base: torch.tensor = torch.eye(4),
         pose_cam_in_world: torch.tensor = None,
+        pose_pc_in_base: torch.tensor = torch.eye(4),
+        pose_pc_in_world: torch.tensor = None,
         image: torch.tensor = None,
         point_cloud: torch.tensor = None,
         grid_map: GridMap = None,
@@ -126,6 +128,10 @@ def __init__(
         self._pose_cam_in_world = (
             self._pose_base_in_world @ self._pose_cam_in_base if pose_cam_in_world is None else pose_cam_in_world
         )
+        self._pose_pc_in_base = pose_pc_in_base
+        self._pose_pc_in_world = (
+            self._pose_base_in_world @ self._pose_pc_in_base if pose_pc_in_world is None else pose_pc_in_world
+        )
         self._image = image
         self._point_cloud = point_cloud
         self._grid_map = grid_map
@@ -285,6 +291,18 @@ def image_projector(self):
     def pose_cam_in_world(self):
         return self._pose_cam_in_world
 
+    @property
+    def pose_cam_in_base(self):
+        return self._pose_cam_in_base
+
+    @property
+    def pose_pc_in_world(self):
+        return self._pose_pc_in_world
+
+    @property
+    def pose_pc_in_base(self):
+        return self._pose_pc_in_base
+
     @property
     def prediction(self):
         return self._prediction
@@ -345,6 +363,18 @@ def image_projector(self, image_projector):
     def pose_cam_in_world(self, pose_cam_in_world):
         self._pose_cam_in_world = pose_cam_in_world
 
+    @pose_cam_in_base.setter
+    def pose_cam_in_base(self, pose_cam_in_base):
+        self._pose_cam_in_base = pose_cam_in_base
+
+    @pose_pc_in_world.setter
+    def pose_pc_in_world(self, pose_pc_in_world):
+        self._pose_pc_in_world = pose_pc_in_world
+
+    @pose_pc_in_base.setter
+    def pose_pc_in_base(self, pose_pc_in_base):
+        self._pose_pc_in_base = pose_pc_in_base
+
     @prediction.setter
     def prediction(self, prediction):
         self._prediction = prediction

wild_visual_navigation/traversability_estimator/traversability_estimator.py

@@ -377,6 +377,8 @@ def add_proprio_node(self, pnode: ProprioceptionNode, projection_mode: str = "im
             )
             pose_camera_in_world = torch.eye(4, device=self._device).repeat(B, 1, 1)
             pose_base_in_world = torch.eye(4, device=self._device).repeat(B, 1, 1)
+            pose_pc_in_base = torch.eye(4, device=self._device).repeat(B, 1, 1)
+            pose_pc_in_world = torch.eye(4, device=self._device).repeat(B, 1, 1)
 
             H = last_mission_node.image_projector.camera.height
             W = last_mission_node.image_projector.camera.width
@@ -386,6 +388,8 @@ def add_proprio_node(self, pnode: ProprioceptionNode, projection_mode: str = "im
                 K[i] = mnode.image_projector.camera.intrinsics
                 pose_camera_in_world[i] = mnode.pose_cam_in_world
                 pose_base_in_world[i] = mnode.pose_base_in_world
+                pose_pc_in_base[i] = mnode.pose_pc_in_base
+                pose_pc_in_world[i] = mnode.pose_pc_in_world
 
                 if (not hasattr(mnode, "supervision_mask")) or (mnode.supervision_mask is None):
                     continue
@@ -458,8 +462,11 @@ def add_proprio_node(self, pnode: ProprioceptionNode, projection_mode: str = "im
                         str(mnode.timestamp).replace(".", "_") + ".jpg",
                     ), img)
 
+                    # Project point cloud to world frame
+                    point_cloud = self.project_pc(mnode.point_cloud, pose_pc_in_base[i])
+
                     # Save point cloud as torch file
-                    torch.save(mnode.point_cloud, os.path.join(
+                    torch.save(point_cloud, os.path.join(
                         self._extraction_store_folder,
                         "pcd",
                         str(mnode.timestamp).replace(".", "_") + ".pt",
@@ -537,6 +544,19 @@ def torch_array_to_grid_map(self, arr, res, layers, timestamp, reference_frame="
 
         return gm_msg
 
+    def project_pc(self, pc, tf):
+
+        tf = tf.cpu()
+        position = np.array(tf[:3, -1])
+        R = np.array(tf[:3, :3])
+
+        # Transform points to frame
+        points_list = []
+        for p in pc:
+            p = np.matmul(R, np.array(p)) + position
+            points_list.append(tuple(p))
+        return torch.tensor(points_list, dtype=torch.float32)
+
     @classmethod
     def load(cls, file_path: str, device="cpu"):
         """Loads pickled file and creates an instance of TraversabilityEstimator,

wild_visual_navigation_ros/scripts/wild_visual_navigation_node.py

@@ -678,18 +678,27 @@ def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_ma
                 self.query_tf(self.fixed_frame, self.base_frame, image_msg.header.stamp), device=self.device
             )
             if not suc:
-                self.system_events["image_callback_cancled"] = {
+                self.system_events["image_callback_canceled"] = {
                     "time": rospy.get_time(),
-                    "value": "cancled due to pose_base_in_world",
+                    "value": "canceled due to pose_base_in_world",
                 }
                 return
             suc, pose_cam_in_base = rc.ros_tf_to_torch(
                 self.query_tf(self.base_frame, image_msg.header.frame_id, image_msg.header.stamp), device=self.device
             )
             if not suc:
-                self.system_events["image_callback_cancled"] = {
+                self.system_events["image_callback_canceled"] = {
+                    "time": rospy.get_time(),
+                    "value": "canceled due to pose_cam_in_base",
+                }
+                return
+            suc, pose_pc_in_base = rc.ros_tf_to_torch(
+                self.query_tf(self.base_frame, point_cloud_msg.header.frame_id, point_cloud_msg.header.stamp), device=self.device
+            )
+            if not suc:
+                self.system_events["image_callback_canceled"] = {
                     "time": rospy.get_time(),
-                    "value": "cancled due to pose_cam_in_base",
+                    "value": "canceled due to pose_pc_in_base",
                 }
                 return
 
@@ -711,6 +720,7 @@ def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_ma
                 timestamp=ts,
                 pose_base_in_world=pose_base_in_world,
                 pose_cam_in_base=pose_cam_in_base,
+                pose_pc_in_base=pose_pc_in_base,
                 image=torch_image,
                 point_cloud=torch_pc,
                 grid_map=grid_map_msg,