RoRD-Layout-Recognation/tests/benchmark_fpn.py

"""
FPN vs 滑窗性能对标脚本

功能：比较 FPN 推理路径与传统图像金字塔滑窗路径的性能差异。

输出指标：
  - 推理时间（ms）
  - 内存占用（MB）
  - 检测到的关键点数
  - 检测精度（匹配内点数）

使用示例：
  uv run python tests/benchmark_fpn.py \
    --layout /path/to/layout.png \
    --template /path/to/template.png \
    --num-runs 5 \
    --output benchmark_results.json
"""

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import psutil
import torch
from PIL import Image

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from models.rord import RoRD
from utils.config_loader import load_config, to_absolute_path
from utils.data_utils import get_transform


def get_memory_usage() -> float:
    """获取当前进程的内存占用（MB）"""
    process = psutil.Process()
    return process.memory_info().rss / 1024 / 1024


def get_gpu_memory_usage() -> float:
    """获取 GPU 显存占用（MB）"""
    if torch.cuda.is_available():
        return torch.cuda.memory_allocated() / 1024 / 1024
    return 0


def benchmark_fpn(
    model: torch.nn.Module,
    layout_image: Image.Image,
    template_image: Image.Image,
    transform,
    matching_cfg,
    num_runs: int = 5,
) -> Dict[str, float]:
    """
    测试 FPN 性能
    
    Args:
        model: RoRD 模型
        layout_image: 大版图
        template_image: 模板
        transform: 图像预处理管道
        matching_cfg: 匹配配置
        num_runs: 运行次数
    
    Returns:
        性能指标字典
    """
    from match import extract_from_pyramid, extract_features_sliding_window, mutual_nearest_neighbor
    
    device = next(model.parameters()).device
    times = []
    keypoint_counts = []
    inlier_counts = []
    
    print(f"\n{'=' * 60}")
    print(f"性能测试：FPN 路径")
    print(f"{'=' * 60}")
    
    for run in range(num_runs):
        # 版图特征提取
        layout_tensor = transform(layout_image).unsqueeze(0).to(device)
        
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        start_time = time.time()
        
        layout_kps, layout_descs = extract_from_pyramid(
            model, 
            layout_tensor, 
            float(matching_cfg.keypoint_threshold),
            getattr(matching_cfg, 'nms', {})
        )
        
        # 模板特征提取（单尺度，取 1.0）
        template_tensor = transform(template_image).unsqueeze(0).to(device)
        template_kps, template_descs = extract_from_pyramid(
            model,
            template_tensor,
            float(matching_cfg.keypoint_threshold),
            getattr(matching_cfg, 'nms', {})
        )
        
        # 匹配
        if len(layout_descs) > 0 and len(template_descs) > 0:
            matches = mutual_nearest_neighbor(template_descs, layout_descs)
            inlier_count = len(matches)
        else:
            inlier_count = 0
        
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        elapsed = (time.time() - start_time) * 1000  # 转换为 ms
        
        times.append(elapsed)
        keypoint_counts.append(len(layout_kps))
        inlier_counts.append(inlier_count)
        
        print(f"  Run {run + 1}/{num_runs}: {elapsed:.2f}ms, KPs: {len(layout_kps)}, Matches: {inlier_count}")
    
    mean_time = np.mean(times)
    std_time = np.std(times)
    mean_kps = np.mean(keypoint_counts)
    mean_inliers = np.mean(inlier_counts)
    gpu_mem = get_gpu_memory_usage()
    
    return {
        "method": "FPN",
        "mean_time_ms": float(mean_time),
        "std_time_ms": float(std_time),
        "min_time_ms": float(np.min(times)),
        "max_time_ms": float(np.max(times)),
        "all_times_ms": [float(t) for t in times],
        "mean_keypoints": float(mean_kps),
        "mean_matches": float(mean_inliers),
        "gpu_memory_mb": float(gpu_mem),
        "num_runs": num_runs,
    }


