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

update preprocessor

Browse Source
This commit is contained in:
hofee 2024-10-05 13:16:14 +08:00
parent fd7614c847
commit 9d6d36f5c2
2 changed files with 36 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
configs/local/view_generate_config.yaml
View File

@ -7,9 +7,9 @@ runner:
name: debug name: debug
root_dir: experiments root_dir: experiments
generate: generate:
port: 5002 port: 5005
from: 2200 from: 2300
to: 2300 # -1 means all to: 2800 # -1 means all
object_dir: /media/hofee/data/data/scaled_object_meshes object_dir: /media/hofee/data/data/scaled_object_meshes
table_model_path: /media/hofee/data/data/others/table.obj table_model_path: /media/hofee/data/data/others/table.obj
output_dir: /media/hofee/repository/new_data_with_normal output_dir: /media/hofee/repository/new_data_with_normal
@ -49,4 +49,4 @@ runner:
Light: Light:
location: [0,0,3.5] location: [0,0,3.5]
orientation: [0,0,0] orientation: [0,0,0]
power: 150 power: 150

58
preprocess/preprocessor.py
View File

@ -30,13 +30,15 @@ def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_t
os.makedirs(os.path.join(root,scene, "target_pts")) os.makedirs(os.path.join(root,scene, "target_pts"))
save_np_pts(pts_path, target_points, file_type) save_np_pts(pts_path, target_points, file_type)
def save_mask_idx(root, scene, frame_idx, mask_idx: np.ndarray,filtered_idx, file_type="txt"): def save_mask_idx(root, scene, frame_idx, mask_train_input: np.ndarray, mask_overlap, file_type="txt"):
indices_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}.{file_type}") mask_train_input_path = os.path.join(root,scene, "mask_idx", f"mask_train_input_{frame_idx}.{file_type}")
mask_overlap_path = os.path.join(root,scene, "mask_idx", f"mask_overlap_{frame_idx}.{file_type}")
if not os.path.exists(os.path.join(root,scene, "mask_idx")): if not os.path.exists(os.path.join(root,scene, "mask_idx")):
os.makedirs(os.path.join(root,scene, "mask_idx")) os.makedirs(os.path.join(root,scene, "mask_idx"))
save_np_pts(indices_path, mask_idx, file_type) save_np_pts(mask_train_input_path, mask_train_input, file_type)
filtered_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}_filtered.{file_type}") save_np_pts(mask_overlap_path, mask_overlap, file_type)
save_np_pts(filtered_path, filtered_idx, file_type) # filtered_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}_filtered.{file_type}")
# save_np_pts(filtered_path, filtered_idx, 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}")
@ -62,14 +64,9 @@ def get_world_points(depth, cam_intrinsic, cam_extrinsic):
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_camera_world return points_camera_world
def get_world_normals(normals, cam_extrinsic):
normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
normals_world = np.dot(cam_extrinsic[:3, :3], normals.T).T
return normals_world
def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic): def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_intrinsic, cam_extrinsic):
scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1)))) scan_points_homogeneous = np.hstack((scan_points, np.ones((scan_points.shape[0], 1))))
points_camera = np.dot(cam_extrinsic, scan_points_homogeneous.T).T[:, :3] points_camera = np.dot(np.linalg.inv(cam_extrinsic), scan_points_homogeneous.T).T[:, :3]
points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T points_image_homogeneous = np.dot(cam_intrinsic, points_camera.T).T
points_image_homogeneous /= points_image_homogeneous[:, 2:] points_image_homogeneous /= points_image_homogeneous[:, 2:]
pixel_x = points_image_homogeneous[:, 0].astype(int) pixel_x = points_image_homogeneous[:, 0].astype(int)
@ -77,11 +74,12 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int
h, w = mask.shape[:2] h, w = mask.shape[:2]
valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h) valid_indices = (pixel_x >= 0) & (pixel_x < w) & (pixel_y >= 0) & (pixel_y < h)
mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]] mask_colors = mask[pixel_y[valid_indices], pixel_x[valid_indices]]
selected_points_indices = mask_colors == display_table_mask_label selected_points_indices = np.where((mask_colors == display_table_mask_label).all(axis=-1))[0]
selected_points_indices = np.where(valid_indices)[0][selected_points_indices]
return selected_points_indices return selected_points_indices
def save_scene_data(root, scene, scene_idx=0, scene_total=1): def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
''' configuration ''' ''' configuration '''
target_mask_label = (0, 255, 0, 255) target_mask_label = (0, 255, 0, 255)
@ -99,7 +97,11 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
''' 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(
@ -107,8 +109,8 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
cam_info["far_plane"], cam_info["far_plane"],
binocular=True binocular=True
) )
mask_L = DataLoadUtil.load_seg(path, binocular=True, left_only=True) mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True) #normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
''' scene points ''' ''' scene points '''
scene_points_L = get_world_points(depth_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) scene_points_L = get_world_points(depth_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"])
@ -131,28 +133,30 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
) )
''' target points ''' ''' target points '''
mask_img = mask_L mask_img_L = mask_L
mask_img_R = mask_R
mask_L = mask_L.reshape(-1, 4) mask_L = mask_L.reshape(-1, 4)
mask_L = (mask_L == target_mask_label).all(axis=-1) mask_L = (mask_L == target_mask_label).all(axis=-1)
mask_overlap = mask_L[random_sample_idx_L][overlap_idx_L] mask_overlap = mask_L[random_sample_idx_L][overlap_idx_L]
scene_normals_L = normal_L.reshape(-1, 3)
target_overlap_normals = scene_normals_L[random_sample_idx_L][overlap_idx_L][mask_overlap]
target_normals = get_world_normals(target_overlap_normals, cam_info["cam_to_world"])
target_points = scene_overlap_points[mask_overlap] target_points = scene_overlap_points[mask_overlap]
filtered_target_points, filtered_idx = PtsUtil.filter_points( filtered_target_points, filtered_idx = PtsUtil.filter_points(
target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True
) )
''' train_input_mask ''' ''' train_input_mask '''
mask_train_input = mask_overlap[train_input_idx] mask_train_input = mask_overlap[train_input_idx]
''' scan points indices ''' ''' scan points indices '''
scan_points_indices = get_scan_points_indices(scan_points, mask_img, display_table_mask_label, cam_info["cam_intrinsic"], cam_info["cam_to_world"]) 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"])
save_full_points(root, scene, frame_id, train_input_points) scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R)
save_target_points(root, scene, frame_id, filtered_target_points) print(scan_points_indices.shape, scan_points_indices_L.shape, scan_points_indices_R.shape)
save_mask_idx(root, scene, frame_id, mask_train_input, filtered_idx=filtered_idx) # np.savetxt(f"{root}/{scene}/scan_points_{frame_id}_L.txt", scan_points[scan_points_indices_L])
np.savetxt(f"{root}/{scene}/scan_points_{frame_id}.txt", scan_points[scan_points_indices])
save_full_points(root, scene, frame_id, train_input_points, file_type=file_type)
save_target_points(root, scene, frame_id, target_points)
save_mask_idx(root, scene, frame_id, mask_train_input, mask_overlap,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)
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
@ -170,8 +174,10 @@ if __name__ == "__main__":
from_idx = 0 from_idx = 0
to_idx = len(scene_list) to_idx = len(scene_list)
cnt = 0 cnt = 0
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]:
save_scene_data(root, scene, cnt, total) save_scene_data(root, scene, cnt, total, "txt")
cnt+=1 cnt+=1