FE2E/evaluation.py

import argparse
import os
import re
import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from infer.seed_all import seed_all

# 设置环境变量消除tokenizers警告
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
os.environ['NCCL_DEBUG'] = 'WARN'
# 消除torchvision警告
os.environ['TORCHVISION_DISABLE_DEPRECATED_WARNING'] = '1'

REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
DEFAULT_QWEN_DIR = os.path.join(REPO_ROOT, "Qwen")
DEFAULT_DEPTH_DATASET_CONFIGS = {
    "nyu_v2": "configs/data_nyu_test.yaml",
    "kitti": "configs/data_kitti_eigen_test.yaml",
    "eth3d": "configs/data_eth3d.yaml",
    "diode": "configs/data_diode_all.yaml",
    "scannet": "configs/data_scannet_val.yaml",
}
DEFAULT_NORMAL_DATASETS = {
    "nyuv2": "test",
    "scannet": "test",
    "ibims": "ibims",
    "sintel": "sintel",
    "oasis": "val",
    "hypersim": "hypersim",
}


def resolve_eval_data_root(args, *required_markers):
    """Resolve the evaluation data root without depending on the launch cwd."""
    candidates = []
    if getattr(args, "eval_data_root", None):
        candidates.append(os.path.abspath(args.eval_data_root))

    candidates.extend(
        [
            os.path.join(REPO_ROOT, "infer"),
            os.path.join(REPO_ROOT, "data"),
            os.path.join(os.path.dirname(REPO_ROOT), "data"),
        ]
    )

    for candidate in candidates:
        if all(os.path.exists(os.path.join(candidate, marker)) for marker in required_markers):
            return candidate

    checked = ", ".join(
        os.path.join(candidate, marker) for candidate in candidates for marker in required_markers
    )
    raise FileNotFoundError(f"未找到评测数据根目录，已检查: {checked}")


def parse_depth_eval_datasets(raw_value):
    requested = [item.strip() for item in raw_value.split(",") if item.strip()]
    if requested == ["all"]:
        requested = list(DEFAULT_DEPTH_DATASET_CONFIGS.keys())
    invalid = [item for item in requested if item not in DEFAULT_DEPTH_DATASET_CONFIGS]
    if invalid:
        raise ValueError(f"不支持的 depth 数据集: {invalid}，可选值: {sorted(DEFAULT_DEPTH_DATASET_CONFIGS)}")
    return {name: DEFAULT_DEPTH_DATASET_CONFIGS[name] for name in requested}


def parse_normal_eval_datasets(raw_value):
    requested = [item.strip() for item in raw_value.split(",") if item.strip()]
    if requested == ["all"]:
        requested = list(DEFAULT_NORMAL_DATASETS.keys())
    invalid = [item for item in requested if item not in DEFAULT_NORMAL_DATASETS]
    if invalid:
        raise ValueError(f"不支持的 normal 数据集: {invalid}，可选值: {sorted(DEFAULT_NORMAL_DATASETS)}")
    return [(name, DEFAULT_NORMAL_DATASETS[name]) for name in requested]


def collect_and_merge_dual_cfg_results(rank, world_size, gathered_metrics_Lpred, gathered_times):
    """
    收集并合并双CFG配置的评估结果
    
    Args:
        rank: 当前进程的rank
        world_size: 总进程数
        gathered_metrics_Lpred
        gathered_times: 处理时间收集结果
    
    Returns:
        tuple: (all_metrics_Lpred, dataset_times)
    """
    if rank != 0:
        return None, None

    # 合并处理时间
    dataset_times = []
    for times_list in gathered_times:
        dataset_times.extend(times_list)

    #先处理L的
    all_metrics_L = {}
    valid_metrics_L = [m for m in gathered_metrics_Lpred if m]
    if valid_metrics_L:
        for key in valid_metrics_L[0].keys():
            values = [m[key] for m in valid_metrics_L if key in m]
            if values:
                all_metrics_L[key] = np.mean(values)

    return all_metrics_L, dataset_times