def benchmark_sliding_window(
    model: torch.nn.Module,
    layout_image: Image.Image,
    template_image: Image.Image,
    transform,
    matching_cfg,
    num_runs: int = 5,
) -> Dict[str, float]:
    """
    测试滑窗性能（图像金字塔路径）
    
    Args:
        model: RoRD 模型
        layout_image: 大版图
        template_image: 模板
        transform: 图像预处理管道
        matching_cfg: 匹配配置
        num_runs: 运行次数
    
    Returns:
        性能指标字典
    """
    from match import extract_features_sliding_window, extract_keypoints_and_descriptors, mutual_nearest_neighbor
    
    device = next(model.parameters()).device
    times = []
    keypoint_counts = []
    inlier_counts = []
    
    print(f"\n{'=' * 60}")
    print(f"性能测试：滑窗路径")
    print(f"{'=' * 60}")
    
    for run in range(num_runs):
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        start_time = time.time()
        
        # 版图滑窗特征提取
        layout_kps, layout_descs = extract_features_sliding_window(
            model, 
            layout_image, 
            transform, 
            matching_cfg
        )
        
        # 模板单尺度特征提取
        template_tensor = transform(template_image).unsqueeze(0).to(device)
        template_kps, template_descs = extract_keypoints_and_descriptors(
            model,
            template_tensor,
            float(matching_cfg.keypoint_threshold)
        )
        
        # 匹配
        if len(layout_descs) > 0 and len(template_descs) > 0:
            matches = mutual_nearest_neighbor(template_descs, layout_descs)
            inlier_count = len(matches)
        else:
            inlier_count = 0
        
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        elapsed = (time.time() - start_time) * 1000  # 转换为 ms
        
        times.append(elapsed)
        keypoint_counts.append(len(layout_kps))
        inlier_counts.append(inlier_count)
        
        print(f"  Run {run + 1}/{num_runs}: {elapsed:.2f}ms, KPs: {len(layout_kps)}, Matches: {inlier_count}")
    
    mean_time = np.mean(times)
    std_time = np.std(times)
    mean_kps = np.mean(keypoint_counts)
    mean_inliers = np.mean(inlier_counts)
    gpu_mem = get_gpu_memory_usage()
    
    return {
        "method": "Sliding Window",
        "mean_time_ms": float(mean_time),
        "std_time_ms": float(std_time),
        "min_time_ms": float(np.min(times)),
        "max_time_ms": float(np.max(times)),
        "all_times_ms": [float(t) for t in times],
        "mean_keypoints": float(mean_kps),
        "mean_matches": float(mean_inliers),
        "gpu_memory_mb": float(gpu_mem),
        "num_runs": num_runs,
    }


def compute_speedup(fpn_result: Dict, sw_result: Dict) -> Dict[str, float]:
    """计算 FPN 相对于滑窗的性能改进"""
    speedup = (sw_result["mean_time_ms"] - fpn_result["mean_time_ms"]) / sw_result["mean_time_ms"] * 100
    memory_saving = (sw_result["gpu_memory_mb"] - fpn_result["gpu_memory_mb"]) / sw_result["gpu_memory_mb"] * 100 if sw_result["gpu_memory_mb"] > 0 else 0
    
    return {
        "speedup_percent": float(speedup),
        "memory_saving_percent": float(memory_saving),
        "fpn_faster": speedup > 0,
        "meets_speedup_target": speedup >= 30,
        "meets_memory_target": memory_saving >= 20,
    }


def print_results(fpn_result: Dict, sw_result: Dict, comparison: Dict) -> None:
    """打印性能对比结果"""
    
    print(f"\n{'=' * 80}")
    print(f"{'性能基准测试结果':^80}")
    print(f"{'=' * 80}\n")
    
    print(f"{'指标':<30} {'FPN':<20} {'滑窗':<20}")
    print("-" * 70)
    
    print(f"{'平均推理时间 (ms)':<30} {fpn_result['mean_time_ms']:<20.2f} {sw_result['mean_time_ms']:<20.2f}")
    print(f"{'标准差 (ms)':<30} {fpn_result['std_time_ms']:<20.2f} {sw_result['std_time_ms']:<20.2f}")
    print(f"{'最小时间 (ms)':<30} {fpn_result['min_time_ms']:<20.2f} {sw_result['min_time_ms']:<20.2f}")
    print(f"{'最大时间 (ms)':<30} {fpn_result['max_time_ms']:<20.2f} {sw_result['max_time_ms']:<20.2f}")
    print()
    
    print(f"{'平均关键点数':<30} {fpn_result['mean_keypoints']:<20.0f} {sw_result['mean_keypoints']:<20.0f}")
    print(f"{'平均匹配数':<30} {fpn_result['mean_matches']:<20.0f} {sw_result['mean_matches']:<20.0f}")
    print()
    
    print(f"{'GPU 内存占用 (MB)':<30} {fpn_result['gpu_memory_mb']:<20.2f} {sw_result['gpu_memory_mb']:<20.2f}")
    print()
    
    print(f"{'=' * 80}")
    print(f"{'对标结果':^80}")
    print(f"{'=' * 80}\n")
    
    speedup = comparison["speedup_percent"]
    memory_saving = comparison["memory_saving_percent"]
    
    print(f"推理速度提升: {speedup:+.2f}% {'✅' if speedup >= 30 else '⚠️'}")
    print(f"  (目标: ≥30% | 达成: {'是' if comparison['meets_speedup_target'] else '否'})")
    print()
    
    print(f"内存节省: {memory_saving:+.2f}% {'✅' if memory_saving >= 20 else '⚠️'}")
    print(f"  (目标: ≥20% | 达成: {'是' if comparison['meets_memory_target'] else '否'})")
    print()
    
    if speedup > 0:
        print(f"🎉 FPN 相比滑窗快 {abs(speedup):.2f}%")
    elif speedup < 0:
        print(f"⚠️ FPN 相比滑窗慢 {abs(speedup):.2f}%")
    else:
        print(f"ℹ️ FPN 与滑窗性能相当")
    print()


