From fd7614c8470f939cec5248bc1985f0b4d1fb0449 Mon Sep 17 00:00:00 2001 From: hofee Date: Thu, 3 Oct 2024 23:36:18 +0800 Subject: [PATCH] update preprocessor --- configs/local/view_generate_config.yaml | 10 ++--- preprocess/preprocessor.py | 58 ++++++++++++++++++------- utils/data_load.py | 11 +++-- utils/pts.py | 6 +++ 4 files changed, 58 insertions(+), 27 deletions(-) diff --git a/configs/local/view_generate_config.yaml b/configs/local/view_generate_config.yaml index d11eef6..c9a5234 100644 --- a/configs/local/view_generate_config.yaml +++ b/configs/local/view_generate_config.yaml @@ -7,16 +7,16 @@ runner: name: debug root_dir: experiments generate: - port: 5003 - from: 3000 - to: -1 # -1 means all + port: 5002 + from: 2200 + to: 2300 # -1 means all object_dir: /media/hofee/data/data/scaled_object_meshes table_model_path: /media/hofee/data/data/others/table.obj - output_dir: /media/hofee/repository/new_full_data + output_dir: /media/hofee/repository/new_data_with_normal binocular_vision: true plane_size: 10 max_views: 512 - min_views: 64 + min_views: 128 random_view_ratio: 0.2 min_cam_table_included_degree: 20 max_diag: 0.7 diff --git a/preprocess/preprocessor.py b/preprocess/preprocessor.py index aa077fc..eea33c9 100644 --- a/preprocess/preprocessor.py +++ b/preprocess/preprocessor.py @@ -1,6 +1,13 @@ import os import json import numpy as np + +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 from utils.pts import PtsUtil @@ -55,8 +62,7 @@ def get_world_points(depth, cam_intrinsic, cam_extrinsic): points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] return points_camera_world -def get_world_normals(normal_image, cam_extrinsic): - normals = normal_image.reshape(-1, 3) +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 @@ -75,7 +81,7 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int return selected_points_indices -def save_scene_data(root, scene): +def save_scene_data(root, scene, scene_idx=0, scene_total=1): ''' configuration ''' target_mask_label = (0, 255, 0, 255) @@ -88,11 +94,12 @@ def save_scene_data(root, scene): ''' scan points ''' display_table_info = DataLoadUtil.get_display_table_info(root, scene) radius = display_table_info["radius"] - scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=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( @@ -104,30 +111,36 @@ def save_scene_data(root, scene): 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_L"], cam_info["cam_extrinsic_L"]) - scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic_R"], cam_info["cam_extrinsic_R"]) - scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud( + 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 ) - scene_points_R = PtsUtil.random_downsample_point_cloud( + 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( - scene_points_L, scene_points_R, voxel_size, require_idx=True + 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 ) ''' 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 = get_world_normals(normal_L, cam_info["cam_extrinsic_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] - target_normals = scene_normals_L[mask_overlap] filtered_target_points, filtered_idx = PtsUtil.filter_points( - target_points, target_normals, cam_info["cam_extrinsic_L"], filter_degree, require_idx=True + target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True ) ''' train_input_mask ''' @@ -135,7 +148,7 @@ def save_scene_data(root, scene): ''' scan points indices ''' - scan_points_indices = get_scan_points_indices(scan_points, mask_L, display_table_mask_label, cam_info["cam_intrinsic_L"], cam_info["cam_extrinsic_L"]) + scan_points_indices = get_scan_points_indices(scan_points, mask_img, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) save_full_points(root, scene, frame_id, train_input_points) save_target_points(root, scene, frame_id, filtered_target_points) @@ -146,6 +159,19 @@ def save_scene_data(root, scene): if __name__ == "__main__": - root = "" - for scene in os.listdir(root): - save_scene_data(root, scene) \ No newline at end of file + #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" + scene_list = [] + + with open(list_path, "r") as f: + for line in f: + scene_list.append(line.strip()) + + from_idx = 0 + to_idx = len(scene_list) + cnt = 0 + total = to_idx - from_idx + for scene in scene_list[from_idx:to_idx]: + save_scene_data(root, scene, cnt, total) + cnt+=1 \ No newline at end of file diff --git a/utils/data_load.py b/utils/data_load.py index e74efe4..33f65c0 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -190,7 +190,7 @@ class DataLoadUtil: mask_path = os.path.join( os.path.dirname(path), "mask", os.path.basename(path) + ".png" ) - mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) + mask_image = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) return mask_image @staticmethod @@ -199,23 +199,22 @@ class DataLoadUtil: normal_path_L = os.path.join( os.path.dirname(path), "normal", os.path.basename(path) + "_L.png" ) - normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_UNCHANGED) + normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_COLOR) normal_path_R = os.path.join( os.path.dirname(path), "normal", os.path.basename(path) + "_R.png" ) - normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_UNCHANGED) - return normal_image_L, normal_image_R + normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_COLOR) + return normal_image_L[:3,:3], normal_image_R[:3,:3] else: if binocular and left_only: normal_path = os.path.join( os.path.dirname(path), "normal", os.path.basename(path) + "_L.png" ) - else: normal_path = os.path.join( os.path.dirname(path), "normal", os.path.basename(path) + ".png" ) - normal_image = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED) + normal_image = cv2.imread(normal_path, cv2.IMREAD_COLOR) return normal_image @staticmethod diff --git a/utils/pts.py b/utils/pts.py index 1cdc659..66258ac 100644 --- a/utils/pts.py +++ b/utils/pts.py @@ -20,6 +20,8 @@ class PtsUtil: @staticmethod def random_downsample_point_cloud(point_cloud, num_points, require_idx=False): if point_cloud.shape[0] == 0: + if require_idx: + return point_cloud, np.array([]) return point_cloud idx = np.random.choice(len(point_cloud), num_points, replace=True) if require_idx: @@ -60,7 +62,11 @@ class PtsUtil: 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