hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update preprocessor.py

Browse Source
This commit is contained in:
hofee 2024-10-06 04:11:49 +08:00
parent dc79f4b313
commit 2af52c64e2
1 changed files with 10 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
preprocess/preprocessor.py
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") scan_points_path = os.path.join(root,scene, "scan_points.txt")
save_np_pts(scan_points_path, scan_points) 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): def get_world_points(depth, mask, cam_intrinsic, cam_extrinsic):
z = depth[mask] z = depth[mask]
i, j = np.nonzero(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) ''' ''' read frame data(depth|mask|normal) '''
frame_num = DataLoadUtil.get_scene_seq_length(root, scene) frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
for frame_id in range(frame_num): for frame_id in range(frame_num):
print(f"[scene({scene_idx}/{scene_total})|frame({frame_id}/{frame_num})]Processing {scene} frame {frame_id}") 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) path = DataLoadUtil.get_path(root, scene, frame_id)
cam_info = DataLoadUtil.load_cam_info(path, binocular=True) cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
depth_L, depth_R = DataLoadUtil.load_depth( 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 sampled_target_points_L, sampled_target_points_R, voxel_size
) )
if has_points:
has_points = target_points.shape[0] > 0
has_points = target_points.shape[0] > 0
if has_points: if has_points:
points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True) points_normals = DataLoadUtil.load_points_normals(root, scene, display_table_as_world_space_origin=True)
target_points = PtsUtil.filter_points( 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: if not has_points:
target_points = np.zeros((0, 3)) target_points = np.zeros((0, 3))
save_target_points(root, scene, frame_id, target_points) save_target_points(root, scene, frame_id, target_points, file_type=file_type)
save_scan_points_indices(root, scene, frame_id, scan_points_indices) 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
if __name__ == "__main__": if __name__ == "__main__":
#root = "/media/hofee/repository/new_data_with_normal" #root = "/media/hofee/repository/new_data_with_normal"
root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample" root = r"/media/hofee/repository/new_data_with_normal"
list_path = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test\test_sample/test_sample_list.txt" list_path = r"/media/hofee/repository/full_list.txt"
scene_list = [] scene_list = []
with open(list_path, "r") as f: with open(list_path, "r") as f:
for line in f: for line in f:
scene_list.append(line.strip()) scene_list.append(line.strip())
from_idx = 0 from_idx = 1000
to_idx = len(scene_list) to_idx = 1500
cnt = 0 cnt = 0
@ -174,7 +155,7 @@ if __name__ == "__main__":
total = to_idx - from_idx total = to_idx - from_idx
for scene in scene_list[from_idx:to_idx]: for scene in scene_list[from_idx:to_idx]:
start = time.time() start = time.time()
save_scene_data(root, scene, cnt, total) save_scene_data(root, scene, cnt, total, "npy")
cnt+=1 cnt+=1
end = time.time() end = time.time()
print(f"Time cost: {end-start}") print(f"Time cost: {end-start}")