meida_backend/scripts/spikes/rtmpose_benchmark.py

"""T9 SPIKE · RTMPose ONNX CPU 单张耗时基准（T-P0-16 验收口径）。

- 输入：scripts/spikes/assets/synthetic_person.png（synthetic.py 生成的合成人形，
  不涉及真实人脸/人体数据）；
- 口径：预热 3 + 正式 20，报「纯模型 session.run」与「预处理+推理+解码」两条链路
  的均值 / P95 / 单次最大（time.perf_counter）；
- GPU 档：先探测 nvidia-smi 与 ORT providers，CUDA EP 不可用则只报 CPU（不折腾）；
- 降级：models/ 下无 RTMPose ONNX 时，用 cv2.HOG 人体检测作关键点占位并在输出声明
  （HOG 无关键点，仅给检测框 + 按人体比例推的占位点，报告需注明）；
- 输出：out/rtmpose_bench.json（耗时）与 out/rtmpose_keypoints.json（关键点）。

用法：.venv/bin/python scripts/spikes/rtmpose_benchmark.py
"""

from __future__ import annotations

import json
import sys
from typing import Any

import cv2
import numpy as np
from synthetic import ASSETS_DIR, ensure_assets
from vision_common import (
    COCO17_NAMES,
    IDX_HIP,
    IDX_SHOULDER,
    MODELS_DIR,
    OUT_DIR,
    POSE_CHAIN,
    RTMPoseRunner,
    bench,
    dump_json,
    gpu_probe,
    load_session,
)

# GT 校验容差：合成图人形位置已知，肩/臀线估计偏差在该范围内视为「关键点可用」
LINE_TOLERANCE_PX = 15
CONF_USABLE = 0.30  # 关键点平均置信度低于该值视为「定位不可信」（软指标，仅提示）


def _print_gpu_verdict(gpu: dict[str, Any]) -> None:
    print("== GPU 探测 ==")
    print(f"  nvidia-smi: {gpu['nvidia_smi'] or '不可用'}")
    print(f"  ORT providers: {gpu['ort_providers']}")
    if gpu["cuda_ep_available"]:
        print("  结论：CUDA EP 可用（但本 spike 约定只装 CPU 版 ORT，仅报 CPU 档）")
    else:
        print("  结论：CUDA EP 不可用 → 只测 CPU 档（任务文档：探测失败不折腾）")


def _print_keypoints(name: str, kps: np.ndarray) -> None:
    print(f"\n== {name} 关键点（画布坐标 {kps.shape[0]} 点）==")
    print(f"  {'idx':>3}  {'name':<14} {'x':>6} {'y':>6} {'conf':>6}")
    for i, (x, y, c) in enumerate(kps):
        kp_name = COCO17_NAMES[i] if i < len(COCO17_NAMES) else f"kp{i}"
        print(f"  {i:>3}  {kp_name:<14} {x:>6.1f} {y:>6.1f} {c:>6.3f}")


def _check_against_gt(name: str, kps: np.ndarray, gt_std: dict[str, Any]) -> dict[str, Any]:
    """合成图真值校验：肩线/臀线高度与置信度（RTMPose 是否真的「看懂」了人形）。"""
    gt_lines: dict[str, int] = gt_std["line_y"]
    y_sh = float(np.mean(kps[list(IDX_SHOULDER), 1]))
    y_hip = float(np.mean(kps[list(IDX_HIP), 1]))
    d_sh = y_sh - gt_lines["shoulder"]
    d_hip = y_hip - gt_lines["hip"]
    conf_mean = float(np.mean(kps[:, 2]))
    usable = (
        abs(d_sh) <= LINE_TOLERANCE_PX
        and abs(d_hip) <= LINE_TOLERANCE_PX
        and conf_mean >= CONF_USABLE
    )
    print(
        f"  [GT 校验] 肩线偏差 {d_sh:+.1f}px（真值 {gt_lines['shoulder']}），"
        f"臀线偏差 {d_hip:+.1f}px（真值 {gt_lines['hip']}），平均置信度 {conf_mean:.3f}"
        f" → {'可用' if usable else '定位不可信（降级口径见 REPORT）'}"
    )
    return {
        "shoulder_line_delta_px": round(d_sh, 1),
        "hip_line_delta_px": round(d_hip, 1),
        "conf_mean": round(conf_mean, 3),
        "keypoints_usable": usable,
    }


