223 lines
10 KiB
Python
223 lines
10 KiB
Python
import os
|
|
import numpy as np
|
|
import time
|
|
import sys
|
|
np.random.seed(0)
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from utils.reconstruction_util import ReconstructionUtil
|
|
from utils.data_load import DataLoadUtil
|
|
from utils.pts_util import PtsUtil
|
|
from PytorchBoot.utils.log_util import Log
|
|
|
|
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_target_points(root, scene, frame_idx, target_points: np.ndarray, file_type="txt"):
|
|
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_scan_points_indices(root, scene, frame_idx, scan_points_indices: np.ndarray, file_type="txt"):
|
|
file_type="npy"
|
|
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, mask, cam_intrinsic, cam_extrinsic, random_downsample_N):
|
|
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)
|
|
sampled_target_points = PtsUtil.random_downsample_point_cloud(
|
|
points_camera, random_downsample_N
|
|
)
|
|
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
|
|
points_camera_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
|
|
|
|
return points_camera_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(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)
|
|
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 = 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
|
|
|
|
def save_scene_data(root, scene, file_type="txt"):
|
|
|
|
''' configuration '''
|
|
target_mask_label = (0, 255, 0, 255)
|
|
display_table_mask_label=(0, 0, 255, 255)
|
|
random_downsample_N = 32768
|
|
voxel_size=0.002
|
|
filter_degree = 75
|
|
min_z = 0.25
|
|
max_z = 0.5
|
|
|
|
''' 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) '''
|
|
frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
|
|
for frame_id in range(frame_num):
|
|
Log.info(f"frame({frame_id}/{frame_num})]Processing {scene} frame {frame_id}")
|
|
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, mask_R = DataLoadUtil.load_seg(path, binocular=True)
|
|
|
|
''' target points '''
|
|
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)
|
|
|
|
|
|
sampled_target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N)
|
|
sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_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
|
|
)
|
|
|
|
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(
|
|
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_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)
|
|
|
|
if not has_points:
|
|
target_points = np.zeros((0, 3))
|
|
|
|
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
|
|
|
|
def process_frame(frame_id, root, scene, scan_points, file_type, target_mask_label, display_table_mask_label, random_downsample_N, voxel_size, filter_degree, min_z, max_z):
|
|
Log.info(f"[frame({frame_id})]Processing {scene} frame {frame_id}")
|
|
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, mask_R = DataLoadUtil.load_seg(path, binocular=True)
|
|
|
|
target_mask_img_L = (mask_L == target_mask_label).all(axis=-1)
|
|
target_mask_img_R = (mask_R == target_mask_label).all(axis=-1)
|
|
|
|
sampled_target_points_L = get_world_points(depth_L, target_mask_img_L, cam_info["cam_intrinsic"], cam_info["cam_to_world"], random_downsample_N)
|
|
sampled_target_points_R = get_world_points(depth_R, target_mask_img_R, cam_info["cam_intrinsic"], cam_info["cam_to_world_R"], random_downsample_N)
|
|
|
|
has_points = sampled_target_points_L.shape[0] > 0 and sampled_target_points_R.shape[0] > 0
|
|
target_points = np.zeros((0, 3))
|
|
|
|
if has_points:
|
|
target_points = PtsUtil.get_overlapping_points(sampled_target_points_L, sampled_target_points_R, voxel_size)
|
|
|
|
if has_points and target_points.shape[0] > 0:
|
|
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_L = get_scan_points_indices(scan_points, mask_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_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_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)
|
|
|
|
def save_scene_data_multithread(root, scene, file_type="txt"):
|
|
target_mask_label = (0, 255, 0, 255)
|
|
display_table_mask_label = (0, 0, 255, 255)
|
|
random_downsample_N = 32768
|
|
voxel_size = 0.002
|
|
filter_degree = 75
|
|
min_z = 0.2
|
|
max_z = 0.5
|
|
|
|
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))
|
|
frame_num = DataLoadUtil.get_scene_seq_length(root, scene)
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
futures = {executor.submit(process_frame, frame_id, root, scene, scan_points, file_type, target_mask_label, display_table_mask_label, random_downsample_N, voxel_size, filter_degree, min_z, max_z): frame_id for frame_id in range(frame_num)}
|
|
|
|
for future in as_completed(futures):
|
|
frame_id = futures[future]
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
Log.error(f"Error processing frame {frame_id}: {e}")
|
|
|
|
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"/media/hofee/data/tempdir/test_real_output"
|
|
# 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())
|
|
scene_list = os.listdir(root)
|
|
from_idx = 0 # 1000
|
|
to_idx = 1 # 1500
|
|
|
|
|
|
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, file_type="npy")
|
|
cnt+=1
|
|
end = time.time()
|
|
print(f"Time cost: {end-start}")
|