def main():
    parser = argparse.ArgumentParser(description="RoRD FPN vs 滑窗性能对标测试")
    parser.add_argument('--config', type=str, default="configs/base_config.yaml", help="YAML 配置文件")
    parser.add_argument('--model_path', type=str, default=None, help="模型权重路径")
    parser.add_argument('--layout', type=str, required=True, help="版图路径")
    parser.add_argument('--template', type=str, required=True, help="模板路径")
    parser.add_argument('--num-runs', type=int, default=5, help="每个方法的运行次数")
    parser.add_argument('--output', type=str, default="benchmark_results.json", help="输出 JSON 文件路径")
    parser.add_argument('--device', type=str, default="cuda", help="使用设备: cuda 或 cpu")
    
    args = parser.parse_args()
    
    # 加载配置
    cfg = load_config(args.config)
    config_dir = Path(args.config).resolve().parent
    matching_cfg = cfg.matching
    
    model_path = args.model_path or str(to_absolute_path(cfg.paths.model_path, config_dir))
    
    # 设置设备
    device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
    print(f"使用设备: {device}")
    
    # 加载模型
    print(f"加载模型: {model_path}")
    model = RoRD().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    
    # 加载图像
    print(f"加载版图: {args.layout}")
    layout_image = Image.open(args.layout).convert('L')
    print(f"  尺寸: {layout_image.size}")
    
    print(f"加载模板: {args.template}")
    template_image = Image.open(args.template).convert('L')
    print(f"  尺寸: {template_image.size}")
    
    # 获取预处理管道
    transform = get_transform()
    
    # 运行基准测试
    print(f"\n{'=' * 80}")
    print(f"{'开始性能基准测试':^80}")
    print(f"{'=' * 80}")
    print(f"运行次数: {args.num_runs}")
    print(f"配置: {args.config}")
    
    with torch.no_grad():
        fpn_result = benchmark_fpn(
            model, layout_image, template_image, transform, matching_cfg, args.num_runs
        )
        
        # 临时禁用 FPN，启用滑窗
        original_use_fpn = getattr(matching_cfg, 'use_fpn', True)
        matching_cfg.use_fpn = False
        
        sw_result = benchmark_sliding_window(
            model, layout_image, template_image, transform, matching_cfg, args.num_runs
        )
        
        # 恢复配置
        matching_cfg.use_fpn = original_use_fpn
    
    # 计算对比指标
    comparison = compute_speedup(fpn_result, sw_result)
    
    # 打印结果
    print_results(fpn_result, sw_result, comparison)
    
    # 保存结果
    results = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "config": str(args.config),
        "model_path": str(model_path),
        "layout_path": str(args.layout),
        "layout_size": list(layout_image.size),
        "template_path": str(args.template),
        "template_size": list(template_image.size),
        "device": str(device),
        "fpn": fpn_result,
        "sliding_window": sw_result,
        "comparison": comparison,
    }
    
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    with open(output_path, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\n✅ 结果已保存至: {output_path}")
    print(f"{'=' * 80}\n")
    
    # 退出状态码
    if comparison["meets_speedup_target"] and comparison["meets_memory_target"]:
        print("🎉 所有性能指标均达到预期目标！")
        return 0
    elif comparison["fpn_faster"]:
        print("✅ FPN 性能优于滑窗，但未完全达到目标。")
        return 1
    else:
        print("⚠️ FPN 性能未优于滑窗，需要优化。")
        return 2


if __name__ == "__main__":
    sys.exit(main())