Make PC processing better configurable.

PiperOrigin-RevId: 582956956
google-research · Nov 16, 2023 · e7fa9bf · e7fa9bf
1 parent 7f35f54
commit e7fa9bf
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/human_scene_transformer/data/jrdb_preprocess_test.py b/human_scene_transformer/data/jrdb_preprocess_test.py
@@ -52,6 +52,13 @@
           ' in the processed dataset.')
 )
 
+_MAX_PC_DISTANCE_TO_ROBOT = flags.DEFINE_float(
+    'max_pc_distance_to_robot',
+    default=10.,
+    help=('Maximum distance of pointcloud point to the robot to be included'
+          ' in the processed dataset.')
+)
+
 _TRACKING_METHOD = flags.DEFINE_string(
     'tracking_method',
     default='ss3d_mot',
@@ -260,7 +267,8 @@ def jrdb_preprocess_test(input_path, output_path):
       )
 
       filtered_pc = utils.filter_agents_and_ground_from_point_cloud(
-          agents_in_odometry_df, scene_pc_odometry, robot_in_odometry_df
+          agents_in_odometry_df, scene_pc_odometry, robot_in_odometry_df,
+          max_dist=_MAX_PC_DISTANCE_TO_ROBOT.value,
       )
 
       scene_pc_ragged_tensor = tf.ragged.stack(filtered_pc)

diff --git a/human_scene_transformer/data/jrdb_preprocess_train.py b/human_scene_transformer/data/jrdb_preprocess_train.py
@@ -53,6 +53,13 @@
           ' in the processed dataset.')
 )
 
+_MAX_PC_DISTANCE_TO_ROBOT = flags.DEFINE_float(
+    'max_pc_distance_to_robot',
+    default=10.,
+    help=('Maximum distance of pointcloud point to the robot to be included'
+          ' in the processed dataset.')
+)
+
 
 AGENT_KEYPOINTS = True
 FROM_DETECTIONS = True
@@ -252,7 +259,8 @@ def jrdb_preprocess_train(input_path, output_path):
       )
 
       filtered_pc = utils.filter_agents_and_ground_from_point_cloud(
-          agents_in_odometry_df, scene_pc_odometry, robot_in_odometry_df
+          agents_in_odometry_df, scene_pc_odometry, robot_in_odometry_df,
+          max_dist=_MAX_PC_DISTANCE_TO_ROBOT.value,
       )
 
       scene_pc_ragged_tensor = tf.ragged.stack(filtered_pc)

diff --git a/human_scene_transformer/data/utils.py b/human_scene_transformer/data/utils.py
@@ -275,12 +275,12 @@ def box_to_hyperplanes(pos, yaw, l, w, h):
 
 
 def filter_agents_and_ground_from_point_cloud(
-    agents_df, pointcloud_dict, robot_in_odometry_df):
+    agents_df, pointcloud_dict, robot_in_odometry_df, max_dist=10.):
   """Filter points which are in human bb or belong to ground."""
   for t, agent_df in agents_df.groupby('timestep'):
     pc_points = pointcloud_dict[t]
     robot_p = robot_in_odometry_df.loc[t]['p'][:2]
-    dist_mask = np.linalg.norm(robot_p - pc_points[..., :2], axis=-1) < 10.
+    dist_mask = np.linalg.norm(robot_p - pc_points[..., :2], axis=-1) < max_dist
     pc_points = pc_points[
         (pc_points[:, -1] > -0.2) & (pc_points[:, -1] < 0.5) & dist_mask]
     for _, row in agent_df.iterrows():