137 lines
6.8 KiB
Python
137 lines
6.8 KiB
Python
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)
|
|
|
|
|
|
|
|
|