def format_dual_cfg_results_table(dataset_name, model_identifier, all_metrics_L, dataset_times):
    """
    格式化双CFG配置的结果表格
    
    Args:
        dataset_name: 数据集名称
        model_identifier: 模型标识符
        all_metrics_L: CFG=1的评估指标
        dataset_times: 处理时间列表
    
    Returns:
        str: 格式化的结果字符串
    """
    eval_metrics_order = ["abs_relative_difference", "squared_relative_difference", "rmse_linear", "rmse_log", "delta1_acc", "delta2_acc", "delta3_acc"]

    # 获取CFG=1的指标值
    mean_errors_L = [all_metrics_L.get(metric, 0.0) for metric in eval_metrics_order]

    # 构建表格
    metrics_header = ["Dataset", "Model", "CFG", "abs_rel", "sq_rel", "rmse", "rmse_log", "a1", "a2", "a3"]

    # CFG=1的结果行
    values_data_L = [dataset_name, model_identifier, "CFG=1"] + [f"{v:.4f}" for v in mean_errors_L]

    header_line = "| " + " | ".join(metrics_header) + " |"
    separator_line = "| " + " | ".join(["---"] * len(metrics_header)) + " |"
    values_line_L = "| " + " | ".join(values_data_L) + " |"

    # 生成输出字符串
    result_str = f"\n数据集 {dataset_name} 评估完成!\n"
    result_str += "-" * 100 + "\n"
    result_str += header_line + "\n"
    result_str += separator_line + "\n"
    result_str += values_line_L + "\n"

    # 添加统计信息
    sample_count = len(dataset_times)
    result_str += f"样本数: {sample_count}\n"
    if dataset_times:
        result_str += f"平均处理时间: {np.mean(dataset_times):.2f}秒/图像\n"
    result_str += "-" * 100 + "\n"

    return result_str


def save_dual_cfg_results_summary(output_dir, all_dataset_results, model_identifier):
    """
    保存所有数据集的双CFG评估结果摘要
    
    Args:
        output_dir: 输出目录
        all_dataset_results: 所有数据集的结果字典
    """
    summary_file = os.path.join(output_dir, f"{model_identifier}.txt")

    with open(summary_file, 'w', encoding='utf-8') as f:
        f.write("=" * 120 + "\n")
        f.write("双CFG配置深度评估结果汇总\n")
        f.write("=" * 120 + "\n\n")

        for dataset_name, result_data in all_dataset_results.items():
            f.write(result_data['formatted_output'])
            f.write(f"结果保存至: {result_data['eval_dir']}\n\n")

    print(f"双CFG评估结果摘要已保存至: {summary_file}")


def parse_args():
    '''Set the Args'''
    parser = argparse.ArgumentParser(description="Run Step...")
    parser.add_argument('--model_path', type=str, default='./pretrain', help='模型路径')
    parser.add_argument('--qwen2vl_model_path', type=str, default=DEFAULT_QWEN_DIR, help='Qwen2.5-VL 模型目录')
    parser.add_argument('--eval_data_root', type=str, default=None, help='评测数据根目录，默认自动在仓库相对路径下查找')
    parser.add_argument("--seed", type=int, default=1234, help="随机种子")
    parser.add_argument("--output_dir", type=str, default="./infer/eval_results", help="Output directory.")
    parser.add_argument('--num_steps', type=int, default=28, help='扩散步数')
    parser.add_argument('--num_samples', type=int, default=-1, help='生成样本数')
    parser.add_argument('--cfg_guidance', type=float, default=6.0, help='CFG引导强度')
    parser.add_argument('--size_level', type=int, default=768, help='输入图像大小')
    parser.add_argument('--num_gpus', type=int, default=torch.cuda.device_count(), help='使用的GPU数量')
    parser.add_argument('--save_viz', action='store_true', help='保存可视化结果')
    parser.add_argument('--offload', action='store_true', help='使用CPU卸载以节省GPU内存')
    parser.add_argument('--quantized', action='store_true', help='使用量化模型')
    parser.add_argument('--lora', type=str, help='LoRA模型路径')
    parser.add_argument('--single_denoise', action='store_true', default=False, help='单步推理')
    parser.add_argument('--old_prompt', action='store_true', default=False, help='使用旧版提示')
    parser.add_argument('--prompt_type', type=str, default='query', help='提示类型')
    parser.add_argument('--prompt', type=str, default='Describe the 3D structure and layout of the scene in the image. Predict the depth of this image.', help='提示')
    parser.add_argument('--norm_type', type=str, default='depth', help='预测结果的归一化方式，目前有 depth、disp、ln')
    parser.add_argument('--task_name', type=str, default='depth', help='任务名称，支持 depth 或 normal')
    parser.add_argument('--depth_eval_datasets', type=str, default='eth3d', help='逗号分隔的 depth 评测数据集')
    parser.add_argument('--normal_eval_datasets', type=str, default='nyuv2,scannet', help='逗号分隔的 normal 评测数据集')
    parser.add_argument('--debug', action='store_true', default=False, help='调试模式')
    args = parser.parse_args()
    if args.single_denoise:
        args.num_steps = 1
    return args


