import os, sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from utils.distance import calculate_distances_numpy, calculate_normal_consistency_score import os import numpy as np from tqdm import tqdm import trimesh import json import glob import multiprocessing # 导入多进程模块 MIN_FACE_AMOUNT = 0 MAX_FACE_AMOUNT = 100000 # =============================================================== # 第一步：你的距离计算函数 (无变化) # =============================================================== pc_num = 1024 def normalize(mesh, file_path): # ... 现有代码 ... bounds = np.array([mesh.vertices.min(axis=0), mesh.vertices.max(axis=0)]) mesh.vertices = mesh.vertices - (bounds[0] + bounds[1])[None, :] / 2 mesh.vertices = mesh.vertices / (bounds[1] - bounds[0]).max() mesh.vertices = mesh.vertices.clip(-0.5, 0.5) # 限制顶点坐标在[-0.5, 0.5] # =============================================================== # 第二步：读取 .obj 文件的函数 (无变化) # =============================================================== def sample_from_obj(file_path, alter=False): # ... 现有代码 ... """ 从 .obj 文件中加载顶点坐标。 """ mesh = trimesh.load(file_path, force="mesh") if alter: mesh.vertices = mesh.vertices[:, [1,2,0]] normalize(mesh, file_path) points, face_idx = mesh.sample(pc_num, return_index=True) normals = mesh.face_normals[face_idx] return points,normals # =============================================================== # [修改] 法向量一致性计算 # =============================================================== # =============================================================== # 第三步：[修改] 多进程工作函数 # =============================================================== def process_common_file_worker(base_name, gt_folder, folder_paths): """ [已修改] 这是一个工作进程执行的任务： 处理单个 base_name 对应的所有文件夹。 它会加载一次GT文件，然后遍历所有预测文件夹进行计算。 返回 L1, L2 CD 距离, |NC| 和 NC。 """ # 1. 查找并加载 GT 文件 gt_filename = base_name + '.glb' # ... 现有代码 ... gt_filepath = os.path.join(gt_folder, gt_filename) if not os.path.exists(gt_filepath): gt_filename = base_name + '.obj' # ... 现有代码 ... gt_filepath = os.path.join(gt_folder, gt_filename) if not os.path.exists(gt_filepath): # 警告：在工作进程中打印 (print) 不是线程安全的， # ... 现有代码 ... # 但返回空列表是安全的，主进程会处理它。 # print(f"\n警告：共同文件 '{base_name}' 对应的GT文件 '{gt_filepath}' 不存在，已跳过。") return [] # 返回空列表表示此文件处理失败 mesh = trimesh.load(gt_filepath, force="mesh") if len(mesh.faces) < MIN_FACE_AMOUNT or len(mesh.faces) > MAX_FACE_AMOUNT: return [] try: pcd_gt, normal_gt = sample_from_obj(gt_filepath) if pcd_gt is None or pcd_gt.size == 0: # ... 现有代码 ... # print(f"\n警告：加载GT文件 {gt_filename} 失败或点云为空。已跳过。") return [] except Exception as e: # ... 现有代码 ... # print(f"\n加载GT文件 {gt_filename} 时发生错误: {e}。已跳过。") return [] # 2. 遍历所有文件夹，加载 Pred 文件并计算 results_for_this_file = [] for folder_path in folder_paths: pred_filename = base_name + '_y.obj' # ... 现有代码 ... pred_filepath = os.path.join(folder_path, pred_filename) try: pcd_pred, normal_pred = sample_from_obj(pred_filepath) if pcd_pred is None or pcd_pred.size == 0: # ... 现有代码 ... # print(f"\n警告：加载文件 {pred_filepath} 失败或点云为空。已跳过。") results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None continue # [修改] 接收 L1 和 L2 CD chamfer_l2, hausdorff, chamfer_l1 = calculate_distances_numpy(pcd_gt, pcd_pred) # [修改] 接收两个 NC 分数 normal_error_abs, normal_error_real = calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred) # [修改] 存储有效结果 (6个值) results_for_this_file.append((folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real)) except Exception as e: # ... 现有代码 ... # print(f"\n计算文件对 {pred_filepath} 时发生错误: {e}。已跳过。") results_for_this_file.append((folder_path, None, None, None, None, None)) # [修改] 6 个 None # 3. 返回此 base_name 对应的所有结果 return results_for_this_file def starmap_helper(args): # ... 现有代码 ... """ [新增] 辅助函数，用于解包 pool.imap 的参数，使其能接受多个参数。 """ return process_common_file_worker(*args) # =============================================================== # 第四步：[修改] 评估函数，使用多进程 # =============================================================== def evaluate_common_files(folder_paths, gt_folder, base_dir): """ [已修改] 评估多个文件夹中 *共有* 的文件，并计算总平均指标和逐文件夹的平均指标。 现在使用 multiprocessing.Pool 来并行处理文件。 """ if not folder_paths: # ... 现有代码 ... print("错误：没有提供用于评估共同文件的文件夹。") return None print("\n" + "="*80) # ... 现有代码 ... print(" 正在计算共同文件指标 ".center(80, "=")) print(f"总共 {len(folder_paths)} 个文件夹参与比较。") # 1. 查找每个文件夹中的文件基础名（不含 '_y.obj'）(无变化) files_by_folder = {} # ... 现有代码 ... for folder_path in folder_paths: try: all_files = os.listdir(folder_path) # ... 现有代码 ... base_names = set([f.replace('_y.obj', '') for f in all_files if (f.endswith('_y.obj') and not f.endswith('_gt_y.obj') and not f.endswith("input_y.obj"))] ) files_by_folder[folder_path] = base_names # ... 现有代码 ... print(f" ‐ 文件夹 '{os.path.basename(folder_path)}' 找到 {len(base_names)} 个文件。") except Exception as e: print(f" ‐ 错误：无法读取文件夹 {folder_path}: {e}") # ... 现有代码 ... files_by_folder[folder_path] = set() # 2. 找到所有文件夹中 *共有* 的文件基础名 (无变化) common_base_names = set.intersection(*files_by_folder.values()) if not common_base_names: # ... 现有代码 ... print("\n错误：在所有文件夹中没有找到任何共有的 `_y.obj` 文件。") print("="*80) return None print(f"\n在所有文件夹中找到了 {len(common_base_names)} 个共同文件。将使用多进程处理...") # 3. [修改] 初始化指标存储 global_chamfer_l2 = [] global_hausdorff = [] global_normal_abs = [] # [修改] global_chamfer_l1 = [] global_normal_real = [] # [新增] per_folder_metrics = {path: {'chamfer_l2': [], 'hausdorff': [], 'normal_abs': [], 'chamfer_l1': [], 'normal_real': []} for path in folder_paths} # [修改] # 4. [修改] 设置并运行多进程任务 # 准备任务参数列表 # ... 现有代码 ... task_args = [(base_name, gt_folder, folder_paths) for base_name in common_base_names] all_results_list = [] # 确定进程数，可以适当调整，例如 os.cpu_count() # ... 现有代码 ... # 留出2个核心给系统，防止卡死 num_processes = max(1, 60) print(f"启动 {num_processes} 个工作进程...") # 创建进程池 # ... 现有代码 ... with multiprocessing.Pool(processes=num_processes) as pool: # 使用 pool.imap 来迭代任务，并用 tqdm 显示进度条 # imap 会按顺序返回结果，同时保持并行执行 # ... 现有代码 ... for result in tqdm(pool.imap(starmap_helper, task_args), total=len(common_base_names), desc="评估共同文件"): all_results_list.append(result) # 5. [修改] 聚合所有进程返回的结果 print("\n正在聚合所有进程的结果...") valid_pairs_processed = 0 for results_for_this_file in all_results_list: if not results_for_this_file: # ... 现有代码 ... # GT 文件加载失败或跳过 continue # [修改] 解包 L1, L2, |NC|, NC for (folder_path, chamfer_l2, hausdorff, normal_error_abs, chamfer_l1, normal_error_real) in results_for_this_file: if chamfer_l2 is not None and hausdorff is not None and normal_error_abs is not None and chamfer_l1 is not None and normal_error_real is not None: # 这是一个有效的结果 global_chamfer_l2.append(chamfer_l2) global_hausdorff.append(hausdorff) global_normal_abs.append(normal_error_abs) # [修改] global_chamfer_l1.append(chamfer_l1) global_normal_real.append(normal_error_real) # [新增] per_folder_metrics[folder_path]['chamfer_l2'].append(chamfer_l2) per_folder_metrics[folder_path]['hausdorff'].append(hausdorff) per_folder_metrics[folder_path]['normal_abs'].append(normal_error_abs) # [修改] per_folder_metrics[folder_path]['chamfer_l1'].append(chamfer_l1) per_folder_metrics[folder_path]['normal_real'].append(normal_error_real) # [新增] valid_pairs_processed += 1 else: # ... 现有代码 ... # 预测文件加载失败或计算错误 # 可以在这里添加日志 pass print(f"聚合完成。总共处理了 {valid_pairs_processed} 个有效的 (GT, Pred) 文件对。") # 6. [修改] 计算最终平均值 if not global_chamfer_l2: # [修改] print("\n错误：未能成功计算任何共同文件的指标。") print("="*80) return None # [修改] 计算总平均 global_avg_results = { "evaluated_pairs_total": len(global_chamfer_l2), "common_files_count": len(common_base_names), "folders_count": len(folder_paths), "avg_chamfer_distance_l2": np.mean(global_chamfer_l2), "avg_hausdorff_distance": np.mean(global_hausdorff), "avg_normal_consistency_abs": np.mean(global_normal_abs), # [修改] "avg_chamfer_distance_l1": np.mean(global_chamfer_l1), "avg_normal_consistency_real": np.mean(global_normal_real) # [新增] } # [修改] 计算逐文件夹平均 per_folder_avg_results = {} for folder_path, metrics in per_folder_metrics.items(): folder_name = os.path.basename(folder_path) if metrics['chamfer_l2']: # [修改] per_folder_avg_results[folder_name] = { "evaluated_pairs": len(metrics['chamfer_l2']), "avg_chamfer_distance_l2": np.mean(metrics['chamfer_l2']), "avg_hausdorff_distance": np.mean(metrics['hausdorff']), "avg_normal_consistency_abs": np.mean(metrics['normal_abs']), # [修改] "avg_chamfer_distance_l1": np.mean(metrics['chamfer_l1']), "avg_normal_consistency_real": np.mean(metrics['normal_real']) # [新增] } else: per_folder_avg_results[folder_name] = { "evaluated_pairs": 0, "avg_chamfer_distance_l2": 0, "avg_hausdorff_distance": 0, "avg_normal_consistency_abs": 0, # [修改] "avg_chamfer_distance_l1": 0, "avg_normal_consistency_real": 0 # [新增] } # 7. [修改] 打印共同文件评估总结 print("\n--- 共同文件评估总结 ---") print(f"\n全局平均 (基于 {len(global_chamfer_l2)} 个总比较对, {len(common_base_names)} 个文件 x {len(folder_paths)} 个文件夹):") print(f" ‐ 平均倒角距离 (L2): {global_avg_results['avg_chamfer_distance_l2']:.6f}") print(f" ‐ 平均倒角距离 (L1): {global_avg_results['avg_chamfer_distance_l1']:.6f}") print(f" ‐ 平均豪斯多夫距离: {global_avg_results['avg_hausdorff_distance']:.6f}") print(f" ‐ 平均法向量 |NC| (Abs): {global_avg_results['avg_normal_consistency_abs']:.6f}") # [修改] print(f" ‐ 平均法向量 NC (Real): {global_avg_results['avg_normal_consistency_real']:.6f}") # [新增] print("\n逐文件夹平均 (仅基于共同文件):") for folder_name, results in per_folder_avg_results.items(): print(f" ‐ 文件夹: {folder_name}") # [修改] 打印所有6个指标 print(f" ‐ 倒角(L2): {results['avg_chamfer_distance_l2']:.6f}, 倒角(L1): {results['avg_chamfer_distance_l1']:.6f}, 豪斯多夫: {results['avg_hausdorff_distance']:.6f}, 法向量|NC|: {results['avg_normal_consistency_abs']:.6f}, 法向量NC: {results['avg_normal_consistency_real']:.6f} (共 {results['evaluated_pairs']} 对)") print("="*80) # 8. 保存结果 (无变化) final_results = { # ... 现有代码 ... "global_average_on_common_files": global_avg_results, "per_folder_average_on_common_files": per_folder_avg_results } output_filepath = os.path.join(base_dir, "common_evaluation_results.json") # ... 现有代码 ... try: with open(output_filepath, 'w', encoding='utf-8') as f: json.dump(final_results, f, ensure_ascii=False, indent=4) # ... 现有代码 ... print(f"\n共同文件评估结果已保存到: {output_filepath}") except Exception as e: print(f"\n错误：保存共同文件结果时失败: {e}") # ... 现有代码 ... return final_results # --- [主函数：run_evaluation_on_common_files_only] (打印部分已修改) --- def run_evaluation_on_common_files_only(target_folders, gt_folder, base_dir=None): """ [无变化] 在手动指定的多个文件夹上 *仅* 运行 *共同文件* 评估。 """ # 1. 验证输入的文件夹列表 # ... 现有代码 ... if not isinstance(target_folders, list) or not target_folders: print("错误：未提供任何文件夹用于评估。请检查 'target_folders' 列表。") return # 2. 检查路径有效性并过滤非文件夹 # ... 现有代码 ... valid_folders = [] for f_path in target_folders: if os.path.isdir(f_path): # ... 现有代码 ... valid_folders.append(f_path) else: print(f"警告：提供的路径 '{f_path}' 不是一个有效的文件夹，已跳过。") if not valid_folders: # ... 现有代码 ... print("错误：提供的列表中没有找到任何有效的文件夹。") return # 确保排序一致 # ... 现有代码 ... valid_folders.sort() # 3. 确定用于保存汇总文件的 base_dir if base_dir is None: # ... 现有代码 ... # 如果未指定，使用第一个有效文件夹的父目录 base_dir = os.path.dirname(valid_folders[0]) print(f"信息：未指定 base_dir，将使用 '{base_dir}' 存放共同文件评估结果。") elif not os.path.isdir(base_dir): # ... 现有代码 ... print(f"警告：指定的 base_dir '{base_dir}' 不存在。将改用第一个文件夹的父目录。") base_dir = os.path.dirname(valid_folders[0]) print(f"\n找到了 {len(valid_folders)} 个有效的文件夹，将评估它们的共同文件：") # ... 现有代码 ... for folder in valid_folders: print(f" - {folder}") # 打印完整路径 # 4. 第1部分：直接进行共同文件评估 # ... 现有代码 ... print("\n" + "#"*80) print(" 开始共同文件评估 (仅计算所有文件夹共有的文件) ".center(80, "#")) print("#"*80) # [调用修改后的函数] # ... 现有代码 ... common_results = evaluate_common_files(valid_folders, gt_folder, base_dir) # 5. [修改] 第2部分：打印最终的汇总报告 print("\n" + "#"*80) print(" 最终汇总报告 ".center(80, "#")) print("#"*80) print("\n--- 共同文件评估结果 (仅基于共同文件) ---") if common_results: # 打印全局平均 global_avg = common_results.get("global_average_on_common_files", {}) if global_avg: print(f"\n [总平均 (基于 {global_avg.get('common_files_count', 0)} 个共同文件)]") print(f" ‐ 倒角(L2): {global_avg.get('avg_chamfer_distance_l2', 0):.6f}, " f"倒角(L1): {global_avg.get('avg_chamfer_distance_l1', 0):.6f}, " f"豪斯多夫: {global_avg.get('avg_hausdorff_distance', 0):.6f}, " f"法向量|NC|: {global_avg.get('avg_normal_consistency_abs', 0):.6f}, " # [修改] f"法向量NC: {global_avg.get('avg_normal_consistency_real', 0):.6f}") # [新增] # 打印逐文件夹平均 per_folder_avg = common_results.get("per_folder_average_on_common_files", {}) if per_folder_avg: print("\n [逐文件夹平均 (仅共同文件)]") for folder_name, results in per_folder_avg.items(): print(f" ‐ 文件夹: {folder_name}") print(f" ‐ 倒角(L2): {results.get('avg_chamfer_distance_l2', 0):.6f}, " f"倒角(L1): {results.get('avg_chamfer_distance_l1', 0):.6f}, " f"豪斯多夫: {results.get('avg_hausdorff_distance', 0):.6f}, " f"法向量|NC|: {results.get('avg_normal_consistency_abs', 0):.6f}, " # [修改] f"法向量NC: {results.get('avg_normal_consistency_real', 0):.6f} " # [新增] f"(共 {results.get('evaluated_pairs', 0)} 对)") else: print(" (没有共同文件的评估结果)") print("\n" + "="*80) print("所有评估任务已完成。") print("="*80) # =============================================================== # 第五步：运行主程序 (无变化) # =============================================================== if __name__ == "__main__": import argparse ap = argparse.ArgumentParser(description="评测: 预测目录(含 _y.obj) vs GT 目录(.glb/.obj)") ap.add_argument("folders", nargs="+", help="一个或多个预测目录(只评估各目录共有的 case)") ap.add_argument("--gt", required=True, help="Ground Truth 目录") ap.add_argument("--out", default=None, help="汇总 JSON 的输出目录(默认第一个预测目录)") a = ap.parse_args() if not os.path.isdir(a.gt): raise SystemExit(f"GT 目录不存在: {a.gt}") run_evaluation_on_common_files_only(a.folders, a.gt, base_dir=a.out)