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

inference finish

Browse Source
This commit is contained in:
hofee
2025-01-15 14:42:54 +08:00
commit 6bb8835d00
30 changed files with 1158013 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
utils/communicate_util.py Normal file
View File

@@ -0,0 +1,41 @@
import requests
import numpy as np
class CommunicateUtil:
VIEW_HOST = "192.168.1.2:7999" #"10.7.250.52:7999" ##
INFERENCE_HOST = "127.0.0.1"
INFERENCE_PORT = 5000
def get_view_data(init = False) -> dict:
params = {}
params["create_scanner"] = init
response = requests.get(f"http://{CommunicateUtil.VIEW_HOST}/api/data", json=params)
data = response.json()
if not data["success"]:
print(f"Failed to get view data")
return None
image_id = data["image_id"]
depth_image = np.array(data["depth_image"], dtype=np.uint16)
color_image = np.array(data["color_image"]).astype(np.uint8)
depth_intrinsics = data["depth_intrinsics"]
depth_extrinsics = np.array(data["depth_extrinsics"])
color_intrinsics = data["color_intrinsics"]
depth_to_color = np.array(data["depth_to_color"])
view_data = {
"image_id": image_id,
"depth_image": depth_image,
"depth_intrinsics": depth_intrinsics,
"depth_extrinsics": depth_extrinsics,
"color_image": color_image,
"color_intrinsics": color_intrinsics,
"depth_to_color": depth_to_color
}
return view_data
def get_inference_data(view_data:dict) -> dict:
data = {}
return data

284
utils/control_util.py Normal file
View File