def extract_model_identifier(lora_path):
    """
    从lora路径中提取模型标识符
    支持多种路径格式：
    - ./log_err/dis-hvsge-log/ckpt.safetensors -> dis-hvsge-log
    - /path/to/folder/ckpt-123 -> folder-epoch123
    - /path/to/model.safetensors -> model
    """
    if not lora_path or not os.path.exists(lora_path):
        return "DefaultModel"
    
    # 规范化路径
    lora_path = os.path.normpath(lora_path)
    
    # 方法1: 匹配 /folder/ckpt-数字 格式
    match = re.search(r'/([^/]+)/ckpt-(\d+)', lora_path)
    if match:
        folder_name = match.group(1)
        epoch_num = int(match.group(2))
        return f"{folder_name}-epoch{epoch_num}"
    
    # 方法2: 匹配 /folder/ckpt.safetensors 格式
    match = re.search(r'/([^/]+)/ckpt\.safetensors$', lora_path)
    if match:
        return match.group(1)
    
    # 方法3: 匹配 ./folder/ckpt.safetensors 格式
    match = re.search(r'[./]*([^/]+)/ckpt\.safetensors$', lora_path)
    if match:
        return match.group(1)
    
    # 方法4: 从文件名中提取（兜底方案）
    filename = os.path.basename(lora_path)
    return filename.split('.')[0] if '.' in filename else filename


def setup(rank, world_size):
    """初始化分布式环境"""
    os.environ.setdefault('MASTER_ADDR', 'localhost')
    os.environ.setdefault('MASTER_PORT', '21256')
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def cleanup():
    """清理分布式环境"""
    dist.destroy_process_group()


def main_worker(rank, world_size, args, dataset_configs):
    """每个进程的主函数"""
    from infer.inference import ImageGenerator
    from infer.inner_evaluation import evaluation_depth_custom_parallel

    setup(rank, world_size)
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    if rank == 0:
        print(f"[main_worker] 开始加载 pipeline, device={device}, datasets={list(dataset_configs.keys())}", flush=True)

    pipeline = ImageGenerator(
        ae_path=os.path.join(args.model_path, 'vae.safetensors'),
        dit_path=os.path.join(args.model_path, "step1x-edit-i1258-FP8.safetensors" if args.quantized else "step1x-edit-i1258.safetensors"),
        qwen2vl_model_path=args.qwen2vl_model_path,
        max_length=640,
        quantized=args.quantized,
        offload=args.offload,
        lora=args.lora,
        device=str(device),
        args=args,
    )

    if rank == 0:
        print(f"Successfully loading pipeline from {args.model_path}.", flush=True)

    test_data_dir = resolve_eval_data_root(args, "configs")

    # 使用新的模型标识符提取函数
    model_identifier = extract_model_identifier(args.lora)

    if rank == 0:
        print(f"模型标识符: {model_identifier}", flush=True)

    all_dataset_results = {}
    aligment_map = {"depth": "least_square", "disp": "least_square_disparity", "ln": "log_space"}
    for dataset_name, config_path in dataset_configs.items():
        # 修改输出目录结构：在数据集名称外添加模型名称层级
        eval_dir = os.path.join(args.output_dir, model_identifier, dataset_name)
        test_dataset_config = os.path.join(test_data_dir, config_path)
        alignment_type = aligment_map[args.norm_type]

        if rank == 0:
            print(f"\n开始评估数据集: {dataset_name}", flush=True)
            print(f"输出目录: {eval_dir}", flush=True)
            print("=" * 80, flush=True)

        metric_tracker_Lpred, metric_tracker_Rpred, processing_times = evaluation_depth_custom_parallel(
            rank,
            world_size,
            eval_dir,
            test_dataset_config,
            args,
            pipeline,
            test_data_dir,
            alignment=alignment_type,
            save_pred_vis=args.save_viz,
        )

        # 同步所有进程
        dist.barrier()

        # 收集两个CFG配置的结果
        gathered_metrics_Lpred = [None] * world_size
        gathered_metrics_Rpred = [None] * world_size
        gathered_times = [None] * world_size

        # 从metric_tracker获取结果字典
        metrics_dict_Lpred = metric_tracker_Lpred.result() if hasattr(metric_tracker_Lpred, 'result') else {}
        # metrics_dict_Rpred = metric_tracker_Rpred.result() if hasattr(metric_tracker_Rpred, 'result') else {}

        dist.all_gather_object(gathered_metrics_Lpred, metrics_dict_Lpred)
        # dist.all_gather_object(gathered_metrics_Rpred, metrics_dict_Rpred)
        dist.all_gather_object(gathered_times, processing_times)

        if rank == 0:
            metrics_dict_Lpred, dataset_times = collect_and_merge_dual_cfg_results(rank, world_size, gathered_metrics_Lpred, gathered_times)

            if metrics_dict_Lpred:
                # 格式化并输出结果表格
                formatted_output = format_dual_cfg_results_table(dataset_name, model_identifier, metrics_dict_Lpred, dataset_times)

                print(formatted_output)
                print(f"结果保存至: {eval_dir}")

                # 存储结果用于后续汇总
                all_dataset_results[dataset_name] = {'metrics_Lpred': metrics_dict_Lpred, 'formatted_output': formatted_output, 'eval_dir': eval_dir, 'processing_times': dataset_times}

    if rank == 0:
        print(f"\n所有数据集评估完成! 结果保存在: {os.path.join(args.output_dir, model_identifier)}")

        # 保存所有数据集的双CFG评估结果摘要
        if all_dataset_results:
            save_dual_cfg_results_summary(os.path.join(args.output_dir, model_identifier), all_dataset_results, model_identifier)

    cleanup()

