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)