From d098c9f9518fdc477b3feb916e749f2e060e0bac Mon Sep 17 00:00:00 2001 From: hofee <64160135+GitHofee@users.noreply.github.com> Date: Sat, 5 Oct 2024 12:24:53 -0500 Subject: [PATCH] optimize preprocessor --- preprocess/preprocessor.py | 147 ++++++++++++++++++------------------- utils/data_load.py | 7 +- utils/pts.py | 32 ++++++-- utils/reconstruction.py | 2 +- 4 files changed, 104 insertions(+), 84 deletions(-) diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py index eea33c9..e3064be 100644 --- a/preprocess/preprocessor.py +++ b/preprocess/preprocessor.py @@ -1,12 +1,9 @@ import os -import json import numpy as np - +import time import sys np.random.seed(0) -# append parent directory to sys.path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -print(sys.path) from utils.reconstruction import ReconstructionUtil from utils.data_load import DataLoadUtil @@ -18,26 +15,12 @@ def save_np_pts(path, pts: np.ndarray, file_type="txt"): else: np.save(path, pts) -def save_full_points(root, scene, frame_idx, full_points: np.ndarray, file_type="txt"): - pts_path = os.path.join(root,scene, "scene_pts", f"{frame_idx}.{file_type}") - if not os.path.exists(os.path.join(root,scene, "scene_pts")): - os.makedirs(os.path.join(root,scene, "scene_pts")) - save_np_pts(pts_path, full_points, file_type) - def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"): - pts_path = os.path.join(root,scene, "target_pts", f"{frame_idx}.{file_type}") - if not os.path.exists(os.path.join(root,scene, "target_pts")): - os.makedirs(os.path.join(root,scene, "target_pts")) + pts_path = os.path.join(root,scene, "pts", f"{frame_idx}.{file_type}") + if not os.path.exists(os.path.join(root,scene, "pts")): + os.makedirs(os.path.join(root,scene, "pts")) save_np_pts(pts_path, target_points, file_type) -def save_mask_idx(root, scene, frame_idx, mask_idx: np.ndarray,filtered_idx, file_type="txt"): - indices_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}.{file_type}") - if not os.path.exists(os.path.join(root,scene, "mask_idx")): - os.makedirs(os.path.join(root,scene, "mask_idx")) - save_np_pts(indices_path, mask_idx, file_type) - filtered_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}_filtered.{file_type}") - save_np_pts(filtered_path, filtered_idx, file_type) - def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"): indices_path = os.path.join(root,scene, "scan_points_indices", f"{frame_idx}.{file_type}") if not os.path.exists(os.path.join(root,scene, "scan_points_indices")): @@ -49,27 +32,36 @@ def save_scan_points(root, scene, scan_points: np.ndarray): save_np_pts(scan_points_path, scan_points) -def get_world_points(depth, cam_intrinsic, cam_extrinsic): +def old_get_world_points(depth, cam_intrinsic, cam_extrinsic): h, w = depth.shape i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy") - + # ----- Debug Trace ----- # + import ipdb; ipdb.set_trace() + # ------------------------ # z = depth x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0] y = (j - 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) + points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] + + return points_camera_world + +def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic): + 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) points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] + return points_camera_world -def get_world_normals(normals, cam_extrinsic): - normals = normals / np.linalg.norm(normals, axis=1, keepdims=True) - normals_world = np.dot(cam_extrinsic[:3, :3], normals.T).T - return normals_world - def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic): scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1)))) - points_camera = np.dot(cam_extrinsic, scan_points_homogeneous.T).T[:, :3] + points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3] points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T points_image_homogeneous /= points_image_homogeneous[:, 2:] pixel_x = points_image_homogeneous[:, 0].astype(int) @@ -77,7 +69,8 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int h, w = mask.shape[:2] valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h) mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]] - selected_points_indices = mask_colors == display_table_mask_label + selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0] + selected_points_indices = np.where(valid_indices)[0][selected_points_indices] return selected_points_indices @@ -86,14 +79,16 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1): ''' configuration ''' target_mask_label = (0, 255, 0, 255) display_table_mask_label=(0, 0, 255, 255) - random_downsample_N = 65536 - train_input_pts_num = 8192 + random_downsample_N = 32768 voxel_size=0.002 filter_degree = 75 + min_z = 0.2 + max_z = 0.45 ''' 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) ''' @@ -107,52 +102,50 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1): cam_info["far_plane"], binocular=True ) - mask_L = DataLoadUtil.load_seg(path, binocular=True, left_only=True) - normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True) - - ''' scene points ''' - scene_points_L = get_world_points(depth_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) - scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"]) - sampled_scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud( - scene_points_L, random_downsample_N, require_idx=True - ) - sampled_scene_points_R = PtsUtil.random_downsample_point_cloud( - scene_points_R, random_downsample_N - ) - scene_overlap_points, overlap_idx_L = PtsUtil.get_overlapping_points( - sampled_scene_points_L, sampled_scene_points_R, voxel_size, require_idx=True - ) - - if scene_overlap_points.shape[0] < 1024: - scene_overlap_points = sampled_scene_points_L - overlap_idx_L = np.arange(sampled_scene_points_L.shape[0]) - train_input_points, train_input_idx = PtsUtil.random_downsample_point_cloud( - scene_overlap_points, train_input_pts_num, require_idx=True - ) + mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True) ''' target points ''' - mask_img = mask_L - mask_L = mask_L.reshape(-1, 4) - mask_L = (mask_L == target_mask_label).all(axis=-1) - mask_overlap = mask_L[random_sample_idx_L][overlap_idx_L] - scene_normals_L = normal_L.reshape(-1, 3) - target_overlap_normals = scene_normals_L[random_sample_idx_L][overlap_idx_L][mask_overlap] - target_normals = get_world_normals(target_overlap_normals, cam_info["cam_to_world"]) - target_points = scene_overlap_points[mask_overlap] - filtered_target_points, filtered_idx = PtsUtil.filter_points( - target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True + 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) + + + 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 + ) + + 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 + ) + - ''' train_input_mask ''' - mask_train_input = mask_overlap[train_input_idx] - - + 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 = get_scan_points_indices(scan_points, mask_img, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) + 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) - save_full_points(root, scene, frame_id, train_input_points) - save_target_points(root, scene, frame_id, filtered_target_points) - save_mask_idx(root, scene, frame_id, mask_train_input, filtered_idx=filtered_idx) + if not has_points: + target_points = np.zeros((0, 3)) + + save_target_points(root, scene, frame_id, target_points) save_scan_points_indices(root, scene, frame_id, scan_points_indices) save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess @@ -160,8 +153,8 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1): if __name__ == "__main__": #root = "/media/hofee/repository/new_data_with_normal" - root = "/media/hofee/repository/test_sample" - list_path = "/media/hofee/repository/test_sample/test_sample_list.txt" + root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample" + list_path = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample/test_sample_list.txt" scene_list = [] with open(list_path, "r") as f: @@ -171,7 +164,11 @@ if __name__ == "__main__": from_idx = 0 to_idx = len(scene_list) cnt = 0 + import time total = to_idx - from_idx for scene in scene_list[from_idx:to_idx]: + start = time.time() save_scene_data(root, scene, cnt, total) - cnt+=1 \ No newline at end of file + cnt+=1 + end = time.time() + print(f"Time cost: {end-start}") \ No newline at end of file diff --git a/utils/data_load.py b/utils/data_load.py index 33f65c0..ed5a144 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -204,7 +204,9 @@ class DataLoadUtil: os.path.dirname(path), "normal", os.path.basename(path) + "_R.png" ) normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_COLOR) - return normal_image_L[:3,:3], normal_image_R[:3,:3] + normalized_normal_image_L = normal_image_L / 255.0 * 2.0 - 1.0 + normalized_normal_image_R = normal_image_R / 255.0 * 2.0 - 1.0 + return normalized_normal_image_L, normalized_normal_image_R else: if binocular and left_only: normal_path = os.path.join( @@ -215,7 +217,8 @@ class DataLoadUtil: os.path.dirname(path), "normal", os.path.basename(path) + ".png" ) normal_image = cv2.imread(normal_path, cv2.IMREAD_COLOR) - return normal_image + normalized_normal_image = normal_image / 255.0 * 2.0 - 1.0 + return normalized_normal_image @staticmethod def load_label(path): diff --git a/utils/pts.py b/utils/pts.py index 66258ac..0551149 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -1,6 +1,7 @@ import numpy as np import open3d as o3d import torch +from scipy.spatial import cKDTree class PtsUtil: @@ -56,17 +57,36 @@ class PtsUtil: return overlapping_points @staticmethod - def filter_points(points, normals, cam_pose, theta=75, require_idx=False): + def new_filter_points(points, normals, cam_pose, theta=75, require_idx=False): camera_axis = -cam_pose[:3, 2] normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True) cos_theta = np.dot(normals_normalized, camera_axis) theta_rad = np.deg2rad(theta) idx = cos_theta > np.cos(theta_rad) - print(cos_theta, theta_rad) filtered_points= points[idx] - # ------ Debug Start ------ - import ipdb;ipdb.set_trace() - # ------ Debug End ------ if require_idx: return filtered_points, idx - return filtered_points \ No newline at end of file + return filtered_points + + @staticmethod + def filter_points(points, points_normals, cam_pose, voxel_size=0.002, theta=45, z_range=(0.2, 0.45)): + + """ filter with z range """ + points_cam = PtsUtil.transform_point_cloud(points, np.linalg.inv(cam_pose)) + idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1]) + z_filtered_points = points[idx] + + """ filter with normal """ + sampled_points = PtsUtil.voxel_downsample_point_cloud(z_filtered_points, voxel_size) + kdtree = cKDTree(points_normals[:,:3]) + _, indices = kdtree.query(sampled_points) + nearest_points = points_normals[indices] + + normals = nearest_points[:, 3:] + camera_axis = -cam_pose[:3, 2] + normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True) + cos_theta = np.dot(normals_normalized, camera_axis) + theta_rad = np.deg2rad(theta) + idx = cos_theta > np.cos(theta_rad) + filtered_sampled_points= sampled_points[idx] + return filtered_sampled_points[:, :3] \ No newline at end of file diff --git a/utils/reconstruction.py b/utils/reconstruction.py index 74f039f..530358d 100644 --- a/utils/reconstruction.py +++ b/utils/reconstruction.py @@ -132,7 +132,7 @@ class ReconstructionUtil: @staticmethod - def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 100, max_attempts = 1000): + def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 500, max_attempts = 1000): points = [] attempts = 0 while len(points) < max_points_num and attempts < max_attempts: