scripts/eval.py

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="评测: 预测目录(含 <case>_y.obj) vs GT 目录(<case>.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)