xsl
446d7a9391
T-P0-16/12/10 可行性验证：CPU 基准 + 代理质检指标 + 合成图体型剖面法 原型 + REPORT.md 三问三答（Sprint 1 门槛的决策依据）。 偏差说明：模型齐全未触发降级（rtmpose-s/m + birefnet_lite fp32 全部首选 候选成功，hf-mirror 实测 4~12MB/s）；huggingface_hub 不在 spike extra 且 禁新增依赖，改用 stdlib urllib 直连 resolve 端点（同 HTTP 语义）；PINTO model zoo 无法脚本化定位小模型、dhkim0115 仓不存在，姿态模型改用 bukuroo/RTMPose-ONNX；BiRefNet 官方仓无 ONNX，用 onnx-community 导出； GPU 档未测（CPU 版 ORT 无 CUDA EP，按文档不折腾）；BiRefNet 单张>3s 正式 次数自适应降为 10 并在输出注明；RTMPose 在合成剪影上置信度低（域差异）， 体型断言按文档允许走真值定高路径，关键点路径仅统计参考一致率。
170 lines
6.5 KiB
Python
170 lines
6.5 KiB
Python
"""T9 SPIKE · 共用小工具：基准统计、ORT 会话加载、RTMPose 推理解码封装。
|
||
|
||
只服务 scripts/spikes/ 下的四个脚本，不属于生产代码路径；网络访问仅存在于
|
||
download_models.py（本批唯一允许联网的环节），本模块与基准脚本全部离线可跑。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import subprocess
|
||
import time
|
||
from collections.abc import Callable
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
|
||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||
MODELS_DIR = REPO_ROOT / "models"
|
||
ASSETS_DIR = Path(__file__).resolve().parent / "assets"
|
||
OUT_DIR = Path(__file__).resolve().parent / "out"
|
||
|
||
# 模型候选链：与 download_models.py 的 TARGETS 保持一致；全缺失时调用方走降级（HOG/Otsu）
|
||
POSE_CHAIN: tuple[str, ...] = ("rtmpose-s.onnx", "rtmpose-m.onnx")
|
||
MATTING_CHAIN: tuple[str, ...] = ("birefnet_lite.onnx", "birefnet_lite_fp16.onnx", "rmbg14.onnx")
|
||
|
||
# RTMPose（mmpose 部署口径）：RGB 0~255 图像直接减均值除标准差
|
||
MMPOSE_MEAN = np.array([123.675, 116.28, 103.53], dtype=np.float32)
|
||
MMPOSE_STD = np.array([58.395, 57.12, 57.375], dtype=np.float32)
|
||
|
||
COCO17_NAMES = (
|
||
"nose",
|
||
"left_eye",
|
||
"right_eye",
|
||
"left_ear",
|
||
"right_ear",
|
||
"left_shoulder",
|
||
"right_shoulder",
|
||
"left_elbow",
|
||
"right_elbow",
|
||
"left_wrist",
|
||
"right_wrist",
|
||
"left_hip",
|
||
"right_hip",
|
||
"left_knee",
|
||
"right_knee",
|
||
"left_ankle",
|
||
"right_ankle",
|
||
)
|
||
# COCO-17 索引（heights_from_keypoints / GT 校验用）
|
||
IDX_SHOULDER = (5, 6)
|
||
IDX_HIP = (11, 12)
|
||
|
||
|
||
def bench(fn: Callable[[], Any], warmup: int = 3, runs: int = 20) -> dict[str, Any]:
|
||
"""预热 warmup 次后正式 runs 次，返回均值/P95/单次最大（秒，perf_counter 口径）。
|
||
|
||
T-P0-16 的验收口径就是这里的输出：预热 3 + 正式 20。
|
||
"""
|
||
for _ in range(warmup):
|
||
fn()
|
||
samples: list[float] = []
|
||
for _ in range(runs):
|
||
t0 = time.perf_counter()
|
||
fn()
|
||
samples.append(time.perf_counter() - t0)
|
||
arr = np.asarray(samples, dtype=np.float64)
|
||
return {
|
||
"runs": runs,
|
||
"mean_s": round(float(arr.mean()), 4),
|
||
"p95_s": round(float(np.percentile(arr, 95)), 4),
|
||
"max_s": round(float(arr.max()), 4),
|
||
}
|
||
|
||
|
||
def load_session(path: Path) -> ort.InferenceSession:
|
||
"""CPU EP 加载 ONNX。spike 约定不装 GPU 版 ORT（CUDA 匹配成本高，见任务文档 §2）。"""
|
||
return ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
|
||
|
||
|
||
def gpu_probe() -> dict[str, Any]:
|
||
"""探测 GPU 档可用性。任务文档：探测失败不折腾，只报 CPU。"""
|
||
info: dict[str, Any] = {"ort_providers": list(ort.get_available_providers())}
|
||
try:
|
||
proc = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader"],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=10,
|
||
check=True,
|
||
)
|
||
info["nvidia_smi"] = proc.stdout.strip()
|
||
except (OSError, subprocess.SubprocessError):
|
||
info["nvidia_smi"] = None
|
||
info["cuda_ep_available"] = "CUDAExecutionProvider" in info["ort_providers"]
|
||
return info
|
||
|
||
|
||
def dump_json(path: Path, payload: Any) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_text(json.dumps(payload, indent=2, ensure_ascii=False), encoding="utf-8")
|
||
|
||
|
||
class RTMPoseRunner:
|
||
"""RTMPose ONNX（SimCC 输出）最小封装：整图 resize → 推理 → 关键点画布坐标。
|
||
|
||
SPIKE 口径说明：自顶向下姿态模型本应由检测器给人体框再裁剪推理；合成图人形
|
||
占满画布且位置已知，故直接整图缩放（等价于"完美检测框"），把变量集中在耗时
|
||
测量上。输出 (K,3)：画布坐标 x、y + 置信度（x/y 两方向 softmax 峰值的较小者）。
|
||
"""
|
||
|
||
def __init__(self, session: ort.InferenceSession) -> None:
|
||
self.session = session
|
||
inp = session.get_inputs()[0]
|
||
self.input_name = inp.name
|
||
# 动态维度（字符串）按 1 处理，仅用于确定输入尺寸
|
||
shape = [d if isinstance(d, int) else 1 for d in inp.shape]
|
||
if len(shape) != 4:
|
||
msg = f"非预期的 RTMPose 输入维度: {inp.shape}"
|
||
raise ValueError(msg)
|
||
self.input_hw = (shape[2], shape[3]) # (H, W)
|
||
|
||
def preprocess(self, img_bgr: np.ndarray) -> np.ndarray:
|
||
h, w = self.input_hw
|
||
resized = cv2.resize(img_bgr, (w, h), interpolation=cv2.INTER_LINEAR)
|
||
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32)
|
||
rgb = (rgb - MMPOSE_MEAN) / MMPOSE_STD
|
||
chw = np.transpose(rgb, (2, 0, 1))
|
||
return np.ascontiguousarray(chw[None])
|
||
|
||
def infer(self, img_bgr: np.ndarray) -> np.ndarray:
|
||
"""完整链路：预处理 + session.run + SimCC 解码 → (K,3)。"""
|
||
tensor = self.preprocess(img_bgr)
|
||
outputs = self.session.run(None, {self.input_name: tensor})
|
||
return self._decode(outputs, orig_w=img_bgr.shape[1], orig_h=img_bgr.shape[0])
|
||
|
||
def _decode(self, outputs: list[np.ndarray], orig_w: int, orig_h: int) -> np.ndarray:
|
||
# SimCC 输出为两个 (1,K,Nx)/(1,K,Ny) 一维分布；Nx∝输入宽、Ny∝输入高
|
||
cand = [o for o in outputs if o.ndim == 3]
|
||
if len(cand) < 2:
|
||
msg = f"非 SimCC 输出，形状: {[o.shape for o in outputs]}"
|
||
raise ValueError(msg)
|
||
in_h, in_w = self.input_hw
|
||
a, b = cand[0], cand[1]
|
||
# 与输入宽成比例的是 x 分布（比较两种配对的残差，取小者）
|
||
fit_ab = abs(a.shape[2] / in_w - b.shape[2] / in_h)
|
||
fit_ba = abs(b.shape[2] / in_w - a.shape[2] / in_h)
|
||
if fit_ba < fit_ab:
|
||
a, b = b, a
|
||
x_simcc, y_simcc = a[0], b[0] # (K,Nx) / (K,Ny)
|
||
scale_x = a.shape[2] / in_w
|
||
scale_y = b.shape[2] / in_h
|
||
xs = x_simcc.argmax(axis=1).astype(np.float32) / scale_x
|
||
ys = y_simcc.argmax(axis=1).astype(np.float32) / scale_y
|
||
conf = np.minimum(_peak_prob(x_simcc), _peak_prob(y_simcc))
|
||
kps = np.zeros((xs.shape[0], 3), dtype=np.float32)
|
||
kps[:, 0] = xs * (orig_w / in_w)
|
||
kps[:, 1] = ys * (orig_h / in_h)
|
||
kps[:, 2] = conf
|
||
return kps
|
||
|
||
|
||
def _peak_prob(simcc: np.ndarray) -> np.ndarray:
|
||
"""每个关键点一维分布的 softmax 峰值概率，作为置信度代理。"""
|
||
e = np.exp(simcc - simcc.max(axis=1, keepdims=True))
|
||
probs = e / e.sum(axis=1, keepdims=True)
|
||
return probs.max(axis=1)
|