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

compute, load, and save covered_scan_pts

Browse Source
This commit is contained in:
hofee 2024-09-30 01:24:48 +08:00
parent cef7ab4429
commit 2633a48b4e
2 changed files with 211 additions and 104 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
runners/strategy_generator.py
View File

@ -77,28 +77,40 @@ class StrategyGenerator(Runner):
model_points_normals = DataLoadUtil.load_points_normals(root, scene_name) model_points_normals = DataLoadUtil.load_points_normals(root, scene_name)
model_pts = model_points_normals[:,:3] model_pts = model_points_normals[:,:3]
down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) 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"]
top = DataLoadUtil.get_display_table_top(root, scene_name)
scan_points = ReconstructionUtil.generate_scan_points(display_table_top=top,display_table_radius=radius)
pts_list = [] pts_list = []
scan_points_indices_list = []
for frame_idx in range(frame_num): for frame_idx in range(frame_num):
if self.load_pts and os.path.exists(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")): if self.load_pts and os.path.exists(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")):
sampled_point_cloud = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt")) sampled_point_cloud = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}.txt"))
indices = np.loadtxt(os.path.join(root,scene_name, "pts", f"{frame_idx}_indices.txt")).astype(np.int32).tolist()
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num)
pts_list.append(sampled_point_cloud) pts_list.append(sampled_point_cloud)
continue scan_points_indices_list.append(indices)
else: else:
path = DataLoadUtil.get_path(root, scene_name, frame_idx) path = DataLoadUtil.get_path(root, scene_name, frame_idx)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True) cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num) status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_idx, frame_num)
point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) point_cloud, display_table_pts = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True, get_display_table_pts=True)
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) 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)
covered_pts, indices = ReconstructionUtil.compute_covered_scan_points(scan_points, display_table_pts)
if self.save_pts: if self.save_pts:
pts_dir = os.path.join(root,scene_name, "pts") pts_dir = os.path.join(root,scene_name, "pts")
covered_pts_dir = os.path.join(pts_dir, "covered_scan_pts")
if not os.path.exists(pts_dir): if not os.path.exists(pts_dir):
os.makedirs(pts_dir) os.makedirs(pts_dir)
if not os.path.exists(covered_pts_dir):
os.makedirs(covered_pts_dir)
np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud) np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud)
np.savetxt(os.path.join(covered_pts_dir, f"{frame_idx}.txt"), covered_pts)
np.savetxt(os.path.join(pts_dir, f"{frame_idx}_indices.txt"), indices)
pts_list.append(sampled_point_cloud) 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) status_manager.set_progress("generate_strategy", "strategy_generator", "loading frame", frame_num, frame_num)
seq_num = min(self.seq_num, len(pts_list)) seq_num = min(self.seq_num, len(pts_list))

215
utils/data_load.py
View File

