add preprocess

This commit is contained in:
hofee 2024-10-03 01:59:13 +08:00
parent f460e6e6b2
commit d7561738c6
8 changed files with 243 additions and 142 deletions

View File

@ -166,7 +166,7 @@ class NBVReconstructionDataset(BaseDataset):
point_cloud_R = PtsUtil.random_downsample_point_cloud(
point_cloud_R, 65536
)
overlap_points = DataLoadUtil.get_overlapping_points(
overlap_points = PtsUtil.get_overlapping_points(
point_cloud_L, point_cloud_R
)
downsampled_target_point_cloud = (

View File

@ -89,7 +89,7 @@ class SeqNBVReconstructionDataset(BaseDataset):
first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536)
first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536)
first_overlap_points = DataLoadUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R)
first_overlap_points = PtsUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R)
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num)
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3]))

151
preprocess/preprocessor.py Normal file
View File

@ -0,0 +1,151 @@
import os
import json
import numpy as np
from utils.reconstruction import ReconstructionUtil
from utils.data_load import DataLoadUtil
from utils.pts import PtsUtil
def save_np_pts(path, pts: np.ndarray, file_type="txt"):
if file_type == "txt":
np.savetxt(path, pts)
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"))
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")):
os.makedirs(os.path.join(root,scene, "scan_points_indices"))
save_np_pts(indices_path, scan_points_indices, file_type)
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, cam_intrinsic, cam_extrinsic):
h, w = depth.shape
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
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_normals(normal_image, cam_extrinsic):
normals = normal_image.reshape(-1, 3)
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_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)
pixel_y = points_image_homogeneous[:, 1].astype(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
return selected_points_indices
def save_scene_data(root, scene):
''' 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
voxel_size=0.002
filter_degree = 75
''' scan points '''
display_table_info = DataLoadUtil.get_display_table_info(root, scene)
radius = display_table_info["radius"]
scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)
''' read frame data(depth|mask|normal) '''
frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
for frame_id in range(frame_num):
path = DataLoadUtil.get_path(root, scene, frame_id)
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
depth_L, depth_R = DataLoadUtil.load_depth(
path, cam_info["near_plane"],
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_L"], cam_info["cam_extrinsic_L"])
scene_points_R = get_world_points(depth_R, cam_info["cam_intrinsic_R"], cam_info["cam_extrinsic_R"])
scene_points_L, random_sample_idx_L = PtsUtil.random_downsample_point_cloud(
scene_points_L, random_downsample_N, require_idx=True
)
scene_points_R = PtsUtil.random_downsample_point_cloud(
scene_points_R, random_downsample_N
)
scene_overlap_points, overlap_idx_L = PtsUtil.get_overlapping_points(
scene_points_L, scene_points_R, voxel_size, require_idx=True
)
train_input_points, train_input_idx = PtsUtil.random_downsample_point_cloud(
scene_overlap_points, train_input_pts_num, require_idx=True
)
''' target points '''
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 = get_world_normals(normal_L, cam_info["cam_extrinsic_L"])
target_points = scene_overlap_points[mask_overlap]
target_normals = scene_normals_L[mask_overlap]
filtered_target_points, filtered_idx = PtsUtil.filter_points(
target_points, target_normals, cam_info["cam_extrinsic_L"], filter_degree, require_idx=True
)
''' train_input_mask '''
mask_train_input = mask_overlap[train_input_idx]
''' scan points indices '''
scan_points_indices = get_scan_points_indices(scan_points, mask_L, display_table_mask_label, cam_info["cam_intrinsic_L"], cam_info["cam_extrinsic_L"])
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)
save_scan_points_indices(root, scene, frame_id, scan_points_indices)
save_scan_points(root, scene, scan_points) # The "done" flag of scene preprocess
if __name__ == "__main__":
root = ""
for scene in os.listdir(root):
save_scene_data(root, scene)

View File

@ -52,9 +52,8 @@ class StrategyGenerator(Runner):
for scene_name in scene_name_list[from_idx:to_idx]:
Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}")
status_manager.set_progress("generate_strategy", "strategy_generator", "scene", cnt, total)
#diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name)
voxel_threshold = 0.002
status_manager.set_status("generate_strategy", "strategy_generator", "voxel_threshold", voxel_threshold)
diag = DataLoadUtil.get_bbox_diag(model_dir, scene_name)
status_manager.set_status("generate_strategy", "strategy_generator", "diagonal", diag)
output_label_path = DataLoadUtil.get_label_path(root_dir, scene_name,0)
if os.path.exists(output_label_path) and not self.overwrite:
Log.info(f"Scene <{scene_name}> Already Exists, Skip")
@ -82,69 +81,14 @@ class StrategyGenerator(Runner):
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
model_pts = model_points_normals[:,:3]
down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
radius = display_table_info["radius"]
scan_points_path = os.path.join(root,scene_name, "scan_points.txt")
if os.path.exists(scan_points_path):
scan_points = np.loadtxt(scan_points_path)
else:
scan_points = ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius)
np.savetxt(scan_points_path, scan_points)
pts_list = []
scan_points_indices_list = []
non_zero_cnt = 0
for frame_idx in range(frame_num):
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num)
pts_path = os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")
if self.load_pts and pts_path:
with open(pts_path, 'r') as f:
pts_str = f.read()
if pts_str == "":
sampled_point_cloud = np.asarray([])
else:
pts_path = os.path.join(root,scene_name, "target_pts", f"{frame_idx}.txt")
sampled_point_cloud = np.loadtxt(pts_path)
indices_path = os.path.join(root,scene_name, "covered_scan_pts", f"{frame_idx}_indices.txt")
with open(indices_path, 'r') as f:
indices_str = f.read()
if indices_str == "":
indices = []
else:
indices = np.loadtxt(indices_path).astype(np.int32).tolist()
if isinstance(indices, int):
indices = [indices]
pts_list.append(sampled_point_cloud)
if sampled_point_cloud.shape[0] != 0:
non_zero_cnt += 1
scan_points_indices_list.append(indices)
else:
path = DataLoadUtil.get_path(root, scene_name, frame_idx)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
point_cloud, display_table_pts = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, get_display_table_pts=True)
if point_cloud.shape[0] != 0:
sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree)
non_zero_cnt += 1
else:
sampled_point_cloud = point_cloud
covered_pts, indices = ReconstructionUtil.compute_covered_scan_points(scan_points, display_table_pts)
if self.save_pts:
pts_dir = os.path.join(root,scene_name, "pts")
#display_dir = os.path.join(root,scene_name, "display_pts")
covered_pts_dir = os.path.join(root,scene_name, "covered_scan_pts")
if not os.path.exists(pts_dir):
os.makedirs(pts_dir)
if not os.path.exists(covered_pts_dir):
os.makedirs(covered_pts_dir)
# if not os.path.exists(display_dir):
# os.makedirs(display_dir)
np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud)
#np.savetxt(os.path.join(display_dir, f"{frame_idx}.txt"), display_table_pts)
np.savetxt(os.path.join(covered_pts_dir, f"{frame_idx}.txt"), covered_pts)
np.savetxt(os.path.join(covered_pts_dir, f"{frame_idx}_indices.txt"), indices)
indices = None # ReconstructionUtil.compute_covered_scan_points(scan_points, display_table_pts)
pts_list.append(sampled_point_cloud)
scan_points_indices_list.append(indices)
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num)

View File

@ -157,8 +157,8 @@ class DataLoadUtil:
return depth_meters
@staticmethod
def load_seg(path, binocular=False):
if binocular:
def load_seg(path, binocular=False, left_only=False):
if binocular and not left_only:
def clean_mask(mask_image):
green = [0, 255, 0, 255]
@ -181,6 +181,11 @@ class DataLoadUtil:
)
mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED))
return mask_image_L, mask_image_R
else:
if binocular and left_only:
mask_path = os.path.join(
os.path.dirname(path), "mask", os.path.basename(path) + "_L.png"
)
else:
mask_path = os.path.join(
os.path.dirname(path), "mask", os.path.basename(path) + ".png"
@ -188,6 +193,31 @@ class DataLoadUtil:
mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
return mask_image
@staticmethod
def load_normal(path, binocular=False, left_only=False):
if binocular and not left_only:
normal_path_L = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
)
normal_image_L = cv2.imread(normal_path_L, cv2.IMREAD_UNCHANGED)
normal_path_R = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_R.png"
)
normal_image_R = cv2.imread(normal_path_R, cv2.IMREAD_UNCHANGED)
return normal_image_L, normal_image_R
else:
if binocular and left_only:
normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + "_L.png"
)
else:
normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + ".png"
)
normal_image = cv2.imread(normal_path, cv2.IMREAD_UNCHANGED)
return normal_image
@staticmethod
def load_label(path):
with open(path, "r") as f:
@ -273,7 +303,7 @@ class DataLoadUtil:
@staticmethod
def get_target_point_cloud(
depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0, 255, 0, 255)
depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0, 255, 0, 255), require_full_points=False
):
h, w = depth.shape
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
@ -293,10 +323,11 @@ class DataLoadUtil:
)
target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
return {
data = {
"points_world": target_points_world,
"points_camera": target_points_camera,
}
return data
@staticmethod
def get_point_cloud(depth, cam_intrinsic, cam_extrinsic):
@ -323,7 +354,8 @@ class DataLoadUtil:
voxel_size=0.005,
target_mask_label=(0, 255, 0, 255),
display_table_mask_label=(0, 0, 255, 255),
get_display_table_pts=False
get_display_table_pts=False,
require_normal=False,
):
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
if binocular:
@ -351,34 +383,9 @@ class DataLoadUtil:
point_cloud_R = PtsUtil.random_downsample_point_cloud(
point_cloud_R, random_downsample_N
)
overlap_points = DataLoadUtil.get_overlapping_points(
overlap_points = PtsUtil.get_overlapping_points(
point_cloud_L, point_cloud_R, voxel_size
)
if get_display_table_pts:
display_pts_L = DataLoadUtil.get_target_point_cloud(
depth_L,
cam_info["cam_intrinsic"],
cam_info["cam_to_world"],
mask_L,
display_table_mask_label,
)["points_world"]
display_pts_R = DataLoadUtil.get_target_point_cloud(
depth_R,
cam_info["cam_intrinsic"],
cam_info["cam_to_world_R"],
mask_R,
display_table_mask_label,
)["points_world"]
display_pts_L = PtsUtil.random_downsample_point_cloud(
display_pts_L, random_downsample_N
)
point_cloud_R = PtsUtil.random_downsample_point_cloud(
display_pts_R, random_downsample_N
)
display_pts_overlap = DataLoadUtil.get_overlapping_points(
display_pts_L, display_pts_R, voxel_size
)
return overlap_points, display_pts_overlap
return overlap_points
else:
depth = DataLoadUtil.load_depth(
@ -390,27 +397,6 @@ class DataLoadUtil:
)["points_world"]
return point_cloud
@staticmethod
def voxelize_points(points, voxel_size):
voxel_indices = np.floor(points / voxel_size).astype(np.int32)
unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
return unique_voxels
@staticmethod
def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005):
voxels_L, indices_L = DataLoadUtil.voxelize_points(point_cloud_L, voxel_size)
voxels_R, _ = DataLoadUtil.voxelize_points(point_cloud_R, voxel_size)
voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
mask_L = np.isin(
indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
)
overlapping_points = point_cloud_L[mask_L]
return overlapping_points
@staticmethod
def load_points_normals(root, scene_name, display_table_as_world_space_origin=True):
points_path = os.path.join(root, scene_name, "points_and_normals.txt")

View File

@ -18,13 +18,49 @@ class PtsUtil:
return points_h[:, :3]
@staticmethod
def random_downsample_point_cloud(point_cloud, num_points):
def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
if point_cloud.shape[0] == 0:
return point_cloud
idx = np.random.choice(len(point_cloud), num_points, replace=True)
if require_idx:
return point_cloud[idx], idx
return point_cloud[idx]
@staticmethod
def random_downsample_point_cloud_tensor(point_cloud, num_points):
idx = torch.randint(0, len(point_cloud), (num_points,))
return point_cloud[idx]
@staticmethod
def voxelize_points(points, voxel_size):
voxel_indices = np.floor(points / voxel_size).astype(np.int32)
unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
return unique_voxels
@staticmethod
def get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size=0.005, require_idx=False):
voxels_L, indices_L = PtsUtil.voxelize_points(point_cloud_L, voxel_size)
voxels_R, _ = PtsUtil.voxelize_points(point_cloud_R, voxel_size)
voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R)
mask_L = np.isin(
indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
)
overlapping_points = point_cloud_L[mask_L]
if require_idx:
return overlapping_points, mask_L
return overlapping_points
@staticmethod
def filter_points(points, normals, cam_pose, theta=75, require_idx=False):
camera_axis = -cam_pose[:3, 2]
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
cos_theta = np.dot(normals_normalized, camera_axis)
theta_rad = np.deg2rad(theta)
idx = cos_theta > np.cos(theta_rad)
filtered_points= points[idx]
if require_idx:
return filtered_points, idx
return filtered_points

View File

@ -130,21 +130,6 @@ class ReconstructionUtil:
sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list))
return view_sequence, remaining_views, down_sampled_combined_point_cloud
@staticmethod
def filter_points(points, points_normals, cam_pose, voxel_size=0.005, theta=75):
sampled_points = PtsUtil.voxel_downsample_point_cloud(points, voxel_size)
kdtree = cKDTree(points_normals[:,:3])
_, indices = kdtree.query(sampled_points)
nearest_points = points_normals[indices]
normals = nearest_points[:, 3:]
camera_axis = -cam_pose[:3, 2]
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
cos_theta = np.dot(normals_normalized, camera_axis)
theta_rad = np.deg2rad(theta)
filtered_sampled_points= sampled_points[cos_theta > np.cos(theta_rad)]
return filtered_sampled_points[:, :3]
@staticmethod
def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 100, max_attempts = 1000):

View File

@ -33,12 +33,11 @@ class RenderUtil:
print(result.stderr)
return None
path = os.path.join(temp_dir, "tmp")
# ------ Debug Start ------
# import ipdb;ipdb.set_trace()
# ------ Debug End ------
point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
filtered_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree)
''' TODO: old code: filter_points api is changed, need to update the code '''
filtered_point_cloud = PtsUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree)
full_scene_point_cloud = None
if require_full_scene:
depth_L, depth_R = DataLoadUtil.load_depth(path, cam_params['near_plane'], cam_params['far_plane'], binocular=True)
@ -47,7 +46,7 @@ class RenderUtil:
point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536)
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536)
full_scene_point_cloud = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R)
full_scene_point_cloud = PtsUtil.get_overlapping_points(point_cloud_L, point_cloud_R)
return filtered_point_cloud, full_scene_point_cloud