@@ -0,0 +1,284 @@
import numpy as np
from frankapy import FrankaArm
from autolab_core import RigidTransform
from scipy.spatial.transform import Rotation as R
import serial
import time
class ControlUtil:
__fa:FrankaArm = None
__ser: serial.Serial = None
curr_rotation = 0
# BASE_TO_WORLD:np.ndarray = np.asarray([
# [1, 0, 0, -0.68834619],
# [0, 1, 0, -0.01785319],
# [0, 0, 1, -0.11924244],
# [0, 0, 0, 1]
# ])
# BASELINE_DIS = 0.1
# LEFT_TO_RIGHT_ANGLE = -25.0
# CAMERA_TO_LEFT_CAMERA:np.ndarray = np.asarray([
# [1, 0, 0, BASELINE_DIS/2 * np.cos(np.radians(-LEFT_TO_RIGHT_ANGLE/2))],
# [0, 1, 0, 0],
# [0, 0, 1, BASELINE_DIS/2 * np.sin(np.radians(-LEFT_TO_RIGHT_ANGLE/2))],
# [0, 0, 0, 1]
# ])
# CAMERA_TO_LEFT_CAMERA[:3, :3] = R.from_euler('y', LEFT_TO_RIGHT_ANGLE/2, degrees=True).as_matrix()
CAMERA_CORRECTION_ANGLE = -13
CAMERA_CORRECTION:np.ndarray = np.eye(4)
CAMERA_CORRECTION[:3, :3] = R.from_euler('y', CAMERA_CORRECTION_ANGLE, degrees=True).as_matrix()
BASE_TO_WORLD:np.ndarray = np.asarray([
[1, 0, 0, -0.49602172],
[0, 1, 0, 0.06688724],
[0, 0, 1, -0.2488704 ],
[0, 0, 0, 1]
])
CAMERA_TO_GRIPPER:np.ndarray = np.asarray(
[[-0.00757086, -0.99984193, 0.01608569, 0.02644546],
[ 0.99993362, -0.00770935, -0.00856502, -0.04860572],
[ 0.00868767, 0.01601977, 0.99983391, -0.02169477],
[ 0., 0. , 0. , 1. , ]
])
INIT_GRIPPER_POSE:np.ndarray = np.asarray(
[[0.41869126 ,0.87596275 , 0.23951774 , 0.36005292],
[ 0.70787907 ,-0.4800251 , 0.51813998 ,-0.26499909],
[ 0.56884584, -0.04739109 ,-0.82107382 ,0.66881103],
[ 0. , 0. , 0. , 1. ]])
# INIT_GRIPPER_POSE:np.ndarray = np.asarray(
# [[1 ,0 , 0 , 0.5],
# [ 0 ,-1 , 0 ,0],
# [ 0, 0 ,-1 ,0.6],
# [ 0. , 0. , 0. , 1. ]])
INIT_JOINTS:np.ndarray = np.asarray([-1.64987928,-0.73557812,-0.02676473,-2.08365335,0.01794927,1.5182423,0.73276058])
@staticmethod
def connect_robot():
if ControlUtil.__fa is None:
ControlUtil.__fa = FrankaArm(robot_num=1)
if ControlUtil.__ser is None:
ControlUtil.__ser = serial.Serial(port="/dev/ttyUSB0", baudrate=115200)
@staticmethod
def franka_reset() -> None:
#ControlUtil.__fa.goto_joints(ControlUtil.INIT_JOINTS, duration=5, use_impedance=False, ignore_errors=False)
ControlUtil.__fa.reset_joints()
@staticmethod
def init() -> None:
ControlUtil.franka_reset()
ControlUtil.set_gripper_pose(ControlUtil.INIT_GRIPPER_POSE)
@staticmethod
def get_pose() -> np.ndarray:
gripper_to_base = ControlUtil.get_curr_gripper_to_base_pose()
cam_to_world = ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
return cam_to_world
@staticmethod
def set_pose(cam_to_world: np.ndarray) -> None:
gripper_to_base = ControlUtil.solve_gripper_to_base(cam_to_world)
ControlUtil.set_gripper_pose(gripper_to_base)
@staticmethod
def rotate_display_table(degree):
turn_directions = {
"left": 1,
"right": 0
}
delta_degree = degree - ControlUtil.curr_rotation
ControlUtil.curr_rotation = degree
print(f"Table rotated {ControlUtil.curr_rotation} degree")
if delta_degree >= 0:
turn_angle = delta_degree
turn_direction = turn_directions["left"]
else:
turn_angle = -delta_degree
turn_direction = turn_directions["right"]
write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({turn_direction},{turn_angle});".encode('utf-8'))
print(f"send command: CT+TRUNSINGLE({turn_direction},{turn_angle});")
return delta_degree
@staticmethod
def get_curr_gripper_to_base_pose() -> np.ndarray:
return ControlUtil.__fa.get_pose().matrix
@staticmethod
def get_curr_joints() -> np.ndarray:
return ControlUtil.__fa.get_joints()
@staticmethod
def set_gripper_pose(gripper_to_base: np.ndarray) -> None:
gripper_to_base = RigidTransform(rotation=gripper_to_base[:3, :3], translation=gripper_to_base[:3, 3], from_frame="franka_tool", to_frame="world")
ControlUtil.__fa.goto_pose(gripper_to_base, duration=6, use_impedance=False, ignore_errors=False, ignore_virtual_walls=True)
@staticmethod
def solve_gripper_to_base(cam_to_world: np.ndarray) -> np.ndarray:
return np.linalg.inv(ControlUtil.BASE_TO_WORLD) @ cam_to_world @ np.linalg.inv(ControlUtil.CAMERA_TO_GRIPPER)
@staticmethod
def sovle_cam_to_world(gripper_to_base: np.ndarray) -> np.ndarray:
return ControlUtil.BASE_TO_WORLD @ gripper_to_base @ ControlUtil.CAMERA_TO_GRIPPER
@staticmethod
def check_limit(new_cam_to_world):
if new_cam_to_world[0,3] > 0 or new_cam_to_world[1,3] > 0:
# if new_cam_to_world[0,3] > 0:
return False
x = abs(new_cam_to_world[0,3])
y = abs(new_cam_to_world[1,3])
tan_y_x = y/x
min_angle = 20 / 180 * np.pi
max_angle = 50 / 180 * np.pi
if tan_y_x < np.tan(min_angle) or tan_y_x > np.tan(max_angle):
return False
return True
@staticmethod
def solve_display_table_rot_and_cam_to_world(cam_to_world: np.ndarray) -> tuple:
if ControlUtil.check_limit(cam_to_world):
return 0, cam_to_world
else:
min_display_table_rot = 180
min_new_cam_to_world = None
for display_table_rot in np.linspace(0.1,360, 1800):
new_world_to_world = ControlUtil.get_z_axis_rot_mat(display_table_rot)
new_cam_to_new_world = cam_to_world
new_cam_to_world = new_world_to_world @ new_cam_to_new_world
if ControlUtil.check_limit(new_cam_to_world):
if display_table_rot < min_display_table_rot:
min_display_table_rot, min_new_cam_to_world = display_table_rot, new_cam_to_world
if abs(display_table_rot - 360) < min_display_table_rot:
min_display_table_rot, min_new_cam_to_world = display_table_rot - 360, new_cam_to_world
if min_new_cam_to_world is None:
raise ValueError("No valid display table rotation found")
return min_display_table_rot, min_new_cam_to_world
@staticmethod
def get_z_axis_rot_mat(degree):
radian = np.radians(degree)
return np.array([
[np.cos(radian), -np.sin(radian), 0, 0],
[np.sin(radian), np.cos(radian), 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
@staticmethod
def get_gripper_to_base_axis_angle(gripper_to_base: np.ndarray) -> bool:
rot_mat = gripper_to_base[:3,:3]
gripper_z_axis = rot_mat[:,2]
base_x_axis = np.array([1,0,0])
angle = np.arccos(np.dot(gripper_z_axis, base_x_axis))
return angle
@staticmethod
def move_to(pose: np.ndarray):
rot_degree, cam_to_world = ControlUtil.solve_display_table_rot_and_cam_to_world(pose)
delta_degree = ControlUtil.rotate_display_table(rot_degree)
exec_time = abs(delta_degree)/9
start_time = time.time()
ControlUtil.set_pose(cam_to_world)
end_time = time.time()
print(f"Move to pose with rotation {rot_degree} degree, exec time: {end_time - start_time}|exec time: {exec_time}")
if end_time - start_time < exec_time:
time.sleep(exec_time - (end_time - start_time))
@staticmethod
def absolute_rotate_display_table(degree):
exec_time = abs(degree)/9
write_len = ControlUtil.__ser.write(f"CT+TRUNSINGLE({0},{degree});".encode('utf-8'))
print(f"send command: CT+TRUNSINGLE({0},{degree});")
time.sleep(exec_time)
# ----------- Debug Test -------------
if __name__ == "__main__":
ControlUtil.connect_robot()
#ControlUtil.franka_reset()
print(ControlUtil.get_curr_gripper_to_base_pose())
# ControlUtil.rotate_display_table(180)
# ControlUtil.init()
# ControlUtil.franka_reset()
# def main_test():
# print(ControlUtil.get_curr_gripper_to_base_pose())
# ControlUtil.init()
# def rotate_back(rotation):
# ControlUtil.rotate_display_table(-rotation)
# #main_test()
# import sys; sys.path.append("/home/yan20/nbv_rec/project/franka_control")
# from utils.communicate_util import CommunicateUtil
# import ipdb
# ControlUtil.init()
# view_data_0 = CommunicateUtil.get_view_data(init=True)
# depth_extrinsics_0 = view_data_0["depth_extrinsics"]
# cam_to_world_0 = ControlUtil.get_pose()
# print("cam_extrinsics_0")
# print(depth_extrinsics_0)
# print("cam_to_world_0")
# print(cam_to_world_0)
# ipdb.set_trace()
# TEST_POSE:np.ndarray = np.asarray([
# [ 0.46532393, 0.62171798, 0.63002284, 0.30230963],
# [ 0.43205618, -0.78075723, 0.45136491, -0.29127173],
# [ 0.77251656, 0.06217437, -0.63193429, 0.559957 ],
# [ 0. , 0. , 0. , 1. ],
# ])
# TEST_POSE_CAM_TO_WORLD = ControlUtil.BASE_TO_WORLD @ TEST_POSE @ ControlUtil.CAMERA_TO_GRIPPER
# ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD)
# view_data_1 = CommunicateUtil.get_view_data()
# depth_extrinsics_1 = view_data_1["depth_extrinsics"]
# depth_extrinsics_1[:3,3] = depth_extrinsics_1[:3,3] / 1000
# cam_to_world_1 = ControlUtil.get_pose()
# print("cam_extrinsics_1")
# print(depth_extrinsics_1)
# print("cam_to_world_1")
# print(TEST_POSE_CAM_TO_WORLD)
# actual_cam_to_world_1 = cam_to_world_0 @ depth_extrinsics_1
# print("actual_cam_to_world_1")
# print(actual_cam_to_world_1)
# ipdb.set_trace()
# TEST_POSE_2:np.ndarray = np.asarray(
# [[ 0.74398544, -0.61922696, 0.251049, 0.47000935],
# [-0.47287207, -0.75338888, -0.45692666, 0.20843903],
# [ 0.47207883 , 0.22123272, -0.85334192, 0.57863381],
# [ 0. , 0. , 0. , 1. , ]]
# )
# TEST_POSE_CAM_TO_WORLD_2 = ControlUtil.BASE_TO_WORLD @ TEST_POSE_2 @ ControlUtil.CAMERA_TO_GRIPPER
# #ControlUtil.move_to(TEST_POSE_CAM_TO_WORLD_2)
# ControlUtil.set_pose(TEST_POSE_CAM_TO_WORLD_2)
# view_data_2 = CommunicateUtil.get_view_data()
# depth_extrinsics_2 = view_data_2["depth_extrinsics"]
# depth_extrinsics_2[:3,3] = depth_extrinsics_2[:3,3] / 1000
# cam_to_world_2 = ControlUtil.get_pose()
# print("cam_extrinsics_2")
# print(depth_extrinsics_2)
# print("cam_to_world_2")
# print(TEST_POSE_CAM_TO_WORLD_2)
# actual_cam_to_world_2 = cam_to_world_0 @ depth_extrinsics_2
# print("actual_cam_to_world_2")
# print(actual_cam_to_world_2)
# ipdb.set_trace()

384
utils/data_load.py Normal file
View File

@@ -0,0 +1,384 @@
import os
import numpy as np
import json
import cv2
import trimesh
import torch
import OpenEXR
import Imath
from utils.pts_util import PtsUtil
class DataLoadUtil:
TABLE_POSITION = np.asarray([0, 0, -0.1])
@staticmethod
def get_display_table_info(root, scene_name):
scene_info = DataLoadUtil.load_scene_info(root, scene_name)
display_table_info = scene_info["display_table"]
return display_table_info
@staticmethod
def get_display_table_top(root, scene_name):
# display_table_height = DataLoadUtil.get_display_table_info(root, scene_name)[
# "height"
# ]
display_table_height = 0.1 #FIXME
display_table_top = DataLoadUtil.TABLE_POSITION + np.asarray(
[0, 0, display_table_height]
)
return display_table_top
@staticmethod
def get_path(root, scene_name, frame_idx):
path = os.path.join(root, scene_name, f"{frame_idx}")
return path
@staticmethod
def get_label_num(root, scene_name):
label_dir = os.path.join(root, scene_name, "label")
return len(os.listdir(label_dir))
@staticmethod
def get_label_path(root, scene_name, seq_idx):
label_dir = os.path.join(root, scene_name, "label")
if not os.path.exists(label_dir):
os.makedirs(label_dir)
path = os.path.join(label_dir, f"{seq_idx}.json")
return path
@staticmethod
def get_label_path_old(root, scene_name):
path = os.path.join(root, scene_name, "label.json")
return path
@staticmethod
def get_scene_seq_length(root, scene_name):
camera_params_path = os.path.join(root, scene_name, "camera_params")
return len(os.listdir(camera_params_path))
@staticmethod
def load_mesh_at(model_dir, object_name, world_object_pose):
model_path = os.path.join(model_dir, object_name, "mesh.obj")
mesh = trimesh.load(model_path)
mesh.apply_transform(world_object_pose)
return mesh
@staticmethod
def get_bbox_diag(model_dir, object_name):
model_path = os.path.join(model_dir, object_name, "mesh.obj")
mesh = trimesh.load(model_path)
bbox = mesh.bounding_box.extents
diagonal_length = np.linalg.norm(bbox)
return diagonal_length
@staticmethod
def load_scene_info(root, scene_name):
scene_info_path = os.path.join(root, scene_name, "scene_info.json")
with open(scene_info_path, "r") as f:
scene_info = json.load(f)
return scene_info
@staticmethod
def load_target_pts_num_dict(root, scene_name):
target_pts_num_path = os.path.join(root, scene_name, "target_pts_num.json")
with open(target_pts_num_path, "r") as f:
target_pts_num_dict = json.load(f)
return target_pts_num_dict
@staticmethod
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):
depth = cv2.imread(real_path, cv2.IMREAD_UNCHANGED)
depth = depth.astype(np.float32) / 65535.0
min_depth = min_depth
max_depth = max_depth
depth_meters = min_depth + (max_depth - min_depth) * depth
return depth_meters
if binocular:
depth_path_L = os.path.join(
os.path.dirname(path), "depth", os.path.basename(path) + "_L.png"
)
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
else:
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)
return depth_meters
@staticmethod
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]
red = [255, 0, 0, 255]
threshold = 2
mask_image = np.where(
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
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_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))
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"
)
mask_image = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
return mask_image
@staticmethod
def load_exr_image(file_path):
exr_file = OpenEXR.InputFile(file_path)
header = exr_file.header()
dw = header['dataWindow']
width = dw.max.x - dw.min.x + 1
height = dw.max.y - dw.min.y + 1
float_channels = ['R', 'G', 'B']
img_data = []
for channel in float_channels:
channel_data = exr_file.channel(channel)
img_data.append(np.frombuffer(channel_data, dtype=np.float16).reshape((height, width)))
img = np.stack(img_data, axis=-1)
return img
@staticmethod
def load_normal(path, binocular=False, left_only=False, file_type="exr"):
if binocular and not left_only:
normal_path_L = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
)
normal_image_L = DataLoadUtil.load_exr_image(normal_path_L)
normal_path_R = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + f"_R.{file_type}"
)
normal_image_R = DataLoadUtil.load_exr_image(normal_path_R)
normalized_normal_image_L = normal_image_L * 2.0 - 1.0
normalized_normal_image_R = normal_image_R * 2.0 - 1.0
return normalized_normal_image_L, normalized_normal_image_R
else:
if binocular and left_only:
normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + f"_L.{file_type}"
)
else:
normal_path = os.path.join(
os.path.dirname(path), "normal", os.path.basename(path) + f".{file_type}"
)
normal_image = DataLoadUtil.load_exr_image(normal_path)
normalized_normal_image = normal_image * 2.0 - 1.0
return normalized_normal_image
@staticmethod
def load_label(path):
with open(path, "r") as f:
label_data = json.load(f)
return label_data
@staticmethod
def load_from_preprocessed_pts(path, file_type="npy"):
npy_path = os.path.join(
os.path.dirname(path), "pts", os.path.basename(path) + "." + file_type
)
if file_type == "txt":
pts = np.loadtxt(npy_path)
else:
pts = np.load(npy_path)
return pts
@staticmethod
def cam_pose_transformation(cam_pose_before):
offset = np.asarray([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
cam_pose_after = cam_pose_before @ offset
return cam_pose_after
@staticmethod
def load_cam_info(path, binocular=False, display_table_as_world_space_origin=True):
scene_dir = os.path.dirname(path)
root_dir = os.path.dirname(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"
)
with open(camera_params_path, "r") as f:
label_data = json.load(f)
cam_to_world = np.asarray(label_data["extrinsic"])
cam_to_world = DataLoadUtil.cam_pose_transformation(cam_to_world)
if display_table_as_world_space_origin:
world_to_display_table = np.eye(4)
world_to_display_table[:3, 3] = -DataLoadUtil.get_display_table_top(
root_dir, scene_name
)
cam_to_world = np.dot(world_to_display_table, cam_to_world)
cam_intrinsic = np.asarray(label_data["intrinsic"])
cam_info = {
"cam_to_world": cam_to_world,
"cam_intrinsic": cam_intrinsic,
"far_plane": label_data["far_plane"],
"near_plane": label_data["near_plane"],
}
if binocular:
cam_to_world_R = np.asarray(label_data["extrinsic_R"])
cam_to_world_R = DataLoadUtil.cam_pose_transformation(cam_to_world_R)
cam_to_world_O = np.asarray(label_data["extrinsic_cam_object"])
cam_to_world_O = DataLoadUtil.cam_pose_transformation(cam_to_world_O)
if display_table_as_world_space_origin:
cam_to_world_O = np.dot(world_to_display_table, cam_to_world_O)
cam_to_world_R = np.dot(world_to_display_table, cam_to_world_R)
cam_info["cam_to_world_O"] = cam_to_world_O
cam_info["cam_to_world_R"] = cam_to_world_R
return cam_info
@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
):
root_dir = os.path.dirname(scene_path)
scene_name = os.path.basename(scene_path)
if isinstance(cam_L, torch.Tensor):
cam_L = cam_L.cpu().numpy()
nO_to_display_table_pose = cam_L @ cam_O_to_cam_L
if display_table_as_world_space_origin:
display_table_to_world = np.eye(4)
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 = DataLoadUtil.cam_pose_transformation(nO_to_world_pose)
return nO_to_world_pose
@staticmethod
def get_target_point_cloud(
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")
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)
mask = mask.reshape(-1, 4)
target_mask = (mask == target_mask_label).all(axis=-1)
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_world = np.dot(cam_extrinsic, target_points_camera_aug.T).T[:, :3]
data = {
"points_world": target_points_world,
"points_camera": target_points_camera,
}
return data
@staticmethod
def get_point_cloud(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_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return {"points_world": points_world, "points_camera": points_camera}
@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),
display_table_mask_label=(0, 0, 255, 255),
get_display_table_pts=False,
require_normal=False,
):
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
if binocular:
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)
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_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 = PtsUtil.get_overlapping_points(
point_cloud_L, point_cloud_R, voxel_size
)
return overlap_points
else:
depth = DataLoadUtil.load_depth(
path, cam_info["near_plane"], cam_info["far_plane"]
)
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"]
return point_cloud
@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")
points_normals = np.loadtxt(points_path)
if display_table_as_world_space_origin:
points_normals[:, :3] = points_normals[
:, :3
] - DataLoadUtil.get_display_table_top(root, scene_name)
return points_normals

