update backend
This commit is contained in:
98
data_load.py
98
data_load.py
@@ -3,23 +3,32 @@ import numpy as np
|
||||
import json
|
||||
import cv2
|
||||
import trimesh
|
||||
import torch
|
||||
from pts import PtsUtil
|
||||
|
||||
class DataLoadUtil:
|
||||
|
||||
DISPLAY_TABLE_POSITION = np.asarray([0,0,0.895])
|
||||
@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_path(root, scene_name):
|
||||
path = os.path.join(root,scene_name, f"label.json")
|
||||
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_sampled_model_points_path(root, scene_name):
|
||||
path = os.path.join(root,scene_name, f"sampled_model_points.txt")
|
||||
def get_label_path_old(root, scene_name):
|
||||
path = os.path.join(root,scene_name,"label.json")
|
||||
return path
|
||||
|
||||
@staticmethod
|
||||
@@ -27,17 +36,6 @@ class DataLoadUtil:
|
||||
camera_params_path = os.path.join(root, scene_name, "camera_params")
|
||||
return len(os.listdir(camera_params_path))
|
||||
|
||||
@staticmethod
|
||||
def load_downsampled_world_model_points(root, scene_name):
|
||||
model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name)
|
||||
model_points = np.loadtxt(model_path)
|
||||
return model_points
|
||||
|
||||
@staticmethod
|
||||
def save_downsampled_world_model_points(root, scene_name, model_points):
|
||||
model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name)
|
||||
np.savetxt(model_path, model_points)
|
||||
|
||||
@staticmethod
|
||||
def load_mesh_at(model_dir, object_name, world_object_pose):
|
||||
model_path = os.path.join(model_dir, object_name, "mesh.obj")
|
||||
@@ -45,6 +43,15 @@ class DataLoadUtil:
|
||||
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 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)
|
||||
@@ -52,11 +59,14 @@ class DataLoadUtil:
|
||||
mesh.export(model_path)
|
||||
|
||||
@staticmethod
|
||||
def save_target_mesh_at_world_space(root, model_dir, scene_name):
|
||||
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]
|
||||
location = transformation["location"]
|
||||
if display_table_as_world_space_origin:
|
||||
location = transformation["location"] - DataLoadUtil.DISPLAY_TABLE_POSITION
|
||||
else:
|
||||
location = transformation["location"]
|
||||
rotation_euler = transformation["rotation_euler"]
|
||||
pose_mat = trimesh.transformations.euler_matrix(*rotation_euler)
|
||||
pose_mat[:3, 3] = location
|
||||
@@ -140,6 +150,12 @@ class DataLoadUtil:
|
||||
rgb_image = cv2.imread(rgb_path, cv2.IMREAD_COLOR)
|
||||
return rgb_image
|
||||
|
||||
@staticmethod
|
||||
def load_from_preprocessed_pts(path):
|
||||
npy_path = os.path.join(os.path.dirname(path), "points", os.path.basename(path) + ".npy")
|
||||
pts = np.load(npy_path)
|
||||
return pts
|
||||
|
||||
@staticmethod
|
||||
def cam_pose_transformation(cam_pose_before):
|
||||
offset = np.asarray([
|
||||
@@ -151,12 +167,16 @@ class DataLoadUtil:
|
||||
return cam_pose_after
|
||||
|
||||
@staticmethod
|
||||
def load_cam_info(path, binocular=False):
|
||||
def load_cam_info(path, binocular=False, display_table_as_world_space_origin=True):
|
||||
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)
|
||||
world_to_display_table = np.eye(4)
|
||||
world_to_display_table[:3, 3] = - DataLoadUtil.DISPLAY_TABLE_POSITION
|
||||
if display_table_as_world_space_origin:
|
||||
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,
|
||||
@@ -167,9 +187,27 @@ class DataLoadUtil:
|
||||
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, display_table_as_world_space_origin=True):
|
||||
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.DISPLAY_TABLE_POSITION
|
||||
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)):
|
||||
h, w = depth.shape
|
||||
@@ -192,6 +230,24 @@ class DataLoadUtil:
|
||||
"points_world": target_points_world,
|
||||
"points_camera": target_points_camera
|
||||
}
|
||||
|
||||
@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)):
|
||||
@@ -232,7 +288,9 @@ class DataLoadUtil:
|
||||
return overlapping_points
|
||||
|
||||
@staticmethod
|
||||
def load_points_normals(root, scene_name):
|
||||
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.DISPLAY_TABLE_POSITION
|
||||
return points_normals
|
||||
Reference in New Issue
Block a user