def main_worker_normal(rank, world_size, args, eval_datasets):
    """normal预测的多进程主函数"""
    from infer.inference import ImageGenerator
    from infer.inner_evaluation import evaluation_normal_custom_parallel

    setup(rank, world_size)
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    pipeline = ImageGenerator(
        ae_path=os.path.join(args.model_path, 'vae.safetensors'),
        dit_path=os.path.join(args.model_path, "step1x-edit-i1258-FP8.safetensors" if args.quantized else "step1x-edit-i1258.safetensors"),
        qwen2vl_model_path=args.qwen2vl_model_path,
        max_length=640,
        quantized=args.quantized,
        offload=args.offload,
        lora=args.lora,
        device=str(device),
        args=args,
    )

    if rank == 0:
        print(f"Successfully loading pipeline from {args.model_path}.")

    test_data_dir = resolve_eval_data_root(args, "dsine_eval")
    dataset_split_path = os.path.join(REPO_ROOT, "infer", "dataset_normal")
    
    # 使用新的模型标识符提取函数
    model_identifier = extract_model_identifier(args.lora)
    
    # 修改输出目录结构：在任务名称外添加模型名称层级
    eval_dir = os.path.join(args.output_dir, model_identifier, args.task_name)

    if rank == 0:
        print(f"模型标识符: {model_identifier}")
        print(f"输出目录: {eval_dir}")

    if rank == 0:
        print(f"\n开始并行Normal评估，使用{world_size}个GPU")
        print("=" * 80)

    # 调用并行评估函数
    all_normal_errors, all_processing_times, all_dataset_metrics = evaluation_normal_custom_parallel(
        rank, world_size, eval_dir, test_data_dir, dataset_split_path, pipeline, args, eval_datasets, save_pred_vis=args.save_viz
    )

    # 同步所有进程
    dist.barrier()

    # 收集所有GPU的结果
    gathered_normal_errors = [None] * world_size
    gathered_processing_times = [None] * world_size
    gathered_dataset_metrics = [None] * world_size

    dist.all_gather_object(gathered_normal_errors, all_normal_errors)
    dist.all_gather_object(gathered_processing_times, all_processing_times)
    dist.all_gather_object(gathered_dataset_metrics, all_dataset_metrics)

    if rank == 0:
        # 合并所有GPU的结果
        final_results = {}
        
        for dataset_name, _ in eval_datasets:
            print(f"\n合并数据集 {dataset_name} 的结果...")
            
            # 合并normal errors
            all_errors_for_dataset = []
            all_times_for_dataset = []
            
            for gpu_errors, gpu_times in zip(gathered_normal_errors, gathered_processing_times):
                if gpu_errors[dataset_name] is not None:
                    all_errors_for_dataset.append(gpu_errors[dataset_name])
                if gpu_times[dataset_name]:
                    all_times_for_dataset.extend(gpu_times[dataset_name])
            
            # 计算最终指标
            if all_errors_for_dataset:
                combined_errors = torch.cat(all_errors_for_dataset, dim=0)
                from infer.util import normal_utils
                final_metrics = normal_utils.compute_normal_metrics(combined_errors)
                
                print(f"数据集 {dataset_name} 最终结果:")
                print("mean median rmse 5 7.5 11.25 22.5 30")
                print("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f" % (
                    final_metrics['mean'], final_metrics['median'], final_metrics['rmse'], 
                    final_metrics['a1'], final_metrics['a2'], final_metrics['a3'], 
                    final_metrics['a4'], final_metrics['a5']
                ))
                
                final_results[dataset_name] = {
                    'metrics': final_metrics,
                    'processing_times': all_times_for_dataset,
                    'sample_count': len(combined_errors)
                }
                
                # 保存结果到文件
                dataset_output_dir = os.path.join(eval_dir, dataset_name)
                os.makedirs(dataset_output_dir, exist_ok=True)
                
                from tabulate import tabulate
                eval_text = f"Evaluation metrics for {dataset_name}:\n"
                eval_text += f"Total samples: {len(combined_errors)}\n"
                eval_text += f"Average processing time: {np.mean(all_times_for_dataset):.2f}s\n"
                eval_text += tabulate([list(final_metrics.keys()), list(final_metrics.values())])
                
                save_path = os.path.join(dataset_output_dir, "eval_metrics.txt")
                with open(save_path, "w+") as f:
                    f.write(eval_text)
                
                print(f"结果已保存至: {save_path}")
            else:
                print(f"数据集 {dataset_name}: 未找到有效数据")
                final_results[dataset_name] = None

        # 保存总体结果摘要
        summary_file = os.path.join(eval_dir, f"{model_identifier}_normal_summary.txt")
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write("=" * 120 + "\n")
            f.write("Normal预测多GPU并行评估结果汇总\n")
            f.write("=" * 120 + "\n\n")
            
            for dataset_name, result in final_results.items():
                if result is not None:
                    f.write(f"数据集: {dataset_name}\n")
                    f.write(f"样本数: {result['sample_count']}\n")
                    f.write(f"平均处理时间: {np.mean(result['processing_times']):.2f}s\n")
                    metrics = result['metrics']
                    f.write("mean median rmse 5 7.5 11.25 22.5 30\n")
                    f.write("%.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (
                        metrics['mean'], metrics['median'], metrics['rmse'], 
                        metrics['a1'], metrics['a2'], metrics['a3'], 
                        metrics['a4'], metrics['a5']
                    ))
                    f.write("-" * 60 + "\n\n")
                else:
                    f.write(f"数据集: {dataset_name} - 无有效数据\n\n")
        
        print(f"\nNormal评估总结已保存至: {summary_file}")

    cleanup()
    