@ -6,8 +6,9 @@ import trimesh
import torch import torch
from utils.pts import PtsUtil from utils.pts import PtsUtil
class DataLoadUtil: class DataLoadUtil:
TABLE_POSITION = np.asarray([0,0,0.8215]) TABLE_POSITION = np.asarray([0, 0, 0.8215])
@staticmethod @staticmethod
def get_display_table_info(root, scene_name): def get_display_table_info(root, scene_name):
@ -17,8 +18,12 @@ class DataLoadUtil:
@staticmethod @staticmethod
def get_display_table_top(root, scene_name): def get_display_table_top(root, scene_name):
display_table_height = DataLoadUtil.get_display_table_info(root, scene_name)["height"] display_table_height = DataLoadUtil.get_display_table_info(root, scene_name)[
display_table_top = DataLoadUtil.TABLE_POSITION + np.asarray([0,0,display_table_height]) "height"
]
display_table_top = DataLoadUtil.TABLE_POSITION + np.asarray(
[0, 0, display_table_height]
)
return display_table_top return display_table_top
@staticmethod @staticmethod
@ -28,20 +33,20 @@ class DataLoadUtil:
@staticmethod @staticmethod
def get_label_num(root, scene_name): def get_label_num(root, scene_name):
label_dir = os.path.join(root,scene_name,"label") label_dir = os.path.join(root, scene_name, "label")
return len(os.listdir(label_dir)) return len(os.listdir(label_dir))
@staticmethod @staticmethod
def get_label_path(root, scene_name, seq_idx): def get_label_path(root, scene_name, seq_idx):
label_dir = os.path.join(root,scene_name,"label") label_dir = os.path.join(root, scene_name, "label")
if not os.path.exists(label_dir): if not os.path.exists(label_dir):
os.makedirs(label_dir) os.makedirs(label_dir)
path = os.path.join(label_dir,f"{seq_idx}.json") path = os.path.join(label_dir, f"{seq_idx}.json")
return path return path
@staticmethod @staticmethod
def get_label_path_old(root, scene_name): def get_label_path_old(root, scene_name):
path = os.path.join(root,scene_name,"label.json") path = os.path.join(root, scene_name, "label.json")
return path return path
@staticmethod @staticmethod
@ -64,7 +69,6 @@ class DataLoadUtil:
diagonal_length = np.linalg.norm(bbox) diagonal_length = np.linalg.norm(bbox)
return diagonal_length return diagonal_length
@staticmethod @staticmethod
def save_mesh_at(model_dir, output_dir, object_name, scene_name, world_object_pose): def save_mesh_at(model_dir, output_dir, object_name, scene_name, world_object_pose):
mesh = DataLoadUtil.load_mesh_at(model_dir, object_name, world_object_pose) mesh = DataLoadUtil.load_mesh_at(model_dir, object_name, world_object_pose)
@ -72,12 +76,16 @@ class DataLoadUtil:
mesh.export(model_path) mesh.export(model_path)
@staticmethod @staticmethod
def save_target_mesh_at_world_space(root, model_dir, scene_name, display_table_as_world_space_origin=True): def save_target_mesh_at_world_space(
root, model_dir, scene_name, display_table_as_world_space_origin=True
):
scene_info = DataLoadUtil.load_scene_info(root, scene_name) scene_info = DataLoadUtil.load_scene_info(root, scene_name)
target_name = scene_info["target_name"] target_name = scene_info["target_name"]
transformation = scene_info[target_name] transformation = scene_info[target_name]
if display_table_as_world_space_origin: if display_table_as_world_space_origin:
location = transformation["location"] - DataLoadUtil.get_display_table_top(root, scene_name) location = transformation["location"] - DataLoadUtil.get_display_table_top(
root, scene_name
)
else: else:
location = transformation["location"] location = transformation["location"]
rotation_euler = transformation["rotation_euler"] rotation_euler = transformation["rotation_euler"]
@ -117,7 +125,7 @@ class DataLoadUtil:
return pose_mat return pose_mat
@staticmethod @staticmethod
def load_depth(path, min_depth=0.01,max_depth=5.0,binocular=False): def load_depth(path, min_depth=0.01, max_depth=5.0, binocular=False):
def load_depth_from_real_path(real_path, min_depth, max_depth): def load_depth_from_real_path(real_path, min_depth, max_depth):
depth = cv2.imread(real_path, cv2.IMREAD_UNCHANGED) depth = cv2.imread(real_path, cv2.IMREAD_UNCHANGED)
@ -128,61 +136,83 @@ class DataLoadUtil:
return depth_meters return depth_meters
if binocular: if binocular:
depth_path_L = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + "_L.png") depth_path_L = os.path.join(
depth_path_R = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + "_R.png") os.path.dirname(path), "depth", os.path.basename(path) + "_L.png"
depth_meters_L = load_depth_from_real_path(depth_path_L, min_depth, max_depth) )
depth_meters_R = load_depth_from_real_path(depth_path_R, min_depth, max_depth) depth_path_R = os.path.join(
os.path.dirname(path), "depth", os.path.basename(path) + "_R.png"
)
depth_meters_L = load_depth_from_real_path(
depth_path_L, min_depth, max_depth
)
depth_meters_R = load_depth_from_real_path(
depth_path_R, min_depth, max_depth
)
return depth_meters_L, depth_meters_R return depth_meters_L, depth_meters_R
else: else:
depth_path = os.path.join(os.path.dirname(path), "depth", os.path.basename(path) + ".png") depth_path = os.path.join(
os.path.dirname(path), "depth", os.path.basename(path) + ".png"
)
depth_meters = load_depth_from_real_path(depth_path, min_depth, max_depth) depth_meters = load_depth_from_real_path(depth_path, min_depth, max_depth)
return depth_meters return depth_meters
@staticmethod @staticmethod
def load_seg(path, binocular=False): def load_seg(path, binocular=False):
if binocular: if binocular:
def clean_mask(mask_image): def clean_mask(mask_image):
green = [0, 255, 0, 255] green = [0, 255, 0, 255]
red = [255, 0, 0, 255] red = [255, 0, 0, 255]
threshold = 2 threshold = 2
mask_image = np.where(np.abs(mask_image - green) <= threshold, green, mask_image) mask_image = np.where(
mask_image = np.where(np.abs(mask_image - red) <= threshold, red, mask_image) np.abs(mask_image - green) <= threshold, green, mask_image
)
mask_image = np.where(
np.abs(mask_image - red) <= threshold, red, mask_image
)
return mask_image return mask_image
mask_path_L = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + "_L.png")
mask_path_L = os.path.join(
os.path.dirname(path), "mask", os.path.basename(path) + "_L.png"
)
mask_image_L = clean_mask(cv2.imread(mask_path_L, cv2.IMREAD_UNCHANGED)) mask_image_L = clean_mask(cv2.imread(mask_path_L, cv2.IMREAD_UNCHANGED))
mask_path_R = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + "_R.png") mask_path_R = os.path.join(
os.path.dirname(path), "mask", os.path.basename(path) + "_R.png"
)
mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED)) mask_image_R = clean_mask(cv2.imread(mask_path_R, cv2.IMREAD_UNCHANGED))
return mask_image_L, mask_image_R return mask_image_L, mask_image_R
else: else:
mask_path = os.path.join(os.path.dirname(path), "mask", os.path.basename(path) + ".png") mask_path = os.path.join(
os.path.dirname(path), "mask", os.path.basename(path) + ".png"
)
mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask_image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
return mask_image return mask_image
@staticmethod @staticmethod
def load_label(path): def load_label(path):
with open(path, 'r') as f: with open(path, "r") as f:
label_data = json.load(f) label_data = json.load(f)
return label_data return label_data
@staticmethod @staticmethod
def load_rgb(path): def load_rgb(path):
rgb_path = os.path.join(os.path.dirname(path), "rgb", os.path.basename(path) + ".png") rgb_path = os.path.join(
os.path.dirname(path), "rgb", os.path.basename(path) + ".png"
)
rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR) rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
return rgb_image return rgb_image
@staticmethod @staticmethod
def load_from_preprocessed_pts(path): def load_from_preprocessed_pts(path):
npy_path = os.path.join(os.path.dirname(path), "points", os.path.basename(path) + ".npy") npy_path = os.path.join(
os.path.dirname(path), "points", os.path.basename(path) + ".npy"
)
pts = np.load(npy_path) pts = np.load(npy_path)
return pts return pts
@staticmethod @staticmethod
def cam_pose_transformation(cam_pose_before): def cam_pose_transformation(cam_pose_before):
offset = np.asarray([ offset = np.asarray([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
[1, 0, 0, 0],
[0, -1, 0, 0],
[0, 0, -1, 0],
[0, 0, 0, 1]])
cam_pose_after = cam_pose_before @ offset cam_pose_after = cam_pose_before @ offset
return cam_pose_after return cam_pose_after
@ -191,13 +221,17 @@ class DataLoadUtil:
scene_dir = os.path.dirname(path) scene_dir = os.path.dirname(path)
root_dir = os.path.dirname(scene_dir) root_dir = os.path.dirname(scene_dir)
scene_name = os.path.basename(scene_dir) scene_name = os.path.basename(scene_dir)
camera_params_path = os.path.join(os.path.dirname(path), "camera_params", os.path.basename(path) + ".json") camera_params_path = os.path.join(
with open(camera_params_path, 'r') as f: os.path.dirname(path), "camera_params", os.path.basename(path) + ".json"
)
with open(camera_params_path, "r") as f:
label_data = json.load(f) label_data = json.load(f)
cam_to_world = np.asarray(label_data["extrinsic"]) cam_to_world = np.asarray(label_data["extrinsic"])
cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world) cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world)
world_to_display_table = np.eye(4) world_to_display_table = np.eye(4)
world_to_display_table[:3, 3] = - DataLoadUtil.get_display_table_top(root_dir, scene_name) world_to_display_table[:3, 3] = -DataLoadUtil.get_display_table_top(
root_dir, scene_name
)
if display_table_as_world_space_origin: if display_table_as_world_space_origin:
cam_to_world = np.dot(world_to_display_table, cam_to_world) cam_to_world = np.dot(world_to_display_table, cam_to_world)
cam_intrinsic = np.asarray(label_data["intrinsic"]) cam_intrinsic = np.asarray(label_data["intrinsic"])
@ -205,7 +239,7 @@ class DataLoadUtil:
"cam_to_world": cam_to_world, "cam_to_world": cam_to_world,
"cam_intrinsic": cam_intrinsic, "cam_intrinsic": cam_intrinsic,
"far_plane": label_data["far_plane"], "far_plane": label_data["far_plane"],
"near_plane": label_data["near_plane"] "near_plane": label_data["near_plane"],
} }
if binocular: if binocular:
cam_to_world_R = np.asarray(label_data["extrinsic_R"]) cam_to_world_R = np.asarray(label_data["extrinsic_R"])
@ -220,7 +254,9 @@ class DataLoadUtil:
return cam_info return cam_info
@staticmethod @staticmethod
def get_real_cam_O_from_cam_L(cam_L, cam_O_to_cam_L, scene_path, display_table_as_world_space_origin=True): def get_real_cam_O_from_cam_L(
cam_L, cam_O_to_cam_L, scene_path, display_table_as_world_space_origin=True
):
root_dir = os.path.dirname(scene_path) root_dir = os.path.dirname(scene_path)
scene_name = os.path.basename(scene_path) scene_name = os.path.basename(scene_path)
if isinstance(cam_L, torch.Tensor): if isinstance(cam_L, torch.Tensor):
@ -228,71 +264,126 @@ class DataLoadUtil:
nO_to_display_table_pose = cam_L @ cam_O_to_cam_L nO_to_display_table_pose = cam_L @ cam_O_to_cam_L
if display_table_as_world_space_origin: if display_table_as_world_space_origin:
display_table_to_world = np.eye(4) display_table_to_world = np.eye(4)
display_table_to_world[:3, 3] = DataLoadUtil.get_display_table_top(root_dir, scene_name) display_table_to_world[:3, 3] = DataLoadUtil.get_display_table_top(
root_dir, scene_name
)
nO_to_world_pose = np.dot(display_table_to_world, nO_to_display_table_pose) nO_to_world_pose = np.dot(display_table_to_world, nO_to_display_table_pose)
nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose) nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose)
return nO_to_world_pose return nO_to_world_pose
@staticmethod @staticmethod
def get_target_point_cloud(depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0,255,0,255)): def get_target_point_cloud(
depth, cam_intrinsic, cam_extrinsic, mask, target_mask_label=(0, 255, 0, 255)
):
h, w = depth.shape h, w = depth.shape
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy') i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
z = depth z = depth
x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0] x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1] 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 = np.stack((x, y, z), axis=-1).reshape(-1, 3)
mask = mask.reshape(-1,4) mask = mask.reshape(-1, 4)
target_mask = (mask == target_mask_label).all(axis=-1) target_mask = (mask == target_mask_label).all(axis=-1)
target_points_camera = points_camera[target_mask] target_points_camera = points_camera[target_mask]
target_points_camera_aug = np.concatenate([target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1) target_points_camera_aug = np.concatenate(
[target_points_camera, np.ones((target_points_camera.shape[0], 1))], axis=-1
)
target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3] target_points_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
return { return {
"points_world": target_points_world, "points_world": target_points_world,
"points_camera": target_points_camera "points_camera": target_points_camera,
} }
@staticmethod @staticmethod
def get_point_cloud(depth, cam_intrinsic, cam_extrinsic): def get_point_cloud(depth, cam_intrinsic, cam_extrinsic):
h, w = depth.shape h, w = depth.shape
i, j = np.meshgrid(np.arange(w), np.arange(h), indexing='xy') i, j = np.meshgrid(np.arange(w), np.arange(h), indexing="xy")
z = depth z = depth
x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0] x = (i - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (j - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1] 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 = 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_aug = np.concatenate(
[points_camera, np.ones((points_camera.shape[0], 1))], axis=-1
)
points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3] points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return { return {"points_world": points_world, "points_camera": points_camera}
"points_world": points_world,
"points_camera": points_camera
}
@staticmethod @staticmethod
def get_target_point_cloud_world_from_path(path, binocular=False, random_downsample_N=65536, voxel_size = 0.005, target_mask_label=(0,255,0,255)): def get_target_point_cloud_world_from_path(
path,
binocular=False,
random_downsample_N=65536,
voxel_size=0.005,
target_mask_label=(0, 255, 0, 255),
display_table_mask_label=(255, 0, 0, 255),
get_display_table_pts=False
):
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular) cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
if binocular: if binocular:
depth_L, depth_R = DataLoadUtil.load_depth(path, cam_info['near_plane'], cam_info['far_plane'], 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) mask_L, mask_R = DataLoadUtil.load_seg(path, binocular=True)
point_cloud_L = DataLoadUtil.get_target_point_cloud(depth_L, cam_info['cam_intrinsic'], cam_info['cam_to_world'], mask_L, target_mask_label)['points_world'] point_cloud_L = DataLoadUtil.get_target_point_cloud(
point_cloud_R = DataLoadUtil.get_target_point_cloud(depth_R, cam_info['cam_intrinsic'], cam_info['cam_to_world_R'], mask_R, target_mask_label)['points_world'] depth_L,
point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, random_downsample_N) cam_info["cam_intrinsic"],
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, random_downsample_N) cam_info["cam_to_world"],
overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R, voxel_size) mask_L,
target_mask_label,
)["points_world"]
point_cloud_R = DataLoadUtil.get_target_point_cloud(
depth_R,
cam_info["cam_intrinsic"],
cam_info["cam_to_world_R"],
mask_R,
target_mask_label,
)["points_world"]
point_cloud_L = PtsUtil.random_downsample_point_cloud(
point_cloud_L, random_downsample_N
)
point_cloud_R = PtsUtil.random_downsample_point_cloud(
point_cloud_R, random_downsample_N
)
overlap_points = DataLoadUtil.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_overlap = DataLoadUtil.get_overlapping_points(
display_pts_L, display_pts_R, voxel_size
)
return overlap_points, display_pts_overlap
return overlap_points return overlap_points
else: else:
depth = DataLoadUtil.load_depth(path, cam_info['near_plane'], cam_info['far_plane']) depth = DataLoadUtil.load_depth(
path, cam_info["near_plane"], cam_info["far_plane"]
)
mask = DataLoadUtil.load_seg(path) mask = DataLoadUtil.load_seg(path)
point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info['cam_intrinsic'], cam_info['cam_to_world'], mask)['points_world'] point_cloud = DataLoadUtil.get_target_point_cloud(
depth, cam_info["cam_intrinsic"], cam_info["cam_to_world"], mask
)["points_world"]
return point_cloud return point_cloud
@staticmethod @staticmethod
def voxelize_points(points, voxel_size): def voxelize_points(points, voxel_size):
@ -305,10 +396,12 @@ class DataLoadUtil:
voxels_L, indices_L = DataLoadUtil.voxelize_points(point_cloud_L, voxel_size) voxels_L, indices_L = DataLoadUtil.voxelize_points(point_cloud_L, voxel_size)
voxels_R, _ = DataLoadUtil.voxelize_points(point_cloud_R, voxel_size) voxels_R, _ = DataLoadUtil.voxelize_points(point_cloud_R, voxel_size)
voxel_indices_L = voxels_L.view([('', voxels_L.dtype)]*3) voxel_indices_L = voxels_L.view([("", voxels_L.dtype)] * 3)
voxel_indices_R = voxels_R.view([('', voxels_R.dtype)]*3) voxel_indices_R = voxels_R.view([("", voxels_R.dtype)] * 3)
overlapping_voxels = np.intersect1d(voxel_indices_L, voxel_indices_R) 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]) mask_L = np.isin(
indices_L, np.where(np.isin(voxel_indices_L, overlapping_voxels))[0]
)
overlapping_points = point_cloud_L[mask_L] overlapping_points = point_cloud_L[mask_L]
return overlapping_points return overlapping_points
@ -317,5 +410,7 @@ class DataLoadUtil:
points_path = os.path.join(root, scene_name, "points_and_normals.txt") points_path = os.path.join(root, scene_name, "points_and_normals.txt")
points_normals = np.loadtxt(points_path) points_normals = np.loadtxt(points_path)
if display_table_as_world_space_origin: if display_table_as_world_space_origin:
points_normals[:,:3] = points_normals[:,:3] - DataLoadUtil.get_display_table_top(root, scene_name) points_normals[:, :3] = points_normals[
:, :3
] - DataLoadUtil.get_display_table_top(root, scene_name)
return points_normals return points_normals