import os import numpy as np import time import sys np.random.seed(0) sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from concurrent.futures import ThreadPoolExecutor, as_completed from utils.reconstruction_util import ReconstructionUtil from utils.data_load import DataLoadUtil from utils.pts_util import PtsUtil from PytorchBoot.utils.log_util import Log def save_np_pts(path, pts: np.ndarray, file_type="txt"): if file_type == "txt": np.savetxt(path, pts) else: np.save(path, pts) def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"): #import ipdb;ipdb.set_trace() 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_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"): file_type="npy" 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")): os.makedirs(os.path.join(root,scene, "scan_points_indices")) save_np_pts(indices_path, scan_points_indices, file_type) 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, 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) 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_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] return points_world def get_world_points_and_normals(depth, normal, 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) normals_camera = normal[mask].reshape(-1, 3) sampled_target_points, idx = PtsUtil.random_downsample_point_cloud( points_camera, random_downsample_N, require_idx=True ) if len(sampled_target_points) == 0: return np.zeros((0, 3)), np.zeros((0, 3)) sampled_target_normals = normals_camera[idx] points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1) points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] return points_world, sampled_target_normals 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(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) pixel_y = points_image_homogeneous[:, 1].astype(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 = 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 def save_scene_data(root, scene, file_type="npy"): ''' 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 = 90 min_z = 0.25 max_z = 0.5 ''' scan points ''' display_table_info = DataLoadUtil.get_display_table_info(root, scene) # radius = display_table_info["radius"] #FIXME radius = 0.25 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): Log.info(f"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) normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=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, sampled_target_normals_L = get_world_points_and_normals(depth_L, normal_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, overlap_idx = PtsUtil.get_overlapping_points( sampled_target_points_L, sampled_target_points_R, voxel_size, require_idx=True ) target_normals = sampled_target_normals_L[overlap_idx] if has_points: has_points = target_points.shape[0] > 0 if has_points: target_points, target_normals = PtsUtil.filter_points( target_points, target_normals, cam_info["cam_to_world"], theta_limit = 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)) target_normals = 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) # 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_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 # 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) # if has_points and target_points.shape[0] > 0: # 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_L = get_scan_points_indices(scan_points, mask_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_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_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_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 # voxel_size = 0.002 # filter_degree = 75 # min_z = 0.2 # max_z = 0.5 # 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: # futures = {executor.submit(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): frame_id for frame_id in range(frame_num)} # for future in as_completed(futures): # frame_id = futures[future] # try: # future.result() # except Exception as e: # Log.error(f"Error processing frame {frame_id}: {e}") # 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"/home/yan20/nbv_rec/project/franka_control/temp_debug" # list_path = r"/media/hofee/repository/full_list.txt" # scene_list = [] # with open(list_path, "r") as f: # for line in f: # scene_list.append(line.strip()) scene_list = os.listdir(root) from_idx = 0 # 1000 to_idx = 1 # 1500 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, file_type="npy") cnt+=1 end = time.time() print(f"Time cost: {end-start}")