diff --git a/utils/preprocess_util.py b/utils/preprocess_util.py index dbef367..154dbe7 100644 --- a/utils/preprocess_util.py +++ b/utils/preprocess_util.py @@ -32,14 +32,17 @@ def save_scan_points(root, scene, scan_points: np.ndarray): scan_points_path = os.path.join(root,scene, "scan_points.txt") save_np_pts(scan_points_path, scan_points) -def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic): +def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic, random_downsample_N): z = depth[mask] i, j = np.nonzero(mask) x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0] y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1] points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3) - points_camera_aug = np.concatenate((points_camera, np.ones((points_camera.shape[0], 1))), axis=-1) + sampled_target_points = PtsUtil.random_downsample_point_cloud( + points_camera, random_downsample_N + ) + points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1) points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] return points_camera_world @@ -58,6 +61,77 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int selected_points_indices = np.where(valid_indices)[0][selected_points_indices] return selected_points_indices +def save_scene_data(root, scene, file_type="txt"): + + ''' configuration ''' + target_mask_label = (0, 255, 0, 255) + display_table_mask_label=(0, 0, 255, 255) + random_downsample_N = 32768 + voxel_size=0.002 + filter_degree = 75 + min_z = 0.2 + max_z = 0.5 + + ''' scan points ''' + display_table_info = DataLoadUtil.get_display_table_info(root, scene) + radius = display_table_info["radius"] + + scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)) + + ''' read frame data(depth|mask|normal) ''' + frame_num = DataLoadUtil.get_scene_seq_length(root, scene) + for frame_id in range(frame_num): + #print(f"[scene({scene_idx}/{scene_total})|frame({frame_id}/{frame_num})]Processing {scene} frame {frame_id}") + path = DataLoadUtil.get_path(root, scene, frame_id) + cam_info = DataLoadUtil.load_cam_info(path, binocular=True) + depth_L, depth_R = DataLoadUtil.load_depth( + path, cam_info["near_plane"], + cam_info["far_plane"], + binocular=True + ) + mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True) + + ''' target points ''' + mask_img_L = mask_L + mask_img_R = mask_R + + target_mask_img_L = (mask_L == target_mask_label).all(axis=-1) + target_mask_img_R = (mask_R == target_mask_label).all(axis=-1) + + + sampled_target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N) + sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N) + + + has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0 + if has_points: + target_points = PtsUtil.get_overlapping_points( + sampled_target_points_L, sampled_target_points_R, voxel_size + ) + + if has_points: + has_points = target_points.shape[0] > 0 + + if has_points: + points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True) + target_points = PtsUtil.filter_points( + target_points, points_normals, cam_info["cam_to_world"],voxel_size=0.002, theta = filter_degree, z_range=(min_z, max_z) + ) + + + ''' scan points indices ''' + scan_points_indices_L = get_scan_points_indices(scan_points, mask_img_L, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) + scan_points_indices_R = get_scan_points_indices(scan_points, mask_img_R, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"]) + scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R) + + if not has_points: + target_points = np.zeros((0, 3)) + + save_target_points(root, scene, frame_id, target_points, file_type=file_type) + save_scan_points_indices(root, scene, frame_id, scan_points_indices, file_type=file_type) + + save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess + def process_frame(frame_id, root, scene, scan_points, file_type, target_mask_label, display_table_mask_label, random_downsample_N, voxel_size, filter_degree, min_z, max_z): Log.info(f"[frame({frame_id})]Processing {scene} frame {frame_id}") path = DataLoadUtil.get_path(root, scene, frame_id) @@ -72,23 +146,14 @@ def process_frame(frame_id, root, scene, scan_points, file_type, target_mask_lab target_mask_img_L = (mask_L == target_mask_label).all(axis=-1) target_mask_img_R = (mask_R == target_mask_label).all(axis=-1) - target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) - target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"]) - - sampled_target_points_L = PtsUtil.random_downsample_point_cloud( - target_points_L, random_downsample_N - ) - sampled_target_points_R = PtsUtil.random_downsample_point_cloud( - target_points_R, random_downsample_N - ) + sampled_target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N) + sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N) has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0 target_points = np.zeros((0, 3)) if has_points: - target_points = PtsUtil.get_overlapping_points( - sampled_target_points_L, sampled_target_points_R, voxel_size - ) + target_points = PtsUtil.get_overlapping_points(sampled_target_points_L, sampled_target_points_R, voxel_size) if has_points and target_points.shape[0] > 0: points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True) @@ -100,13 +165,10 @@ def process_frame(frame_id, root, scene, scan_points, file_type, target_mask_lab scan_points_indices_R = get_scan_points_indices(scan_points, mask_R, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"]) scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R) - if not has_points: - target_points = np.zeros((0, 3)) - save_target_points(root, scene, frame_id, target_points, file_type=file_type) save_scan_points_indices(root, scene, frame_id, scan_points_indices, file_type=file_type) -def save_scene_data(root, scene, file_type="txt"): +def save_scene_data_multithread(root, scene, file_type="txt"): target_mask_label = (0, 255, 0, 255) display_table_mask_label = (0, 0, 255, 255) random_downsample_N = 32768 @@ -115,7 +177,9 @@ def save_scene_data(root, scene, file_type="txt"): min_z = 0.2 max_z = 0.5 - scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0, display_table_radius=0.25)) + display_table_info = DataLoadUtil.get_display_table_info(root, scene) + radius = display_table_info["radius"] + scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0, display_table_radius=radius)) frame_num = DataLoadUtil.get_scene_seq_length(root, scene) with ThreadPoolExecutor() as executor: @@ -131,6 +195,7 @@ def save_scene_data(root, scene, file_type="txt"): save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess + if __name__ == "__main__": #root = "/media/hofee/repository/new_data_with_normal" root = r"/media/hofee/data/tempdir/test_real_output"