import os import sys # 设置项目根目录 project_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) sys.path.append(project_dir) os.chdir(project_dir) # 导入日志系统 from utils.logger import logger import numpy as np from scipy.spatial import cKDTree from scipy.spatial.distance import directed_hausdorff import trimesh import pandas as pd import csv import math import pickle import argparse # parse args first and set gpu id parser = argparse.ArgumentParser() parser.add_argument('--gt_path', type=str, default=os.path.join(project_dir, '../data/eval_data'), help='ground truth data path') parser.add_argument('--pred_path', type=str, default=os.path.join(project_dir, '../data/output_data'), help='converted data path') parser.add_argument('--name_list', type=str, default='broken_bullet_name.txt', help='names of models to be evaluated, if you want to evaluate the whole dataset, please set it as all_names.txt') parser.add_argument('--nsample', type=int, default=50000, help='point batch size') parser.add_argument('--regen', default = False, action="store_true", help = 'regenerate feature curves') args = parser.parse_args() def distance_p2p(points_src, normals_src, points_tgt, normals_tgt): ''' Computes minimal distances of each point in points_src to points_tgt. Args: points_src (numpy array [N, 3]): source points normals_src (numpy array [N, 3]): source normals points_tgt (numpy array [M, 3]): target points normals_tgt (numpy array [M, 3]): target Returns: dist (numpy array [N]): minimal distances of each point in points_src to points_tgt normals_dot_product (numpy array [N]): dot product of normals of points_src and points_tgt ''' kdtree = cKDTree(points_tgt) dist, idx = kdtree.query(points_src) if normals_src is not None and normals_tgt is not None: normals_src = \ normals_src / np.linalg.norm(normals_src, axis=-1, keepdims=True) normals_tgt = \ normals_tgt / np.linalg.norm(normals_tgt, axis=-1, keepdims=True) normals_dot_product = (normals_tgt[idx] * normals_src).sum(axis=-1) # Handle normals that point into wrong direction gracefully # (mostly due to mehtod not caring about this in generation) normals_dot_product = np.abs(normals_dot_product) return dist, normals_dot_product def distance_feature2mesh(points, mesh): prox = trimesh.proximity.ProximityQuery(mesh) signed_distance = prox.signed_distance(points) return np.abs(signed_distance) def distance_p2mesh(points_src, normals_src, mesh): points_tgt, idx = mesh.sample(args.nsample, return_index=True) points_tgt = points_tgt.astype(np.float32) normals_tgt = mesh.face_normals[idx] cd1, nc1 = distance_p2p(points_src, normals_src, points_tgt, normals_tgt) #pred2gt hd1 = cd1.max() cd1 = cd1.mean() nc1 = np.clip(nc1, -1.0, 1.0) angles1 = np.arccos(nc1) / math.pi * 180.0 angles1_mean = angles1.mean() angles1_std = np.std(angles1) cd2, nc2 = distance_p2p(points_tgt, normals_tgt, points_src, normals_src) #gt2pred hd2 = cd2.max() cd2 = cd2.mean() nc2 = np.clip(nc2, -1.0, 1.0) angles2 = np.arccos(nc2)/ math.pi * 180.0 angles2_mean = angles2.mean() angles2_std = np.std(angles2) cd = 0.5 * (cd1 + cd2) hd = max(hd1, hd2) angles_mean = 0.5 * (angles1_mean + angles2_mean) angles_std = 0.5 * (angles1_std + angles2_std) return cd, hd, angles_mean, angles_std, hd1, hd2 def distance_fea(gt_pa, pred_pa): """计算特征点之间的距离和角度差异 Args: gt_pa: 真实特征点和角度 [N, 4] pred_pa: 预测特征点和角度 [N, 4] Returns: dfg2p: 真实到预测的距离 dfp2g: 预测到真实的距离 fag2p: 真实到预测的角度差 fap2g: 预测到真实的角度差 """ gt_points = gt_pa[:,:3] pred_points = pred_pa[:,:3] gt_angle = gt_pa[:,3] pred_angle = pred_pa[:,3] dfg2p = 0.0 dfp2g = 0.0 fag2p = 0.0 fap2g = 0.0 pred_kdtree = cKDTree(pred_points) dist1, idx1 = pred_kdtree.query(gt_points) dfg2p = dist1.mean() assert(idx1.shape[0] == gt_points.shape[0]) fag2p = np.abs(gt_angle - pred_angle[idx1]) gt_kdtree = cKDTree(gt_points) dist2, idx2 = gt_kdtree.query(pred_points) dfp2g = dist2.mean() fap2g = np.abs(pred_angle - gt_angle[idx2]) fag2p = fag2p.mean() fap2g = fap2g.mean() return dfg2p, dfp2g, fag2p, fap2g def load_and_process_single_model(line, gt_path, pred_mesh_path, args): """处理单个模型的评估 Args: line (str): 模型名称 gt_path (str): 真值路径 pred_mesh_path (str): 预测网格路径 args: 参数配置 Returns: dict: 包含该模型所有评估指标的字典 """ try: line = line.strip()[:-4] result = {'name': line} # 加载点云数据 test_xyz = os.path.join(gt_path, line+'_50k.xyz') ptnormal = np.loadtxt(test_xyz) # 加载预测网格 meshfile = os.path.join(pred_mesh_path, '{}_50k.ply'.format(line)) if not os.path.exists(meshfile): logger.warning(f'File not exists: {meshfile}') return None # 检查缓存 stat_file = meshfile + "_stat" if not args.regen and os.path.exists(stat_file) and os.path.getsize(stat_file) > 0: with open(stat_file, 'rb') as f: return pickle.load(f) # 计算网格距离指标 mesh = trimesh.load(meshfile) cd, hd, adm, ads, hd_pred2gt, hd_gt2pred = distance_p2mesh( ptnormal[:,:3], ptnormal[:,3:], mesh) result.update({ 'CD': cd, 'HD': hd, 'HDpred2gt': hd_pred2gt, 'HDgt2pred': hd_gt2pred, 'AngleDiffMean': adm, 'AngleDiffStd': ads }) # 计算特征点指标 gt_ptangle = np.loadtxt(os.path.join(gt_path, line + '_detectfea4e-3.ptangle')) pred_ptangle_path = meshfile[:-4]+'_4e-3.ptangle' if not os.path.exists(pred_ptangle_path) or args.regen: os.system('./evaluation/MeshFeatureSample/build/SimpleSample -i {} -o {} -s 4e-3'.format( meshfile, pred_ptangle_path)) pred_ptangle = np.loadtxt(pred_ptangle_path).reshape(-1,4) # 处理特征点结果 if len(gt_ptangle) == 0 or len(pred_ptangle) == 0: result.update({ 'FeaDfgt2pred': 0.0, 'FeaDfpred2gt': 0.0, 'FeaAnglegt2pred': 0.0, 'FeaAnglepred2gt': 0.0, 'FeaDf': 0.0, 'FeaAngle': 0.0 }) else: dfg2p, dfp2g, fag2p, fap2g = distance_fea(gt_ptangle, pred_ptangle) result.update({ 'FeaDfgt2pred': dfg2p, 'FeaDfpred2gt': dfp2g, 'FeaAnglegt2pred': fag2p, 'FeaAnglepred2gt': fap2g, 'FeaDf': (dfg2p + dfp2g) / 2.0, 'FeaAngle': (fag2p + fap2g) / 2.0 }) # 保存缓存 with open(stat_file, "wb") as f: pickle.dump(result, f) return result except Exception as e: logger.error(f"Error processing {line}: {str(e)}") return None def compute_all(): """计算所有模型的评估指标""" try: # 初始化结果字典 results = [] # 读取模型列表 with open(os.path.join(project_dir, 'evaluation', args.name_list), 'r') as f: lines = f.readlines() # 处理每个模型 for line in lines: result = load_and_process_single_model(line, args.gt_path, args.pred_path, args) if result: results.append(result) # 计算平均值 mean_result = {'name': 'mean'} for key in results[0].keys(): if key != 'name': mean_result[key] = sum(r[key] for r in results) / len(results) results.append(mean_result) # 保存结果 df = pd.DataFrame(results) df.to_csv('eval_results.csv', index=False) logger.info(f"Evaluation completed. Results saved to {os.path.abspath('eval_results.csv')}") except Exception as e: logger.error(f"Error in compute_all: {str(e)}") raise if __name__ == '__main__': compute_all()