update preprocessor.py

This commit is contained in:
hofee 2024-10-06 04:11:49 +08:00
parent dc79f4b313
commit 2af52c64e2

View File

@ -31,22 +31,6 @@ 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 old_get_world_points(depth, cam_intrinsic, cam_extrinsic):
h, w = depth.shape
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
# ----- Debug Trace ----- #
import ipdb; ipdb.set_trace()
# ------------------------ #
z = depth
x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (j - 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)
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_camera_world
def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic):
z = depth[mask]
i, j = np.nonzero(mask)
@ -94,11 +78,7 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
''' 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}")
if frame_id != 126:
continue
path = DataLoadUtil.get_path(root, scene, frame_id)
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
depth_L, depth_R = DataLoadUtil.load_depth(
@ -132,8 +112,9 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
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(
@ -149,24 +130,24 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
if not has_points:
target_points = np.zeros((0, 3))
save_target_points(root, scene, frame_id, target_points)
save_scan_points_indices(root, scene, frame_id, scan_points_indices)
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
if __name__ == "__main__":
#root = "/media/hofee/repository/new_data_with_normal"
root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample"
list_path = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample/test_sample_list.txt"
root = r"/media/hofee/repository/new_data_with_normal"
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())
from_idx = 0
to_idx = len(scene_list)
from_idx = 1000
to_idx = 1500
cnt = 0
@ -174,7 +155,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)
save_scene_data(root, scene, cnt, total, "npy")
cnt+=1
end = time.time()
print(f"Time cost: {end-start}")