optimize
This commit is contained in:
hofee
parent
f6c4db859e
commit
8fd2d6b1e1
@ -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"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user