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

optimize preprocessor

Browse Source
This commit is contained in:
hofee
2024-10-05 12:24:53 -05:00
parent ee7537c315
commit d098c9f951
4 changed files with 104 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
preprocess/preprocessor.py
View File

@@ -1,12 +1,9 @@
import os
import json
import numpy as np
import time
import sys
np.random.seed(0)
# append parent directory to sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
print(sys.path)
from utils.reconstruction import ReconstructionUtil
from utils.data_load import DataLoadUtil
@@ -18,26 +15,12 @@ def save_np_pts(path, pts: np.ndarray, file_type="txt"):
else:
np.save(path, pts)
def save_full_points(root, scene, frame_idx, full_points: np.ndarray, file_type="txt"):
pts_path = os.path.join(root,scene, "scene_pts", f"{frame_idx}.{file_type}")
if not os.path.exists(os.path.join(root,scene, "scene_pts")):
os.makedirs(os.path.join(root,scene, "scene_pts"))
save_np_pts(pts_path, full_points, file_type)
def save_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"):
pts_path = os.path.join(root,scene, "target_pts", f"{frame_idx}.{file_type}")
if not os.path.exists(os.path.join(root,scene, "target_pts")):
os.makedirs(os.path.join(root,scene, "target_pts"))
pts_path = os.path.join(root,scene, "pts", f"{frame_idx}.{file_type}")
if not os.path.exists(os.path.join(root,scene, "pts")):
os.makedirs(os.path.join(root,scene, "pts"))
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"):
indices_path = os.path.join(root,scene, "mask_idx", f"{frame_idx}.{file_type}")
if not os.path.exists(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)
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"):
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")):
@@ -49,27 +32,36 @@ def save_scan_points(root, scene, scan_points: np.ndarray):
save_np_pts(scan_points_path, scan_points)
def get_world_points(depth, cam_intrinsic, cam_extrinsic):
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)
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)
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
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):
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 /= points_image_homogeneous[:, 2:]
pixel_x = points_image_homogeneous[:, 0].astype(int)
@@ -77,7 +69,8 @@ def get_scan_points_indices(scan_points, mask, display_table_mask_label, cam_int
h, w = mask.shape[:2]
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]]
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
@@ -86,14 +79,16 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
''' configuration '''
target_mask_label = (0, 255, 0, 255)
display_table_mask_label=(0, 0, 255, 255)
random_downsample_N = 65536
train_input_pts_num = 8192
random_downsample_N = 32768
voxel_size=0.002
filter_degree = 75
min_z = 0.2
max_z = 0.45
''' scan points '''
display_table_info = DataLoadUtil.get_display_table_info(root, scene)
radius = display_table_info["radius"]
scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
''' read frame data(depth|mask|normal) '''
@@ -107,52 +102,50 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
cam_info["far_plane"],
binocular=True
)
mask_L = DataLoadUtil.load_seg(path, binocular=True, left_only=True)
normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
''' scene points '''
scene_points_L = get_world_points(depth_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"])
scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"])
sampled_scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud(
scene_points_L, random_downsample_N, require_idx=True
)
sampled_scene_points_R = PtsUtil.random_downsample_point_cloud(
scene_points_R, random_downsample_N
)
scene_overlap_points, overlap_idx_L = PtsUtil.get_overlapping_points(
sampled_scene_points_L, sampled_scene_points_R, voxel_size, require_idx=True
)
if scene_overlap_points.shape[0] < 1024:
scene_overlap_points = sampled_scene_points_L
overlap_idx_L = np.arange(sampled_scene_points_L.shape[0])
train_input_points, train_input_idx = PtsUtil.random_downsample_point_cloud(
scene_overlap_points, train_input_pts_num, require_idx=True
)
mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
''' target points '''
mask_img = mask_L
mask_L = mask_L.reshape(-1, 4)
mask_L = (mask_L == target_mask_label).all(axis=-1)
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]
filtered_target_points, filtered_idx = PtsUtil.filter_points(
target_points, target_normals, cam_info["cam_to_world"], filter_degree, require_idx=True
mask_img_L = mask_L
mask_img_R = mask_R
target_mask_img_L = (mask_L == target_mask_label).all(axis=-1)
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 = 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:
target_points = PtsUtil.get_overlapping_points(
sampled_target_points_L, sampled_target_points_R, voxel_size
)
''' train_input_mask '''
mask_train_input = mask_overlap[train_input_idx]
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(
target_points, points_normals, cam_info["cam_to_world"],voxel_size=0.002, theta = filter_degree, z_range=(min_z, max_z)
)
''' 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"])
scan_points_indices = np.intersect1d(scan_points_indices_L, scan_points_indices_R)
save_full_points(root, scene, frame_id, train_input_points)
save_target_points(root, scene, frame_id, filtered_target_points)
save_mask_idx(root, scene, frame_id, mask_train_input, filtered_idx=filtered_idx)
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_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess
@@ -160,8 +153,8 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1):
if __name__ == "__main__":
#root = "/media/hofee/repository/new_data_with_normal"
root = "/media/hofee/repository/test_sample"
list_path = "/media/hofee/repository/test_sample/test_sample_list.txt"
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"
scene_list = []
with open(list_path, "r") as f:
@@ -171,7 +164,11 @@ if __name__ == "__main__":
from_idx = 0
to_idx = len(scene_list)
cnt = 0
import time
total = to_idx - from_idx
for scene in scene_list[from_idx:to_idx]:
start = time.time()
save_scene_data(root, scene, cnt, total)
cnt+=1
cnt+=1
end = time.time()
print(f"Time cost: {end-start}")