add save preprocessed normals

This commit is contained in:
hofee 2024-10-23 00:42:18 -05:00
parent cd56d9ea58
commit e25f7b3334
2 changed files with 11 additions and 4 deletions

View File

@ -22,6 +22,12 @@ def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_t
os.makedirs(os.path.join(root,scene, "pts")) os.makedirs(os.path.join(root,scene, "pts"))
save_np_pts(pts_path, target_points, file_type) save_np_pts(pts_path, target_points, file_type)
def save_target_normals(root, scene, frame_idx, target_normals: np.ndarray, file_type="txt"):
pts_path = os.path.join(root,scene, "normals", f"{frame_idx}.{file_type}")
if not os.path.exists(os.path.join(root,scene, "normals")):
os.makedirs(os.path.join(root,scene, "normals"))
save_np_pts(pts_path, target_normals, file_type)
def save_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"): 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}") 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")): if not os.path.exists(os.path.join(root,scene, "scan_points_indices")):
@ -135,7 +141,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
has_points = target_points.shape[0] > 0 has_points = target_points.shape[0] > 0
if has_points: if has_points:
target_points = PtsUtil.filter_points( target_points, target_normals = PtsUtil.filter_points(
target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z) target_points, sampled_target_normal_L, cam_info["cam_to_world"], theta_limit = filter_degree, z_range=(min_z, max_z)
) )
@ -149,6 +155,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
target_points = np.zeros((0, 3)) target_points = np.zeros((0, 3))
save_target_points(root, scene, frame_id, target_points, file_type=file_type) save_target_points(root, scene, frame_id, target_points, file_type=file_type)
save_target_normals(root, scene, frame_id, target_normals, file_type=file_type)
save_scan_points_indices(root, scene, frame_id, scan_points_indices, 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 save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess

View File

@ -84,14 +84,14 @@ class PtsUtil:
theta = np.arccos(cos_theta) * 180 / np.pi theta = np.arccos(cos_theta) * 180 / np.pi
idx = theta < theta_limit idx = theta < theta_limit
filtered_sampled_points = points[idx] filtered_sampled_points = points[idx]
filtered_normals = normals[idx]
""" filter with z range """ """ filter with z range """
points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose)) points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose))
idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1]) idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
z_filtered_points = filtered_sampled_points[idx] z_filtered_points = filtered_sampled_points[idx]
z_filtered_normals = filtered_normals[idx]
return z_filtered_points[:, :3] return z_filtered_points[:, :3], z_filtered_normals
@staticmethod @staticmethod
def point_to_hash(point, voxel_size): def point_to_hash(point, voxel_size):