Extract all point clouds

scripts/dataset_generation/extract_binary_maps.py

@@ -124,6 +124,8 @@ def do(n, dry_run):
     total_time_state = 0
     n = 0
 
+    point_clouds = {}
+
     with rosbag.Bag(output_bag_wvn, "r") as bag:
         if rospy.is_shutdown():
             return
@@ -149,16 +151,19 @@ def do(n, dry_run):
                     state_msg_valid = True
 
                 if topic == "/depth_camera_front_upper/point_cloud_self_filtered":
-                    point_cloud_msg = msg
-                    # print("point_cloud_msg", msg.header.stamp.to_sec())
+                    point_clouds["depth_camera_front_upper"] = msg
+
+                if topic == "/depth_camera_rear_upper/point_cloud_self_filtered":
+                    point_clouds["depth_camera_rear_upper"] = msg
+
+                if topic == "/depth_camera_left/point_cloud_self_filtered":
+                    point_clouds["depth_camera_left"] = msg
 
-                if topic == "/elevation_mapping/elevation_map_raw":
-                    elevation_map_msg = msg
+                if topic == "/depth_camera_right/point_cloud_self_filtered":
+                    point_clouds["depth_camera_right"] = msg
 
                 elif topic == "/wide_angle_camera_front/img_out":
                     image_msg = msg
-                    # print("Received /wide_angle_camera_front/img_out")
-                    # print("image_msg", msg.header.stamp.to_sec())
 
                     info_msg.header = msg.header
                     camera_options = {}
@@ -167,7 +172,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)
+                        wvn_ros_interface.image_callback(image_msg, point_clouds, info_msg, camera_options)
                         # 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/cfg/extraction_params.py