0
utils/inference_util.py Normal file
View File

164
utils/pose_util.py Normal file
View File

@@ -0,0 +1,164 @@
import numpy as np
class PoseUtil:
ROTATION = 1
TRANSLATION = 2
SCALE = 3
@staticmethod
def get_uniform_translation(trans_m_min, trans_m_max, trans_unit, debug=False):
if isinstance(trans_m_min, list):
x_min, y_min, z_min = trans_m_min
x_max, y_max, z_max = trans_m_max
else:
x_min, y_min, z_min = trans_m_min, trans_m_min, trans_m_min
x_max, y_max, z_max = trans_m_max, trans_m_max, trans_m_max
x = np.random.uniform(x_min, x_max)
y = np.random.uniform(y_min, y_max)
z = np.random.uniform(z_min, z_max)
translation = np.array([x, y, z])
if trans_unit == "cm":
translation = translation / 100
if debug:
print("uniform translation:", translation)
return translation
@staticmethod
def get_uniform_rotation(rot_degree_min=0, rot_degree_max=180, debug=False):
axis = np.random.randn(3)
axis /= np.linalg.norm(axis)
theta = np.random.uniform(
rot_degree_min / 180 * np.pi, rot_degree_max / 180 * np.pi
)
K = np.array(
[[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]]
)
R = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * (K @ K)
if debug:
print("uniform rotation:", theta * 180 / np.pi)
return R
@staticmethod
def get_uniform_pose(
trans_min, trans_max, rot_min=0, rot_max=180, trans_unit="cm", debug=False
):
translation = PoseUtil.get_uniform_translation(
trans_min, trans_max, trans_unit, debug
)
rotation = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
pose = np.eye(4)
pose[:3, :3] = rotation
pose[:3, 3] = translation
return pose
@staticmethod
def get_n_uniform_pose(
trans_min,
trans_max,
rot_min=0,
rot_max=180,
n=1,
trans_unit="cm",
fix=None,
contain_canonical=True,
debug=False,
):
if fix == PoseUtil.ROTATION:
translations = np.zeros((n, 3))
for i in range(n):
translations[i] = PoseUtil.get_uniform_translation(
trans_min, trans_max, trans_unit, debug
)
if contain_canonical:
translations[0] = np.zeros(3)
rotations = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
elif fix == PoseUtil.TRANSLATION:
rotations = np.zeros((n, 3, 3))
for i in range(n):
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
if contain_canonical:
rotations[0] = np.eye(3)
translations = PoseUtil.get_uniform_translation(
trans_min, trans_max, trans_unit, debug
)
else:
translations = np.zeros((n, 3))
rotations = np.zeros((n, 3, 3))
for i in range(n):
translations[i] = PoseUtil.get_uniform_translation(
trans_min, trans_max, trans_unit, debug
)
for i in range(n):
rotations[i] = PoseUtil.get_uniform_rotation(rot_min, rot_max, debug)
if contain_canonical:
translations[0] = np.zeros(3)
rotations[0] = np.eye(3)
pose = np.eye(4, 4, k=0)[np.newaxis, :].repeat(n, axis=0)
pose[:, :3, :3] = rotations
pose[:, :3, 3] = translations
return pose
@staticmethod
def get_n_uniform_pose_batch(
trans_min,
trans_max,
rot_min=0,
rot_max=180,
n=1,
batch_size=1,
trans_unit="cm",
fix=None,
contain_canonical=False,
debug=False,
):
batch_poses = []
for i in range(batch_size):
pose = PoseUtil.get_n_uniform_pose(
trans_min,
trans_max,
rot_min,
rot_max,
n,
trans_unit,
fix,
contain_canonical,
debug,
)
batch_poses.append(pose)
pose_batch = np.stack(batch_poses, axis=0)
return pose_batch
@staticmethod
def get_uniform_scale(scale_min, scale_max, debug=False):
if isinstance(scale_min, list):
x_min, y_min, z_min = scale_min
x_max, y_max, z_max = scale_max
else:
x_min, y_min, z_min = scale_min, scale_min, scale_min
x_max, y_max, z_max = scale_max, scale_max, scale_max
x = np.random.uniform(x_min, x_max)
y = np.random.uniform(y_min, y_max)
z = np.random.uniform(z_min, z_max)
scale = np.array([x, y, z])
if debug:
print("uniform scale:", scale)
return scale
@staticmethod
def rotation_6d_to_matrix_numpy(d6):
a1, a2 = d6[:3], d6[3:]
b1 = a1 / np.linalg.norm(a1)
b2 = a2 - np.dot(b1, a2) * b1
b2 = b2 / np.linalg.norm(b2)
b3 = np.cross(b1, b2)
return np.stack((b1, b2, b3), axis=-2)
@staticmethod
def matrix_to_rotation_6d_numpy(matrix):
return np.copy(matrix[:2, :]).reshape((6,))

