diff --git a/human_scene_transformer/data/jrdb_preprocess_test.py b/human_scene_transformer/data/jrdb_preprocess_test.py index e61de9a..fc3ca95 100644 --- 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 index 4397906..8432829 100644 --- 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 index 55f6def..88eddad 100644 --- 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():