@@ -6,7 +6,9 @@ class ExtractionParams:
     wvn_topics: List[str] = field(default_factory=lambda: ["/state_estimator/anymal_state",
                                                            "/wide_angle_camera_front/img_out",
                                                            "/depth_camera_front_upper/point_cloud_self_filtered",
-                                                           "/elevation_mapping/elevation_map_raw"])
+                                                           "/depth_camera_rear_upper/point_cloud_self_filtered",
+                                                           "/depth_camera_left/point_cloud_self_filtered",
+                                                           "/depth_camera_right/point_cloud_self_filtered"])
     wvn_bags: List[str] = field(default_factory=lambda: ["/home/rschmid/RosBags/6_proc/images.bag",
                                                          "/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-jetson_mission_0.bag",
                                                          "/home/rschmid/RosBags/6_proc/2023-03-02-11-13-08_anymal-d020-jetson_mission_1.bag",

wild_visual_navigation/traversability_estimator/nodes.py

@@ -6,7 +6,6 @@
 import torch
 import torch.nn.functional as F
 from typing import Optional
-from grid_map_msgs.msg import GridMap, GridMapInfo
 
 
 class BaseNode:
@@ -113,11 +112,10 @@ 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,
+        pose_pc_in_base: dict = None,
+        pose_pc_in_world: dict = None,
         image: torch.tensor = None,
-        point_cloud: torch.tensor = None,
-        grid_map: GridMap = None,
+        point_clouds: dict = None,
         image_projector: ImageProjector = None,
         camera_name="cam",
         use_for_training=True,
@@ -129,12 +127,14 @@ def __init__(
             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._pose_pc_in_world = {}
+        for key in self._pose_pc_in_base:
+            self._pose_pc_in_world[key] = (
+                self._pose_base_in_world @ self._pose_pc_in_base[key] 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
+        self._point_clouds = point_clouds
         self._image_projector = image_projector
         self._camera_name = camera_name
         self._use_for_training = use_for_training
@@ -276,12 +276,8 @@ def image(self):
         return self._image
 
     @property
-    def point_cloud(self):
-        return self._point_cloud
-
-    @property
-    def grid_map(self):
-        return self._grid_map
+    def point_clouds(self):
+        return self._point_clouds
 
     @property
     def image_projector(self):

wild_visual_navigation/traversability_estimator/traversability_estimator.py

@@ -377,8 +377,10 @@ 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)
+            # 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)
+            pose_pc_in_base = {}
+            pose_pc_in_world = {}
 
             H = last_mission_node.image_projector.camera.height
             W = last_mission_node.image_projector.camera.width
@@ -388,8 +390,6 @@ 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
@@ -462,8 +462,12 @@ 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])
+                    # Project point cloud to world frame, concatenate all points
+                    point_cloud = []
+                    for key in mnode.point_clouds:
+                        # print(key)
+                        point_cloud.append(self.project_pc(mnode.point_clouds[key], mnode.pose_pc_in_base[key]))
+                    point_cloud = torch.cat(point_cloud, dim=0)
 
                     # Save point cloud as torch file
                     torch.save(point_cloud, os.path.join(

wild_visual_navigation_ros/scripts/wild_visual_navigation_node.py

@@ -298,7 +298,10 @@ def read_params(self):
             self.proprio_graph_dist_thr = 0.1
 
             os.makedirs(os.path.join(self.extraction_store_folder, "image"), exist_ok=True)
-            os.makedirs(os.path.join(self.extraction_store_folder, "supervision_mask"), exist_ok=True)
+            os.makedirs(os.path.join(self.extraction_store_folder, "image_jpg"), exist_ok=True)
+            os.makedirs(os.path.join(self.extraction_store_folder, "mask"), exist_ok=True)
+            os.makedirs(os.path.join(self.extraction_store_folder, "mask_jpg"), exist_ok=True)
+            os.makedirs(os.path.join(self.extraction_store_folder, "pcd"), exist_ok=True)
 
         # Experiment file
         exp_file = rospy.get_param("~exp")
@@ -650,7 +653,7 @@ def robot_state_callback(self, state_msg, desired_twist_msg: TwistStamped):
             raise Exception("Error in robot state callback")
 
     @accumulate_time
-    def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_map_msg: GridMap, info_msg: CameraInfo, camera_options: dict):
+    def image_callback(self, image_msg: Image, point_clouds: dict, info_msg: CameraInfo, camera_options: dict):
         """Main callback to process incoming images
 
         Args:
@@ -692,15 +695,18 @@ def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_ma
                     "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": "canceled due to pose_pc_in_base",
-                }
-                return
+
+            pose_pc_in_base = {}
+            for key in point_clouds.keys():
+                suc, pose_pc_in_base[key] = rc.ros_tf_to_torch(
+                    self.query_tf(self.base_frame, point_clouds[key].header.frame_id, point_clouds[key].header.stamp), device=self.device
+                )
+                if not suc:
+                    self.system_events["image_callback_canceled"] = {
+                        "time": rospy.get_time(),
+                        "value": "canceled due to pose_pc_in_base",
+                    }
+                    return
 
             # Prepare image projector
             K, H, W = rc.ros_cam_info_to_tensors(info_msg, device=self.device)
@@ -713,7 +719,9 @@ def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_ma
             torch_image = rc.ros_image_to_torch(image_msg, device=self.device)
             torch_image = image_projector.resize_image(torch_image)
 
-            torch_pc = self.rospcmsg_to_pctorch(point_cloud_msg)
+            torch_pc = {}
+            for key in point_clouds.keys():
+                torch_pc[key] = self.rospcmsg_to_pctorch(point_clouds[key])
 
             # Create mission node for the graph
             mission_node = MissionNode(
@@ -722,8 +730,7 @@ def image_callback(self, image_msg: Image, point_cloud_msg: PointCloud2, grid_ma
                 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,
+                point_clouds=torch_pc,
                 image_projector=image_projector,
                 camera_name=camera_options["name"],
                 use_for_training=camera_options["use_for_training"],