247
utils/preprocess_util.py Normal file
View File

@@ -0,0 +1,247 @@
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"):
#import ipdb;ipdb.set_trace()
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_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_world
def get_world_points_and_normals(depth, normal, 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)
normals_camera = normal[mask].reshape(-1, 3)
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N, require_idx=True
)
if len(sampled_target_points) == 0:
return np.zeros((0, 3)), np.zeros((0, 3))
sampled_target_normals = normals_camera[idx]
points_camera_aug = np.concatenate((sampled_target_points, np.ones((sampled_target_points.shape[0], 1))), axis=-1)
points_world = np.dot(cam_extrinsic, points_camera_aug.T).T[:, :3]
return points_world, sampled_target_normals
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="npy"):
''' 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 = 90
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"] #FIXME
radius = 0.25
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)
normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=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, sampled_target_normals_L = get_world_points_and_normals(depth_L, normal_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, overlap_idx = PtsUtil.get_overlapping_points(
sampled_target_points_L, sampled_target_points_R, voxel_size, require_idx=True
)
target_normals = sampled_target_normals_L[overlap_idx]
if has_points:
has_points = target_points.shape[0] > 0
if has_points:
target_points, target_normals = PtsUtil.filter_points(
target_points, target_normals, cam_info["cam_to_world"], theta_limit = 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))
target_normals = 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"/home/yan20/nbv_rec/project/franka_control/temp_debug"
# 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, file_type="npy")
cnt+=1
end = time.time()
print(f"Time cost: {end-start}")

308
utils/pts_util.py Normal file
View File

@@ -0,0 +1,308 @@
import numpy as np
import open3d as o3d
import torch
import trimesh
from scipy.spatial import cKDTree
from utils.pose_util import PoseUtil
class PtsUtil:
@staticmethod
def voxel_downsample_point_cloud_old(point_cloud, voxel_size=0.005):
o3d_pc = o3d.geometry.PointCloud()
o3d_pc.points = o3d.utility.Vector3dVector(point_cloud)
downsampled_pc = o3d_pc.voxel_down_sample(voxel_size)
return np.asarray(downsampled_pc.points)
@staticmethod
def voxel_downsample_point_cloud(point_cloud, voxel_size=0.005, require_idx=False):
voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
if require_idx:
_, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
idx_sort = np.argsort(inverse)
idx_unique = idx_sort[np.cumsum(counts)-counts]
downsampled_points = point_cloud[idx_unique]
return downsampled_points, idx_unique
else:
unique_voxels = np.unique(voxel_indices, axis=0, return_inverse=True)
return unique_voxels[0]*voxel_size
@staticmethod
def random_downsample_point_cloud(point_cloud, num_points, require_idx=False):
if point_cloud.shape[0] == 0:
if require_idx:
return point_cloud, np.array([])
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 fps_downsample_point_cloud(point_cloud, num_points, require_idx=False):
N = point_cloud.shape[0]
mask = np.zeros(N, dtype=bool)
sampled_indices = np.zeros(num_points, dtype=int)
sampled_indices[0] = np.random.randint(0, N)
distances = np.linalg.norm(
point_cloud - point_cloud[sampled_indices[0]], axis=1
)
for i in range(1, num_points):
farthest_index = np.argmax(distances)
sampled_indices[i] = farthest_index
mask[farthest_index] = True
new_distances = np.linalg.norm(
point_cloud - point_cloud[farthest_index], axis=1
)
distances = np.minimum(distances, new_distances)
sampled_points = point_cloud[sampled_indices]
if require_idx:
return sampled_points, sampled_indices
return sampled_points
@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 transform_point_cloud(points, pose_mat):
points_h = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)
points_h = np.dot(pose_mat, points_h.T).T
return points_h[:, :3]
@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_limit=45, z_range=(0.2, 0.45)):
""" filter with normal """
normals_normalized = normals / np.linalg.norm(normals, axis=1, keepdims=True)
cos_theta = np.dot(normals_normalized, np.array([0, 0, 1]))
theta = np.arccos(cos_theta) * 180 / np.pi
idx = theta < theta_limit
filtered_sampled_points = points[idx]
filtered_normals = normals[idx]
""" filter with z range """
points_cam = PtsUtil.transform_point_cloud(filtered_sampled_points, np.linalg.inv(cam_pose))
idx = (points_cam[:, 2] > z_range[0]) & (points_cam[:, 2] < z_range[1])
z_filtered_points = filtered_sampled_points[idx]
z_filtered_normals = filtered_normals[idx]
return z_filtered_points[:, :3], z_filtered_normals
@staticmethod
def multi_scale_icp(
source, target, voxel_size_range, init_transformation=None, steps=20
):
pipreg = o3d.pipelines.registration
if init_transformation is not None:
current_transformation = init_transformation
else:
current_transformation = np.identity(4)
cnt = 0
best_score = 1e10
best_reg = None
voxel_sizes = []
for i in range(steps):
voxel_sizes.append(
voxel_size_range[0]
+ i * (voxel_size_range[1] - voxel_size_range[0]) / steps
)
for voxel_size in voxel_sizes:
radius_normal = voxel_size * 2
source_downsampled = source.voxel_down_sample(voxel_size)
source_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
target_downsampled = target.voxel_down_sample(voxel_size)
target_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
reg_icp = pipreg.registration_icp(
source_downsampled,
target_downsampled,
voxel_size * 2,
current_transformation,
pipreg.TransformationEstimationPointToPlane(),
pipreg.ICPConvergenceCriteria(max_iteration=500),
)
cnt += 1
if reg_icp.fitness == 0:
score = 1e10
else:
score = reg_icp.inlier_rmse / reg_icp.fitness
if score < best_score:
best_score = score
best_reg = reg_icp
return best_reg, best_score
@staticmethod
def multi_scale_ransac(source_downsampled, target_downsampled, source_fpfh, model_fpfh, voxel_size_range, steps=20):
pipreg = o3d.pipelines.registration
cnt = 0
best_score = 1e10
best_reg = None
voxel_sizes = []
for i in range(steps):
voxel_sizes.append(
voxel_size_range[0]
+ i * (voxel_size_range[1] - voxel_size_range[0]) / steps
)
for voxel_size in voxel_sizes:
reg_ransac = pipreg.registration_ransac_based_on_feature_matching(
source_downsampled,
target_downsampled,
source_fpfh,
model_fpfh,
mutual_filter=True,
max_correspondence_distance=voxel_size*2,
estimation_method=pipreg.TransformationEstimationPointToPoint(False),
ransac_n=4,
checkers=[pipreg.CorrespondenceCheckerBasedOnEdgeLength(0.9)],
criteria=pipreg.RANSACConvergenceCriteria(8000000, 500),
)
cnt += 1
if reg_ransac.fitness == 0:
score = 1e10
else:
score = reg_ransac.inlier_rmse / reg_ransac.fitness
if score < best_score:
best_score = score
best_reg = reg_ransac
return best_reg, best_score
@staticmethod
def register(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.01):
radius_normal = voxel_size * 2
pipreg = o3d.pipelines.registration
model_pcd = o3d.geometry.PointCloud()
model_pcd.points = o3d.utility.Vector3dVector(model.vertices)
model_downsampled = model_pcd.voxel_down_sample(voxel_size)
model_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
model_fpfh = pipreg.compute_fpfh_feature(
model_downsampled,
o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
)
source_pcd = o3d.geometry.PointCloud()
source_pcd.points = o3d.utility.Vector3dVector(pcl)
source_downsampled = source_pcd.voxel_down_sample(voxel_size)
source_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
source_fpfh = pipreg.compute_fpfh_feature(
source_downsampled,
o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=100),
)
reg_ransac, ransac_best_score = PtsUtil.multi_scale_ransac(
source_downsampled,
model_downsampled,
source_fpfh,
model_fpfh,
voxel_size_range=(0.03, 0.005),
steps=3,
)
reg_icp, icp_best_score = PtsUtil.multi_scale_icp(
source_downsampled,
model_downsampled,
voxel_size_range=(0.02, 0.001),
init_transformation=reg_ransac.transformation,
steps=50,
)
return reg_icp.transformation
@staticmethod
def register_fine(pcl: np.ndarray, model: trimesh.Trimesh, voxel_size=0.01):
radius_normal = voxel_size * 2
pipreg = o3d.pipelines.registration
model_pcd = o3d.geometry.PointCloud()
model_pcd.points = o3d.utility.Vector3dVector(model.vertices)
model_downsampled = model_pcd.voxel_down_sample(voxel_size)
model_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
source_pcd = o3d.geometry.PointCloud()
source_pcd.points = o3d.utility.Vector3dVector(pcl)
source_downsampled = source_pcd.voxel_down_sample(voxel_size)
source_downsampled.estimate_normals(
search_param=o3d.geometry.KDTreeSearchParamHybrid(
radius=radius_normal, max_nn=30
)
)
reg_icp, icp_best_score = PtsUtil.multi_scale_icp(
source_downsampled,
model_downsampled,
voxel_size_range=(0.02, 0.001),
init_transformation=np.eye(4),
steps=50,
)
return reg_icp.transformation
@staticmethod
def get_pts_from_depth(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)
mask = mask.reshape(-1, 4)
points_camera = np.concatenate(
[points_camera, np.ones((points_camera.shape[0], 1))], axis=-1
)
points_world = np.dot(cam_extrinsic, points_camera.T).T[:, :3]
data = {
"points_world": points_world,
"points_camera": points_camera,
}
return data

205
utils/reconstruction_util.py Normal file
View File

@@ -0,0 +1,205 @@
import numpy as np
from scipy.spatial import cKDTree
from utils.pts_util import PtsUtil
class ReconstructionUtil:
@staticmethod
def compute_coverage_rate(target_point_cloud, combined_point_cloud, threshold=0.01):
kdtree = cKDTree(combined_point_cloud)
distances, _ = kdtree.query(target_point_cloud)
covered_points_num = np.sum(distances < threshold*2)
coverage_rate = covered_points_num / target_point_cloud.shape[0]
return coverage_rate, covered_points_num
@staticmethod
def check_overlap(new_point_cloud, combined_point_cloud, overlap_area_threshold=25, voxel_size=0.01):
kdtree = cKDTree(combined_point_cloud)
distances, _ = kdtree.query(new_point_cloud)
overlapping_points = np.sum(distances < voxel_size*2)
cm = 0.01
voxel_size_cm = voxel_size / cm
overlap_area = overlapping_points * voxel_size_cm * voxel_size_cm
return overlap_area > overlap_area_threshold
@staticmethod
def get_new_added_points(old_combined_pts, new_pts, threshold=0.005):
if old_combined_pts.size == 0:
return new_pts
if new_pts.size == 0:
return np.array([])
tree = cKDTree(old_combined_pts)
distances, _ = tree.query(new_pts, k=1)
new_added_points = new_pts[distances > threshold]
return new_added_points
@staticmethod
def compute_next_best_view_sequence_with_overlap(target_point_cloud, point_cloud_list, scan_points_indices_list, threshold=0.01, soft_overlap_threshold=0.5, hard_overlap_threshold=0.7, init_view = 0, scan_points_threshold=5, status_info=None):
selected_views = [init_view]
combined_point_cloud = point_cloud_list[init_view]
history_indices = [scan_points_indices_list[init_view]]
max_rec_pts = np.vstack(point_cloud_list)
downsampled_max_rec_pts = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold)
max_rec_pts_num = downsampled_max_rec_pts.shape[0]
max_real_rec_pts_coverage, _ = ReconstructionUtil.compute_coverage_rate(target_point_cloud, downsampled_max_rec_pts, threshold)
new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, combined_point_cloud, threshold)
current_coverage = new_coverage
current_covered_num = new_covered_num
remaining_views = list(range(len(point_cloud_list)))
view_sequence = [(init_view, current_coverage)]
cnt_processed_view = 0
remaining_views.remove(init_view)
curr_rec_pts_num = combined_point_cloud.shape[0]
while remaining_views:
best_view = None
best_coverage_increase = -1
best_combined_point_cloud = None
best_covered_num = 0
for view_index in remaining_views:
if point_cloud_list[view_index].shape[0] == 0:
continue
if selected_views:
new_scan_points_indices = scan_points_indices_list[view_index]
if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold):
overlap_threshold = hard_overlap_threshold
else:
overlap_threshold = soft_overlap_threshold
overlap_rate = ReconstructionUtil.compute_overlap_rate(point_cloud_list[view_index],combined_point_cloud, threshold)
if overlap_rate < overlap_threshold:
continue
new_combined_point_cloud = np.vstack([combined_point_cloud, point_cloud_list[view_index]])
new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold)
new_coverage, new_covered_num = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold)
coverage_increase = new_coverage - current_coverage
if coverage_increase > best_coverage_increase:
best_coverage_increase = coverage_increase
best_view = view_index
best_covered_num = new_covered_num
best_combined_point_cloud = new_downsampled_combined_point_cloud
if best_view is not None:
if best_coverage_increase <=1e-3 or best_covered_num - current_covered_num <= 5:
break
selected_views.append(best_view)
best_rec_pts_num = best_combined_point_cloud.shape[0]
print(f"Current rec pts num: {curr_rec_pts_num}, Best rec pts num: {best_rec_pts_num}, Best cover pts: {best_covered_num}, Max rec pts num: {max_rec_pts_num}")
print(f"Current coverage: {current_coverage}, Best coverage increase: {best_coverage_increase}, Max Real coverage: {max_real_rec_pts_coverage}")
current_covered_num = best_covered_num
curr_rec_pts_num = best_rec_pts_num
combined_point_cloud = best_combined_point_cloud
remaining_views.remove(best_view)
history_indices.append(scan_points_indices_list[best_view])
current_coverage += best_coverage_increase
cnt_processed_view += 1
if status_info is not None:
sm = status_info["status_manager"]
app_name = status_info["app_name"]
runner_name = status_info["runner_name"]
sm.set_status(app_name, runner_name, "current coverage", current_coverage)
sm.set_progress(app_name, runner_name, "processed view", cnt_processed_view, len(point_cloud_list))
view_sequence.append((best_view, current_coverage))
else:
break
if status_info is not None:
sm = status_info["status_manager"]
app_name = status_info["app_name"]
runner_name = status_info["runner_name"]
sm.set_progress(app_name, runner_name, "processed view", len(point_cloud_list), len(point_cloud_list))
return view_sequence, remaining_views, combined_point_cloud
@staticmethod
def compute_next_best_view_with_overlap(scanned_pts, point_cloud_list, selected_view, history_indices, scan_points_indices_list, threshold=0.01, overlap_area_threshold=25, scan_points_threshold=5):
max_rec_pts = np.vstack(point_cloud_list)
downsampled_max_rec_pts = PtsUtil.voxel_downsample_point_cloud(max_rec_pts, threshold)
best_view = None
new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(scanned_pts,threshold)
best_coverage,_ = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold)
best_covered_num = new_downsampled_combined_point_cloud.shape[0]
#import ipdb; ipdb.set_trace()
for view in range(len(point_cloud_list)):
print(f"Processing view {view}/{len(point_cloud_list)}: {best_coverage},{best_covered_num}")
if point_cloud_list[view].shape[0] == 0:
continue
if view in selected_view:
continue
new_scan_points_indices = scan_points_indices_list[view]
# if not ReconstructionUtil.check_scan_points_overlap(history_indices, new_scan_points_indices, scan_points_threshold):
# curr_overlap_area_threshold = overlap_area_threshold
# else:
# curr_overlap_area_threshold = overlap_area_threshold * 0.8
# if not ReconstructionUtil.check_overlap(point_cloud_list[view], scanned_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=threshold):
# continue
new_combined_point_cloud = np.vstack([scanned_pts ,point_cloud_list[view]])
new_downsampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(new_combined_point_cloud,threshold)
new_coverage, _ = ReconstructionUtil.compute_coverage_rate(downsampled_max_rec_pts, new_downsampled_combined_point_cloud, threshold)
new_covered_num = new_downsampled_combined_point_cloud.shape[0]
print("the new coverage_rate: ", new_coverage, " the new covered num: ", new_covered_num, " the fake covered num:" , _)
if new_coverage > best_coverage :
#import ipdb; ipdb.set_trace()
best_coverage = new_coverage
best_covered_num = new_covered_num
best_view = view
return best_view, best_coverage, best_covered_num
@staticmethod
def generate_scan_points(display_table_top, display_table_radius, min_distance=0.03, max_points_num = 500, max_attempts = 1000):
points = []
attempts = 0
while len(points) < max_points_num and attempts < max_attempts:
angle = np.random.uniform(0, 2 * np.pi)
r = np.random.uniform(0, display_table_radius)
x = r * np.cos(angle)
y = r * np.sin(angle)
z = display_table_top
new_point = (x, y, z)
if all(np.linalg.norm(np.array(new_point) - np.array(existing_point)) >= min_distance for existing_point in points):
points.append(new_point)
attempts += 1
return points
@staticmethod
def compute_covered_scan_points(scan_points, point_cloud, threshold=0.01):
tree = cKDTree(point_cloud)
covered_points = []
indices = []
for i, scan_point in enumerate(scan_points):
if tree.query_ball_point(scan_point, threshold):
covered_points.append(scan_point)
indices.append(i)
return covered_points, indices
@staticmethod
def check_scan_points_overlap(history_indices, indices2, threshold=5):
try:
if len(indices2) == 0:
return False
for indices1 in history_indices:
if len(set(indices1).intersection(set(indices2))) >= threshold:
return True
except Exception as e:
print(e)
import ipdb; ipdb.set_trace()
return False

45
utils/render_util.py Normal file
View File

@@ -0,0 +1,45 @@
import os
import json
import subprocess
import tempfile
import shutil
from utils.data_load import DataLoadUtil
from utils.pts_util import PtsUtil
class RenderUtil:
@staticmethod
def render_pts(cam_pose, object_name, script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False):
nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path)
with tempfile.TemporaryDirectory() as temp_dir:
params = {
"cam_pose": nO_to_world_pose.tolist(),
"object_name": scene_path
}
params_data_path = os.path.join(temp_dir, "params.json")
with open(params_data_path, 'w') as f:
json.dump(params, f)
result = subprocess.run([
'blender', '-b', '-P', script_path, '--', temp_dir
], capture_output=True, text=True)
if result.returncode != 0:
print("Blender script failed:")
print(result.stderr)
return None
path = os.path.join(temp_dir, "tmp")
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 = 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)
point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_params['cam_intrinsic'], cam_params['cam_to_world'])['points_world']
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_params['cam_intrinsic'], cam_params['cam_to_world_R'])['points_world']
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 = PtsUtil.get_overlapping_points(point_cloud_L, point_cloud_R)
return filtered_point_cloud, full_scene_point_cloud

132
utils/view_util.py Normal file
View File

@@ -0,0 +1,132 @@
import numpy as np
from scipy.spatial.transform import Rotation as R
from dataclasses import dataclass
@dataclass
class CameraIntrinsics:
width: int
height: int
fx: float
fy: float
cx: float
cy: float
@property
def intrinsic_matrix(self):
return np.array([[self.fx, 0, self.cx], [0, self.fy, self.cy], [0, 0, 1]])
@dataclass
class CameraExtrinsics:
def __init__(self, rotation: np.ndarray, translation: np.ndarray, rot_type: str):
"""
rotation: 3x3 rotation matrix or 1x3 euler angles or 1x4 quaternion
translation: 1x3 or 3x1 translation vector
rot_type: "mat", "euler_xyz", "quat_xyzw"
"""
assert rot_type in ["mat", "euler_xyz", "quat_xyzw"]
if rot_type == "mat":
self._rot = R.from_matrix(rotation)
elif rot_type == "euler_xyz":
self._rot = R.from_euler('xyz', rotation, degrees=True)
elif rot_type == "quat_xyzw":
self._rot = R.from_quat(rotation)
self._translation = translation
@property
def extrinsic_matrix(self):
return np.vstack([np.hstack([self._rot.as_matrix(), self._translation.reshape(3, 1)]), [0, 0, 0, 1]])
@property
def rotation_euler_xyz(self):
return self._rot.as_euler('xyz', degrees=True)
@property
def rotation_quat_xyzw(self):
return self._rot.as_quat()
@property
def rotation_matrix(self):
return self._rot.as_matrix()
@property
def translation(self):
return self._translation
@dataclass
class CameraData:
def __init__(self, depth_image: np.ndarray, image_id: int, intrinsics: CameraIntrinsics, extrinsics: CameraExtrinsics):
self._depth_image = depth_image
self._image_id = image_id
self._intrinsics = intrinsics
self._extrinsics = extrinsics
@property
def depth_image(self):
return self._depth_image
@property
def image_id(self):
return self._image_id
@property
def intrinsics(self):
return self._intrinsics.intrinsic_matrix
@property
def extrinsics(self):
return self._extrinsics.extrinsic_matrix
@property
def projection_matrix(self):
return self.intrinsics @ self.extrinsics[:3, :4]
@property
def pts_camera(self):
height, width = self.depth_image.shape
v, u = np.indices((height, width))
points = np.vstack([u.flatten(), v.flatten(), np.ones_like(u.flatten())]) # 3xN
points = np.linalg.inv(self.intrinsics) @ points # 3xN
points = points.T # Nx3
points = points * self.depth_image.flatten()[:, np.newaxis] # Nx3
points = points[points[:, 2] > 0] # Nx3
return points
@property
def pts_world(self):
homogeneous_pts = np.hstack([self.pts_camera, np.ones((self.pts_camera.shape[0], 1))]) # Nx4
pts_world = self.extrinsics @ homogeneous_pts.T # 4xN
return pts_world[:3, :].T
class ViewUtil:
def get_pts(view_data):
image_id = view_data["image_id"]
depth_intrinsics = view_data["depth_intrinsics"]
depth_extrinsics = view_data["depth_extrinsics"]
depth_image = np.array(view_data["depth_image"], dtype=np.uint16)
if image_id is None:
return None
else:
camera_intrinsics = CameraIntrinsics(
depth_intrinsics['width'],
depth_intrinsics['height'],
depth_intrinsics['fx'],
depth_intrinsics['fy'],
depth_intrinsics['cx'],
depth_intrinsics['cy']
)
camera_extrinsics = CameraExtrinsics(
depth_extrinsics[:3, :3],
depth_extrinsics[:3, 3],
rot_type="mat"
)
camera_data = CameraData(depth_image, image_id, camera_intrinsics, camera_extrinsics)
pts = camera_data.pts_world
return pts/1000
def get_camera_pose(view_data, translation_scale=0.001):
depth_extrinsics = view_data["depth_extrinsics"]
depth_extrinsics[:3, 3] *= translation_scale
return depth_extrinsics

172
utils/vis.py Normal file
View File

@@ -0,0 +1,172 @@
import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import trimesh
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.data_load import DataLoadUtil
from utils.pts_util import PtsUtil
class visualizeUtil:
@staticmethod
def save_all_cam_pos_and_cam_axis(root, scene, output_dir):
length = DataLoadUtil.get_scene_seq_length(root, scene)
all_cam_pos = []
all_cam_axis = []
for i in range(length):
path = DataLoadUtil.get_path(root, scene, i)
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
cam_pose = cam_info["cam_to_world"]
cam_pos = cam_pose[:3, 3]
cam_axis = cam_pose[:3, 2]
num_samples = 10
sample_points = [cam_pos + 0.02*t * cam_axis for t in range(num_samples)]
sample_points = np.array(sample_points)
all_cam_pos.append(cam_pos)
all_cam_axis.append(sample_points)
all_cam_pos = np.array(all_cam_pos)
all_cam_axis = np.array(all_cam_axis).reshape(-1, 3)
np.savetxt(os.path.join(output_dir, "all_cam_pos.txt"), all_cam_pos)
np.savetxt(os.path.join(output_dir, "all_cam_axis.txt"), all_cam_axis)
@staticmethod
def save_all_combined_pts(root, scene, output_dir):
length = DataLoadUtil.get_scene_seq_length(root, scene)
all_combined_pts = []
for i in range(length):
path = DataLoadUtil.get_path(root, scene, i)
pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy")
if pts.shape[0] == 0:
continue
all_combined_pts.append(pts)
all_combined_pts = np.vstack(all_combined_pts)
downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.001)
np.savetxt(os.path.join(output_dir, "all_combined_pts.txt"), downsampled_all_pts)
@staticmethod
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)
target_name = scene_info["target_name"]
transformation = scene_info[target_name]
if display_table_as_world_space_origin:
location = transformation["location"] - DataLoadUtil.get_display_table_top(
root, scene_name
)
else:
location = transformation["location"]
rotation_euler = transformation["rotation_euler"]
pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
pose_mat[:3, 3] = location
mesh = DataLoadUtil.load_mesh_at(model_dir, target_name, pose_mat)
mesh_dir = os.path.join(root, scene_name, "mesh")
if not os.path.exists(mesh_dir):
os.makedirs(mesh_dir)
model_path = os.path.join(mesh_dir, "world_target_mesh.obj")
mesh.export(model_path)
@staticmethod
def save_points_and_normals(root, scene, frame_idx, output_dir, binocular=False):
target_mask_label = (0, 255, 0, 255)
path = DataLoadUtil.get_path(root, scene, frame_idx)
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular, display_table_as_world_space_origin=False)
depth = DataLoadUtil.load_depth(
path, cam_info["near_plane"],
cam_info["far_plane"],
binocular=binocular,
)
if isinstance(depth, tuple):
depth = depth[0]
mask = DataLoadUtil.load_seg(path, binocular=binocular, left_only=True)
normal = DataLoadUtil.load_normal(path, binocular=binocular, left_only=True)
''' target points '''
if mask is None:
target_mask_img = np.ones_like(depth, dtype=bool)
else:
target_mask_img = (mask == target_mask_label).all(axis=-1)
cam_intrinsic = cam_info["cam_intrinsic"]
z = depth[target_mask_img]
i, j = np.nonzero(target_mask_img)
x = (j - cam_intrinsic[0, 2]) * z / cam_intrinsic[0, 0]
y = (i - cam_intrinsic[1, 2]) * z / cam_intrinsic[1, 1]
random_downsample_N = 1000
points_camera = np.stack((x, y, z), axis=-1).reshape(-1, 3)
normal_camera = normal[target_mask_img].reshape(-1, 3)
sampled_target_points, idx = PtsUtil.random_downsample_point_cloud(
points_camera, random_downsample_N, require_idx=True
)
if len(sampled_target_points) == 0:
print("No target points")
sampled_normal_camera = normal_camera[idx]
sampled_visualized_normal = []
sampled_normal_camera[:, 2] = -sampled_normal_camera[:, 2]
sampled_normal_camera[:, 1] = -sampled_normal_camera[:, 1]
num_samples = 10
for i in range(len(sampled_target_points)):
sampled_visualized_normal.append([sampled_target_points[i] + 0.02*t * sampled_normal_camera[i] for t in range(num_samples)])
sampled_visualized_normal = np.array(sampled_visualized_normal).reshape(-1, 3)
np.savetxt(os.path.join(output_dir, "target_pts.txt"), sampled_target_points)
np.savetxt(os.path.join(output_dir, "target_normal.txt"), sampled_visualized_normal)
@staticmethod
def visualize_pts_and_camera(pts, camera):
import plotly.graph_objects as go
fig = go.Figure()
if pts is not None:
fig.add_trace(go.Scatter3d(
x=pts[:, 0], y=pts[:, 1], z=pts[:, 2],
mode='markers', marker=dict(size=1, color='gray', opacity=0.5), name='Input Points'
))
colors = ['aggrnyl', 'agsunset', 'algae', 'amp', 'armyrose', 'balance',
'blackbody', 'bluered', 'blues', 'blugrn', 'bluyl', 'brbg']
color = colors[0]
origin_candidate = camera[:3, 3]
z_axis_candidate = camera[:3, 2]
fig.add_trace(go.Cone(
x=[origin_candidate[0]], y=[origin_candidate[1]], z=[origin_candidate[2]],
u=[z_axis_candidate[0]], v=[z_axis_candidate[1]], w=[z_axis_candidate[2]],
colorscale=color,
sizemode="absolute", sizeref=0.1, anchor="tail", showscale=False
))
fig.update_layout(
title="Clustered Poses and Input Points",
scene=dict(
xaxis_title='X',
yaxis_title='Y',
zaxis_title='Z'
),
margin=dict(l=0, r=0, b=0, t=40),
scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
)
fig.show()
# ------ Debug ------
if __name__ == "__main__":
root = r"/home/yan20/nbv_rec/project/franka_control/output"
model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes"
scene = "box_1"
output_dir = f"/home/yan20/nbv_rec/project/franka_control/output/{scene}/visualize"
if not os.path.exists(output_dir):
os.makedirs(output_dir)
visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir)
visualizeUtil.save_all_combined_pts(root, scene, output_dir)
#visualizeUtil.save_target_mesh_at_world_space(root, model_dir, scene)
#visualizeUtil.save_points_and_normals(root, scene,"10", output_dir, binocular=True)