import torch import os import json import numpy as np import subprocess import tempfile from utils.data_load import DataLoadUtil from utils.reconstruction import ReconstructionUtil from utils.pose import PoseUtil from utils.pts import PtsUtil import PytorchBoot.stereotype as stereotype import PytorchBoot.namespace as namespace from PytorchBoot.utils.log_util import Log def render_pts(cam_pose, scene_path,script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None): nO_to_world_pose = cam_pose.cpu().numpy() @ nO_to_nL_pose nO_to_world_pose = DataLoadUtil.cam_pose_transformation(nO_to_world_pose) with tempfile.TemporaryDirectory() as temp_dir: params = { "cam_pose": nO_to_world_pose.tolist(), "scene_path": scene_path } params_data_path = os.path.join(temp_dir, "params.json") with open(params_data_path, 'w') as f: json.dump(params, f) result = subprocess.run([ 'blender', '-b', '-P', script_path, '--', temp_dir ], capture_output=True, text=True) if result.returncode != 0: print("Blender script failed:") print(result.stderr) return None path = os.path.join(temp_dir, "tmp") point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) cam_params = DataLoadUtil.load_cam_info(path, binocular=True) sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) return sampled_point_cloud @stereotype.evaluation_method("pose_diff") class PoseDiff: def __init__(self, _): pass def evaluate(self, output_list, data_list): results = {namespace.TensorBoard.SCALAR: {}} rot_angle_list = [] trans_dist_list = [] for output, data in zip(output_list, data_list): gt_pose_9d = data['best_to_1_pose_9d'] pred_pose_9d = output['pred_pose_9d'] gt_rot_6d = gt_pose_9d[:, :6] gt_trans = gt_pose_9d[:, 6:] pred_rot_6d = pred_pose_9d[:, :6] pred_trans = pred_pose_9d[:, 6:] gt_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(gt_rot_6d) pred_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_rot_6d) rotation_angles = PoseUtil.rotation_angle_distance(gt_rot_mat, pred_rot_mat) rot_angle_list.extend(list(rotation_angles)) trans_dist = torch.norm(gt_trans-pred_trans) trans_dist_list.append(trans_dist) results[namespace.TensorBoard.SCALAR]["rot_diff"] = float(sum(rot_angle_list) / len(rot_angle_list)) results[namespace.TensorBoard.SCALAR]["trans_diff"] = float(sum(trans_dist_list) / len(trans_dist_list)) return results @stereotype.evaluation_method("coverage_rate_increase") class ConverageRateIncrease: def __init__(self, config): self.config = config self.renderer_path = config["renderer_path"] def evaluate(self, output_list, data_list): results = {namespace.TensorBoard.SCALAR: {}} gt_coverate_increase_list = [] pred_coverate_increase_list = [] cr_diff_list = [] for output, data in zip(output_list, data_list): scanned_cr = data['scanned_coverage_rate'] gt_cr = data["best_coverage_rate"] scene_path_list = data['scene_path'] model_points_normals_list = data['model_points_normals'] scanned_view_pts_list = data['scanned_target_pts_list'] pred_pose_9ds = output['pred_pose_9d'] nO_to_nL_pose_batch = data["nO_to_nL_pose"] voxel_threshold_list = data["voxel_threshold"] filter_degree_list = data["filter_degree"] first_frame_to_world = data["first_frame_to_world"] pred_n_to_1_pose_mats = torch.eye(4, device=pred_pose_9ds.device).unsqueeze(0).repeat(pred_pose_9ds.shape[0], 1, 1) pred_n_to_1_pose_mats[:,:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6]) pred_n_to_1_pose_mats[:,:3,3] = pred_pose_9ds[:, 6:] pred_n_to_world_pose_mats = torch.matmul(first_frame_to_world, pred_n_to_1_pose_mats) for idx in range(len(scanned_cr)): model_points_normals = model_points_normals_list[idx] scanned_view_pts = scanned_view_pts_list[idx] voxel_threshold = voxel_threshold_list[idx] model_pts = model_points_normals[:,:3] down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold) old_scanned_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold) gt_coverate_increase_list.append(gt_cr[idx]-old_scanned_cr) scene_path = scene_path_list[idx] pred_pose = pred_n_to_world_pose_mats[idx] filter_degree = filter_degree_list[idx] nO_to_nL_pose = nO_to_nL_pose_batch[idx] try: new_pts = render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose) pred_cr = self.compute_coverage_rate(scanned_view_pts, new_pts, down_sampled_model_pts, threshold=voxel_threshold) except Exception as e: Log.warning(f"Error in scene {scene_path}, {e}") pred_cr = old_scanned_cr pred_coverate_increase_list.append(pred_cr-old_scanned_cr) cr_diff_list.append(gt_cr[idx]-pred_cr) results[namespace.TensorBoard.SCALAR]["gt_cr_increase"] = float(sum(gt_coverate_increase_list) / len(gt_coverate_increase_list)) results[namespace.TensorBoard.SCALAR]["pred_cr_increase"] = float(sum(pred_coverate_increase_list) / len(pred_coverate_increase_list)) results[namespace.TensorBoard.SCALAR]["cr_diff"] = float(sum(cr_diff_list) / len(cr_diff_list)) return results def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005): if new_pts is not None: new_scanned_view_pts = scanned_view_pts + [new_pts] else: new_scanned_view_pts = scanned_view_pts combined_point_cloud = np.vstack(new_scanned_view_pts) down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold) return ReconstructionUtil.compute_coverage_rate(model_pts, down_sampled_combined_point_cloud, threshold)