first commit
This commit is contained in:
208
utils/vis.py
Normal file
208
utils/vis.py
Normal file
@@ -0,0 +1,208 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user