nbv_reconstruction/utils/vis.py

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

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,"txt")
            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):
        target_mask_label = (0, 255, 0, 255)
        path = DataLoadUtil.get_path(root, scene, frame_idx)
        cam_info = DataLoadUtil.load_cam_info(path, binocular=True)
        depth_L,_ = DataLoadUtil.load_depth(
                path, cam_info["near_plane"], 
                cam_info["far_plane"], 
                binocular=True,
            )
        mask_L = DataLoadUtil.load_seg(path, binocular=True, left_only=True)
        normal_L = DataLoadUtil.load_normal(path, binocular=True, left_only=True)
        ''' target points '''
        target_mask_img_L = (mask_L == target_mask_label).all(axis=-1)
        cam_intrinsic = cam_info["cam_intrinsic"]
        z = depth_L[target_mask_img_L]
        i, j = np.nonzero(target_mask_img_L)
        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_L[target_mask_img_L].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")
            
        offset = np.asarray([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
        sampled_normal_camera  = normal_camera[idx]
        sampled_normal_camera = np.dot(sampled_normal_camera, offset)
        sampled_visualized_normal = []
        
        
        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)
        


# ------ 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 = "omniobject3d-box_030"
    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_target_mesh_at_world_space(root, model_dir, scene)
    visualizeUtil.save_points_and_normals(root, scene, 0, output_dir)