def main():
    args = parse_args()
    if args.seed is not None:
        seed_all(args.seed)

    os.makedirs(args.output_dir, exist_ok=True)

    # 检查GPU数量
    world_size = min(args.num_gpus, torch.cuda.device_count())
    if world_size <= 0:
        print("错误：未检测到可用的GPU。")
        return

    # 设置多进程相关环境变量
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'

    print(f"即将使用 {world_size} 个GPU进行并行推理...")

    if args.task_name == 'depth':
        test_depth_dataset_configs = parse_depth_eval_datasets(args.depth_eval_datasets)
        print(f"Depth评估数据集: {list(test_depth_dataset_configs.keys())}")

        if world_size == 1:
            # 单GPU情况，直接运行
            main_worker(0, 1, args, test_depth_dataset_configs)
        else:
            # 多GPU情况，使用multiprocessing
            try:
                mp.spawn(main_worker, args=(world_size, args, test_depth_dataset_configs), nprocs=world_size, join=True)
            except Exception as e:
                print(f"多进程执行出错: {e}")
                print("尝试降级到单GPU模式...")
                # 降级到单GPU模式
                main_worker(0, 1, args, test_depth_dataset_configs)

    elif args.task_name == 'normal':
        eval_datasets = parse_normal_eval_datasets(args.normal_eval_datasets)
        print(f"Normal评估数据集: {eval_datasets}")

        if world_size == 1:
            main_worker_normal(0, 1, args, eval_datasets)
        else:
            try:
                mp.spawn(main_worker_normal, args=(world_size, args, eval_datasets), nprocs=world_size, join=True)
            except Exception as e:
                print(f"多进程执行出错: {e}")
                print("尝试降级到单GPU模式...")
                # 降级到单GPU模式
                main_worker_normal(0, 1, args, eval_datasets)
    else:
        raise ValueError(f"不支持的 task_name: {args.task_name}，仅支持 depth 或 normal")

if __name__ == '__main__':
    main()