209 lines
9.4 KiB
Python
209 lines
9.4 KiB
Python
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 import PtsUtil
|
|
from utils.pose import PoseUtil
|
|
|
|
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 get_cam_pose_and_cam_axis(cam_pose, is_6d_pose):
|
|
if is_6d_pose:
|
|
matrix_cam_pose = np.eye(4)
|
|
matrix_cam_pose[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(cam_pose[:6])
|
|
matrix_cam_pose[:3, 3] = cam_pose[6:]
|
|
else:
|
|
matrix_cam_pose = cam_pose
|
|
cam_pos = matrix_cam_pose[:3, 3]
|
|
cam_axis = matrix_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)
|
|
return cam_pos, sample_points
|
|
|
|
@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_seq_cam_pos_and_cam_axis(root, scene, frame_idx_list, output_dir):
|
|
all_cam_pos = []
|
|
all_cam_axis = []
|
|
for i in frame_idx_list:
|
|
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, "seq_cam_pos.txt"), all_cam_pos)
|
|
np.savetxt(os.path.join(output_dir, "seq_cam_axis.txt"), all_cam_axis)
|
|
|
|
@staticmethod
|
|
def save_seq_combined_pts(root, scene, frame_idx_list, output_dir):
|
|
all_combined_pts = []
|
|
for i in frame_idx_list:
|
|
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, "seq_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 save_pts_nrm(root, scene, frame_idx, output_dir, binocular=False):
|
|
path = DataLoadUtil.get_path(root, scene, frame_idx)
|
|
pts_world = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
|
|
nrm_camera = DataLoadUtil.load_from_preprocessed_nrm(path, "npy")
|
|
cam_info = DataLoadUtil.load_cam_info(path, binocular=binocular)
|
|
cam_to_world = cam_info["cam_to_world"]
|
|
nrm_world = nrm_camera @ cam_to_world[:3, :3].T
|
|
visualized_nrm = []
|
|
num_samples = 10
|
|
for i in range(len(pts_world)):
|
|
for t in range(num_samples):
|
|
visualized_nrm.append(pts_world[i] - 0.02 * t * nrm_world[i])
|
|
|
|
visualized_nrm = np.array(visualized_nrm)
|
|
np.savetxt(os.path.join(output_dir, "nrm.txt"), visualized_nrm)
|
|
np.savetxt(os.path.join(output_dir, "pts.txt"), pts_world)
|
|
|
|
# ------ Debug ------
|
|
|
|
if __name__ == "__main__":
|
|
root = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\temp"
|
|
model_dir = r"H:\\AI\\Datasets\\scaled_object_box_meshes"
|
|
scene = "box"
|
|
output_dir = r"C:\Document\Local Project\nbv_rec\nbv_reconstruction\test"
|
|
|
|
#visualizeUtil.save_all_cam_pos_and_cam_axis(root, scene, output_dir)
|
|
# visualizeUtil.save_all_combined_pts(root, scene, output_dir)
|
|
# visualizeUtil.save_seq_combined_pts(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], output_dir)
|
|
# visualizeUtil.save_seq_cam_pos_and_cam_axis(root, scene, [0, 121, 286, 175, 111,366,45,230,232,225,255,17,199,78,60], 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)
|
|
visualizeUtil.save_pts_nrm(root, scene, "116", output_dir, binocular=True)
|