optimize preproess
This commit is contained in:
@@ -31,14 +31,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
|
||||
@@ -96,15 +99,9 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
|
||||
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 = 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)
|
||||
|
||||
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:
|
||||
@@ -138,7 +135,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
|
||||
|
||||
if __name__ == "__main__":
|
||||
#root = "/media/hofee/repository/new_data_with_normal"
|
||||
root = r"/media/hofee/data/tempdir/test_real_output"
|
||||
root = r"/media/hofee/data/data/box_output"
|
||||
# list_path = r"/media/hofee/repository/full_list.txt"
|
||||
# scene_list = []
|
||||
|
||||
@@ -147,7 +144,8 @@ if __name__ == "__main__":
|
||||
# scene_list.append(line.strip())
|
||||
scene_list = os.listdir(root)
|
||||
from_idx = 0 # 1000
|
||||
to_idx = 1 # 1500
|
||||
to_idx = len(scene_list) # 1500
|
||||
print(scene_list)
|
||||
|
||||
|
||||
cnt = 0
|
||||
@@ -155,7 +153,7 @@ if __name__ == "__main__":
|
||||
total = to_idx - from_idx
|
||||
for scene in scene_list[from_idx:to_idx]:
|
||||
start = time.time()
|
||||
save_scene_data(root, scene, cnt, total, file_type="npy")
|
||||
save_scene_data(root, scene, cnt, total, file_type="txt")
|
||||
cnt+=1
|
||||
end = time.time()
|
||||
print(f"Time cost: {end-start}")
|
||||
|
Reference in New Issue
Block a user