def _bench_model(runner: RTMPoseRunner, tensor: np.ndarray, img: np.ndarray) -> dict[str, Any]:
    def fn_model() -> None:
        runner.session.run(None, {runner.input_name: tensor})

    def fn_full() -> None:
        runner.infer(img)

    return {
        "model_only": bench(fn_model, warmup=3, runs=20),
        "full_pipeline": bench(fn_full, warmup=3, runs=20),
    }


def _hog_fallback(img: np.ndarray) -> int:
    """无 RTMPose 时的降级路径：HOG 人体检测作关键点占位（报告必须注明此口径）。"""
    print("\n== 降级路径：cv2.HOG 人体检测（关键点占位，非 RTMPose）==")
    hog = cv2.HOGDescriptor()
    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())

    def fn() -> None:
        hog.detectMultiScale(img, winStride=(8, 8), padding=(4, 4), scale=1.05)

    stats = bench(fn, warmup=2, runs=10)
    found, _weights = hog.detectMultiScale(img, winStride=(8, 8), padding=(4, 4), scale=1.05)
    if len(found):
        x, y, w, h = found[0]
        print(f"  检测到人体框: x={x} y={y} w={w} h={h}；耗时 {stats}")
        # 占位关键点：按人体比例从检测框推（HOG 不输出关键点，仅维持管线可跑通）
        placeholder = {
            "left_shoulder": [x + 0.2 * w, y + 0.25 * h],
            "right_shoulder": [x + 0.8 * w, y + 0.25 * h],
            "left_hip": [x + 0.3 * w, y + 0.55 * h],
            "right_hip": [x + 0.7 * w, y + 0.55 * h],
        }
    else:
        print("  未检测到人体：占位关键点不可用")
        placeholder = {}
    dump_json(
        OUT_DIR / "rtmpose_keypoints.json",
        {
            "source": "hog-placeholder",
            "detections": [list(map(int, b)) for b in found],
            "placeholder": placeholder,
        },
    )
    dump_json(OUT_DIR / "rtmpose_bench.json", {"source": "hog-placeholder", "latency": stats})
    return 0


def main() -> int:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    ensure_assets()
    img = cv2.imread(str(ASSETS_DIR / "synthetic_person.png"))
    gt = json.loads((ASSETS_DIR / "ground_truth.json").read_text(encoding="utf-8"))

    gpu = gpu_probe()
    _print_gpu_verdict(gpu)

    available = [name for name in POSE_CHAIN if (MODELS_DIR / name).exists()]
    if not available:
        print("[降级] models/ 下无 RTMPose ONNX，改用 cv2.HOG 作关键点占位（见 REPORT）")
        return _hog_fallback(img)

    results: dict[str, Any] = {
        "gpu": gpu,
        "input": "synthetic_person.png (384x512 合成人形)",
        "models": {},
    }
    keypoints_dump: dict[str, Any] = {}
    for name in available:
        path = MODELS_DIR / name
        print(f"\n==== {name}（{path.stat().st_size / 1048576:.1f} MB，CPUExecutionProvider）====")
        runner = RTMPoseRunner(load_session(path))
        print(f"  输入: {runner.input_name} {runner.input_hw}(H,W)")
        kps = runner.infer(img)
        _print_keypoints(name, kps)
        gt_check = _check_against_gt(name, kps, gt["variants"]["STD"])
        stats = _bench_model(runner, runner.preprocess(img), img)
        for lane, s in stats.items():
            print(
                f"  [{lane}] runs={s['runs']} 均值={s['mean_s'] * 1000:.1f}ms "
                f"P95={s['p95_s'] * 1000:.1f}ms 最大={s['max_s'] * 1000:.1f}ms"
            )
        results["models"][name] = {
            "input_hw": list(runner.input_hw),
            "num_keypoints": int(kps.shape[0]),
            "gt_check": gt_check,
            "latency": stats,
        }
        keypoints_dump[name] = {
            "names": list(COCO17_NAMES)[: kps.shape[0]],
            "keypoints": [
                [round(float(x), 1), round(float(y), 1), round(float(c), 3)] for x, y, c in kps
            ],
        }

    dump_json(OUT_DIR / "rtmpose_bench.json", results)
    dump_json(OUT_DIR / "rtmpose_keypoints.json", keypoints_dump)
    print(f"\n结果已写入 {OUT_DIR / 'rtmpose_bench.json'} 与 {OUT_DIR / 'rtmpose_keypoints.json'}")
    return 0


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