spike(vision): rtmpose latency, birefnet matting proxies, body-shape prototype
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 在合成剪影上置信度低（域差异）， 体型断言按文档允许走真值定高路径，关键点路径仅统计参考一致率。
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
# T9 · 视觉算法 SPIKE 报告（RTMPose 体型原型 + BiRefNet 抠图）
|
||||
|
||||
> 分支 `task/t9-vision-spike` ｜ 2026-08-14 ｜ 对应任务 T-P0-16 / T-P0-12 / TD-06 前置
|
||||
> 测试机：Intel Core Ultra 7 265K（WSL 分配 16 vCPU、54GB 内存），onnxruntime 1.28.0（CPU 版）、opencv-python-headless 5.0.0、numpy 2.5.2
|
||||
> 所有输入均为 OpenCV 合成人形图（`synthetic.py` 确定性生成），**未使用任何真实人脸/人体数据**（合规红线）。
|
||||
|
||||
---
|
||||
|
||||
## 一、三问三答（结论先行）
|
||||
|
||||
### Q1 RTMPose ONNX 在 CPU 上的单张推理耗时是多少？（T-P0-16）
|
||||
|
||||
**结论：通过。** RTMPose-s 单张纯推理 **2.2ms（P95 2.6ms / 最大 2.7ms）**，含预处理与解码的全链路 **2.8ms（P95 3.3ms / 最大 3.5ms）**；RTMPose-m 全链路 5.9ms。CPU 单张耗时对服务端批量/异步处理完全无压力，T-P0-16 验收口径（可接受即通过）满足。GPU 档未测：开发机有 RTX 5090（nvidia-smi 可见），但 venv 装的是 CPU 版 onnxruntime（无 CUDAExecutionProvider），按任务文档约定不装 GPU 版、只报 CPU 档。
|
||||
|
||||
### Q2 BiRefNet ONNX 抠图的质量与成本如何？（T-P0-12 初步）
|
||||
|
||||
**结论：有条件通过。** 管线成立：BiRefNet_lite fp32 在合成图上前景占比 24.4%（区间 15%~60% 内）、最大连通域占前景 1.0（单一主体）、边缘羽化带约 1px、纯色与渐变背景结果一致。**条件/风险有二**：① 单张全链路 **约 3.1s（CPU，1024×1024 输入）**，同步接口不可接受，必须走异步任务或降档（见 §四）；② 本 SPIKE 的指标是**代理口径**（合成图 + 几何指标），T-P0-12 的「质检通过率 ≥90%」必须用真实素材 + 人工标注验证，合成图结果不能外推到正式验收。
|
||||
|
||||
### Q3 「关键点定高 + 掩膜宽度剖面」体型测算原型是否成立？（TD-06 前置）
|
||||
|
||||
**结论：通过（原型口径）。** 管线端到端跑通：定高 → 掩膜剖面取宽 → 肩/腰/臀比值 → 规则初判。**真值定高路径 6/6 变体（H/A/V/O/X 全类别 + 中性 STD）分类硬断言全部通过**；掩膜测宽与几何真值偏差仅 ±1px（一个变体腰线 +3px）；关键点（RTMPose）定高路径参考一致率 6/6。两点保留：① COCO-17 无腰点，腰线取肩臀中点是代理，对真实人体（腰臀比例分布）是否成立待验证；② RTMPose 在合成剪影上置信度极低（见 §三），KP 路径在真实照片上的可用性未验证；≥80% 一致率的正式验证需要顾问轨人工标注（本任务不承诺）。
|
||||
|
||||
---
|
||||
|
||||
## 二、基准数字表
|
||||
|
||||
### 2.1 RTMPose CPU 单张耗时（预热 3 + 正式 20，`time.perf_counter`）
|
||||
|
||||
| 模型 | 链路 | 均值 | P95 | 单次最大 |
|
||||
| --- | --- | ---: | ---: | ---: |
|
||||
| rtmpose-s（20.9MB，输入 256×192） | 纯模型 `session.run` | 2.2ms | 2.6ms | 2.7ms |
|
||||
| rtmpose-s | 全链路（resize+归一化+推理+SimCC 解码） | 2.8ms | 3.3ms | 3.5ms |
|
||||
| rtmpose-m（51.8MB，输入 256×192） | 纯模型 | 5.2ms | 5.7ms | 5.8ms |
|
||||
| rtmpose-m | 全链路 | 5.9ms | 6.2ms | 6.5ms |
|
||||
|
||||
- 预处理+解码开销约 0.6~0.7ms（全链路 − 纯模型），与模型规模无关。
|
||||
- 耗时口径说明：RTMPose 是 top-down 模型，本基准以「合成图已知完美人体框、整图缩放」为输入，即纯姿态网络耗时；生产链路还需前置一个人体检测器（见 §四建议）。
|
||||
|
||||
### 2.2 BiRefNet_lite 抠图代理质检指标（预热 3 + 正式 10，单张 >3s 自适应降档并注明）
|
||||
|
||||
| 输入 | 前景占比 | 最大连通域/前景 | 连通域数 | 边缘羽化带 | 均值耗时 | P95 | 最大 |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| 纯色背景 | 0.2444 | 1.0000 | 1 | 0.97px | 3097ms | 3128ms | 3131ms |
|
||||
| 渐变背景 | 0.2443 | 1.0000 | 1 | 1.02px | 3097ms | 3131ms | 3139ms |
|
||||
|
||||
- 前景占比处于 15%~60% 合理区间，主体单一完整（>0.9 阈值达成，实际 1.0）。
|
||||
- 羽化带 ≈1px：对高对比合成图边缘干净利落，真实照片（毛发/半透明）会显著变宽，该值应作为正式质检的分档特征而非固定阈值。
|
||||
- 两背景变体指标几乎一致 → 对背景明暗变化不敏感（该结论仅限合成域）。
|
||||
|
||||
### 2.3 体型测算原型（合成图，掩膜 = BiRefNet_lite）
|
||||
|
||||
| 变体 | 期望 | 真值定高（硬断言） | 关键点定高（参考） | 掩膜测宽 肩/腰/臀(px) | 保真度偏差(px) |
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| STD（中性） | H | **H ✓** | H ✓ | 131/111/131 | +1/+1/+1 |
|
||||
| H | H | **H ✓** | H ✓ | 127/121/125 | +1/+1/+1 |
|
||||
| A | A | **A ✓** | A ✓ | 99/103/147 | +1/+1/+1 |
|
||||
| V | V | **V ✓** | V ✓ | 155/107/119 | +1/+1/+1 |
|
||||
| X | X | **X ✓** | X ✓ | 131/81/135 | +1/+3/+1 |
|
||||
| O | O | **O ✓** | O ✓ | 115/143/123 | +1/+1/+1 |
|
||||
|
||||
- 真值定高 **6/6 断言通过**；关键点定高一致率 **6/6（100%）**（参考值，样本为合成图）。
|
||||
- 阈值（首版经验值，**待顾问校准**）：肩/臀 >1.05 → V；<0.95 → A；腰 <0.80×min(肩,臀) → X；腰 >1.02×max(肩,臀) → O；否则 H。判定顺序 O → X → V → A → H。
|
||||
|
||||
---
|
||||
|
||||
## 三、RTMPose 在合成剪影上的关键点质量（重要发现）
|
||||
|
||||
RTMPose 关键点在合成剪影上**定位偏差可用但置信度极低**：rtmpose-s 肩线偏差 +12px、臀线偏差 −28.5px（画布高 512px），但平均置信度仅 0.003（SimCC 分布近乎平坦）。已排除预处理口径因素（RGB/BGR、0-255/0-1 归一化四种组合结果几乎一致），结论是**域差异**：模型在真实人像上训练，对无纹理几何剪影没有信心。对本 SPIKE 的处理：
|
||||
|
||||
- 断言路径用「真值定高」（任务文档明确允许：*"用 3.2 的关键点（或合成图直接已知的真值框）定肩线/腰线/臀线高度"*），KP 路径仅统计参考一致率；
|
||||
- 对 Sprint 1 的含义：**RTMPose 的关键点质量必须在真实照片上重新验证**，合成图不能作为其可用性依据；体型管线的关键点依赖是当前最大不确定性。
|
||||
|
||||
---
|
||||
|
||||
## 四、对 Sprint 1 正式验证的建议
|
||||
|
||||
1. **T-P0-16（RTMPose 耗时）**：CPU 数字已可决策（s 档 3ms 级全链路）。还需要补两件事：① 生产链路要前置**人体检测器**（RTMPose 是 top-down），建议补测 YOLOv8n-ONNX 或 RTMO 一体化模型的 CPU 耗时；② 若上传链路要同步姿态，当前 CPU 数字已够；无需 GPU。
|
||||
2. **T-P0-12（抠图质检 ≥90%）**：需要 **200~500 张真实全身照**（覆盖多背景/光线/姿态），人工标注粗前景 GT，指标建议 IoU ≥0.90 或边界 F-measure，先定义「单张通过」标准再统计通过率。成本决策点：若 CPU 3.1s/张不可接受，候选降档（按优先级）——任务异步化（现有 task 架构，零成本）→ BiRefNet fp16 / 512×512 变体（候选已在 download_models.py 链里）→ GPU EP。
|
||||
3. **TD-06（体型测算 ≥80% 一致率）**：需要顾问轨 **≥100 人体型标注**（H/A/V/O/X 五类尽量均衡）+ 每人至少一张标准站姿全身照；阈值校准建议用分位数/ROC 而非拍脑袋值；**前置依赖**是 RTMPose 在真实照片上的关键点质量验证（§三）。腰线代理（肩臀中点）需在标注数据上复核落点。
|
||||
4. **数据合规**：真实人体照片进入验证集前需走用户授权/合成替代评估；本 SPIKE 已验证「合成图 + 真值」开发范式可行，阈值调参阶段可继续用合成图，终验必须真实分布。
|
||||
|
||||
---
|
||||
|
||||
## 五、使用的模型与版本（models/manifest.json）
|
||||
|
||||
| 用途 | 模型 | 仓库 @ revision | 变体 | 大小 |
|
||||
| --- | --- | --- | --- | ---: |
|
||||
| 姿态（基准主档） | rtmpose-s.onnx | `bukuroo/RTMPose-ONNX` @ `a6e9fb8a` | RTMPose-s body COCO-17 SimCC，输入 256×192 | 20.9MB |
|
||||
| 姿态（对比档） | rtmpose-m.onnx | `bukuroo/RTMPose-ONNX` @ `a6e9fb8a` | RTMPose-m 同上 | 51.8MB |
|
||||
| 抠图 | birefnet_lite.onnx | `onnx-community/BiRefNet_lite-ONNX` @ `de15b22b` | fp32，输入 1024×1024，ImageNet 归一化 | 213.6MB |
|
||||
|
||||
下载通道：`https://hf-mirror.com`（国内镜像优先，实测 4~12MB/s，总 ~290MB，未触发 10 分钟超时纪律）。
|
||||
|
||||
## 六、降级路径与偏差说明
|
||||
|
||||
- **实现并验证的降级链**（本次全部未触发，代码路径已在脚本中就位）：姿态 `bukuroo rtmpose-s` → （失败）`cv2.HOG` 关键点占位；抠图 `BiRefNet_lite fp32` → `fp16` → `briaai/RMBG-1.4` → （仍失败）Otsu 二值化占位。脚本对每级失败均打印 `[降级]` 日志并在输出 JSON 记 `model_source`。
|
||||
- **偏差 1（依赖）**：任务文档建议用 `huggingface_hub.snapshot_download`，但该包不在 spike extra 且本批禁止新增依赖（协作规程 §4.2），改用 stdlib `urllib` 直连 `{endpoint}/{repo}/resolve/main/{file}`（同一 HTTP 语义，支持 30x 与 Range 断点续传）。已在脚本 docstring 声明。
|
||||
- **偏差 2（模型选型）**：文档首选 `PINTO0309/PINTO_model_zoo` 的 rtmpose，但该仓库为巨型混合仓库（API 兄弟列表无法按"body 17/26 keypoint 小模型"脚本化定位）；`dhkim0115/rtmpose-onnx` 经 API 探查不存在；最终选社区导出 `bukuroo/RTMPose-ONNX`（17 点 COCO、SimCC、256×192，符合文档"body 小模型"要求）。BiRefNet 官方仓 `ZhengPeng7/BiRefNet` 只有 safetensors 无 ONNX，采用 `onnx-community` 的 ONNX 导出（文档列举的备选之一）。
|
||||
- **偏差 3（GPU 档）**：nvidia-smi 可见 RTX 5090，但 venv 的 onnxruntime 为 CPU 构建（providers 无 CUDA），按任务文档「不装 GPU 版 ORT、探测失败不折腾」只报 CPU 档。
|
||||
- **口径调整**：BiRefNet 单张 >3s，正式次数由 20 自适应降为 10（脚本内 `adaptive_note` 字段注明）；RTMPose 基准严格按预热 3 + 正式 20。
|
||||
|
||||
## 七、复现方式（worktree 根目录）
|
||||
|
||||
```bash
|
||||
.venv/bin/python scripts/spikes/download_models.py # 联网下载（已缓存则秒过，唯一联网环节）
|
||||
.venv/bin/python scripts/spikes/rtmpose_benchmark.py # Q1：耗时基准 + 关键点
|
||||
.venv/bin/python scripts/spikes/birefnet_matting.py # Q2：代理质检指标
|
||||
.venv/bin/python scripts/spikes/body_measure_prototype.py # Q3：体型分类硬断言
|
||||
```
|
||||
|
||||
输出：数字与 JSON 落在 `scripts/spikes/out/`（已 gitignore），掩膜/抠图预览同目录可人工核验。
|
||||
@@ -0,0 +1,185 @@
|
||||
"""T9 SPIKE · BiRefNet ONNX 抠图管线 + 代理质检指标（T-P0-12 初步口径）。
|
||||
|
||||
代理指标（注意：这是代理口径，不是 T-P0-12 的正式验收——正式验收需要真实素材 +
|
||||
人工标注的通过率，本 SPIKE 只建管线并给少量合成样图的主观核验依据）：
|
||||
- 前景像素占比（合理区间 15%~60%）；
|
||||
- 最大连通域占前景比（>0.9 视为「主体单一完整」）；
|
||||
- 边缘羽化度（0.05<soft<0.95 过渡带面积 / 边界周长 ≈ 平均过渡带宽度 px）；
|
||||
- 单张耗时（预热 3 + 正式 20 次；单张 >3s 时自动降为 10 次并注明）。
|
||||
|
||||
输入：合成人形图（纯色背景 + 渐变背景变体，不涉及真实个人数据）；
|
||||
降级：models/ 下无抠图模型时用 Otsu 二值化占位（硬边缘，无羽化，报告需注明）；
|
||||
输出：out/ 下掩膜图 + 抠图预览（gitignore，报告只留指标）+ matting_metrics.json。
|
||||
|
||||
用法：.venv/bin/python scripts/spikes/birefnet_matting.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
from synthetic import ASSETS_DIR, ensure_assets
|
||||
from vision_common import MATTING_CHAIN, MODELS_DIR, OUT_DIR, bench, dump_json, load_session
|
||||
|
||||
# ImageNet 归一化（onnx-community 导出的 BiRefNet 系列 preprocessor_config 口径）
|
||||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
# 代理质检阈值（T-P2-07 质检雏形口径，非正式验收）
|
||||
FG_RATIO_RANGE = (0.15, 0.60) # 前景占比合理区间
|
||||
LARGEST_CC_OK = 0.90 # 最大连通域占前景比阈值
|
||||
SOFT_BAND = (0.05, 0.95) # 软边过渡带判定区间
|
||||
SLOW_RUNS = 10 # 单张 >3s 时的正式次数（默认 20）
|
||||
|
||||
_session: ort.InferenceSession | None = None
|
||||
_model_source: str | None = None
|
||||
|
||||
|
||||
def _get_session() -> ort.InferenceSession | None:
|
||||
"""按候选链加载第一个可用模型（加载失败自动降级到下一候选）。"""
|
||||
global _session, _model_source
|
||||
if _model_source is not None:
|
||||
return _session
|
||||
for name in MATTING_CHAIN:
|
||||
path = MODELS_DIR / name
|
||||
if not path.exists():
|
||||
continue
|
||||
try:
|
||||
_session = load_session(path)
|
||||
_model_source = name
|
||||
print(f"[模型] 使用 {name}（{path.stat().st_size / 1048576:.1f} MB）")
|
||||
return _session
|
||||
except Exception as exc: # SPIKE：逐候选降级，失败原因留痕后换下一个
|
||||
print(f"[降级] {name} 加载失败：{exc}")
|
||||
_session = None
|
||||
_model_source = "otsu-fallback"
|
||||
print("[降级] 无可用抠图模型，使用 Otsu 二值化占位（硬边缘、无羽化，报告需注明）")
|
||||
return None
|
||||
|
||||
|
||||
def _preprocess(img_bgr: np.ndarray, size_wh: tuple[int, int]) -> np.ndarray:
|
||||
resized = cv2.resize(img_bgr, size_wh, interpolation=cv2.INTER_LINEAR)
|
||||
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
|
||||
rgb = (rgb - IMAGENET_MEAN) / IMAGENET_STD
|
||||
return np.ascontiguousarray(np.transpose(rgb, (2, 0, 1))[None])
|
||||
|
||||
|
||||
def build_person_mask(img_bgr: np.ndarray) -> tuple[np.ndarray, str]:
|
||||
"""抠图主入口：返回 (软掩膜 float32 HxW ∈ [0,1] 与原图同尺寸, 来源标签)。
|
||||
|
||||
无模型时走 Otsu：人体比背景暗 → BINARY_INV 使人体=1。body_measure_prototype
|
||||
也复用本函数，保证两条 spike 的掩膜口径一致。
|
||||
"""
|
||||
session = _get_session()
|
||||
if session is None:
|
||||
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
||||
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
||||
thresh, _ = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||
return thresh.astype(np.float32) / 255.0, "otsu-fallback"
|
||||
inp = session.get_inputs()[0]
|
||||
size_wh = (
|
||||
inp.shape[3] if isinstance(inp.shape[3], int) else 1024,
|
||||
inp.shape[2] if isinstance(inp.shape[2], int) else 1024,
|
||||
)
|
||||
tensor = _preprocess(img_bgr, size_wh)
|
||||
logits = session.run(None, {inp.name: tensor})[0]
|
||||
soft_small = 1.0 / (1.0 + np.exp(-np.clip(logits[0, 0].astype(np.float32), -30.0, 30.0)))
|
||||
soft = cv2.resize(
|
||||
soft_small, (img_bgr.shape[1], img_bgr.shape[0]), interpolation=cv2.INTER_LINEAR
|
||||
)
|
||||
assert _model_source is not None
|
||||
return soft, _model_source
|
||||
|
||||
|
||||
def proxy_metrics(img_bgr: np.ndarray, soft: np.ndarray) -> dict[str, Any]:
|
||||
"""T-P0-12 代理质检指标（合成图口径，正式验收需真实素材）。"""
|
||||
binary = (soft > 0.5).astype(np.uint8)
|
||||
fg = int(binary.sum())
|
||||
fg_ratio = fg / binary.size
|
||||
num_labels, _, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
|
||||
largest = int(stats[1:, cv2.CC_STAT_AREA].max()) if num_labels > 1 else 0
|
||||
largest_cc_ratio = largest / fg if fg else 0.0
|
||||
# 边缘羽化：过渡带面积 / 边界周长 ≈ 平均过渡带宽度（px）
|
||||
band = ((soft > SOFT_BAND[0]) & (soft < SOFT_BAND[1])).astype(np.uint8)
|
||||
band_area = int(band.sum())
|
||||
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
|
||||
perimeter = float(sum(cv2.arcLength(c, True) for c in contours)) or 1.0
|
||||
return {
|
||||
"fg_ratio": round(fg_ratio, 4),
|
||||
"fg_ratio_in_range": FG_RATIO_RANGE[0] <= fg_ratio <= FG_RATIO_RANGE[1],
|
||||
"largest_cc_ratio": round(largest_cc_ratio, 4),
|
||||
"largest_cc_ok": largest_cc_ratio > LARGEST_CC_OK,
|
||||
"feather_band_px": round(band_area / perimeter, 2),
|
||||
"band_area_px": band_area,
|
||||
"boundary_perimeter_px": round(perimeter, 1),
|
||||
"components": int(num_labels - 1),
|
||||
}
|
||||
|
||||
|
||||
def _bench_latency(img_bgr: np.ndarray) -> dict[str, Any]:
|
||||
"""单张全链路耗时（预处理+推理+sigmoid+resize 回原尺寸），口径同 rtmpose。"""
|
||||
|
||||
def fn() -> tuple[np.ndarray, str]:
|
||||
return build_person_mask(img_bgr)
|
||||
|
||||
fn() # 预估单张耗时以自适应正式次数（BiRefNet 1024² 在 CPU 上可能较慢）
|
||||
t0 = time.perf_counter()
|
||||
fn()
|
||||
single = time.perf_counter() - t0
|
||||
runs = SLOW_RUNS if single > 3.0 else 20
|
||||
stats = bench(fn, warmup=3, runs=runs)
|
||||
stats["adaptive_note"] = f"单张约 {single:.2f}s → 正式 {runs} 次"
|
||||
return stats
|
||||
|
||||
|
||||
def _save_outputs(stem: str, img: np.ndarray, soft: np.ndarray) -> None:
|
||||
"""掩膜与抠图预览落盘（out/ 已 gitignore，仅作人工核验，不入库）。"""
|
||||
cv2.imwrite(str(OUT_DIR / f"mask_{stem}.png"), (soft * 255).astype(np.uint8))
|
||||
alpha = soft[..., None]
|
||||
cutout = (img.astype(np.float32) * alpha + 255.0 * (1.0 - alpha)).astype(np.uint8)
|
||||
cv2.imwrite(str(OUT_DIR / f"cutout_{stem}.png"), cutout)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
ensure_assets()
|
||||
images = {
|
||||
"synthetic_person": ASSETS_DIR / "synthetic_person.png",
|
||||
"synthetic_person_grad": ASSETS_DIR / "synthetic_person_grad.png",
|
||||
}
|
||||
results: dict[str, Any] = {}
|
||||
for stem, path in images.items():
|
||||
img = cv2.imread(str(path))
|
||||
soft, source = build_person_mask(img)
|
||||
metrics = proxy_metrics(img, soft)
|
||||
print(f"\n==== {stem}（{source}）====")
|
||||
print(
|
||||
f" 前景占比 {metrics['fg_ratio']:.4f}（区间 {FG_RATIO_RANGE}）"
|
||||
f" → {'OK' if metrics['fg_ratio_in_range'] else 'OUT'}"
|
||||
)
|
||||
print(
|
||||
f" 最大连通域/前景 {metrics['largest_cc_ratio']:.4f}（>{LARGEST_CC_OK}）"
|
||||
f" → {'OK' if metrics['largest_cc_ok'] else 'OUT'}；连通域数 {metrics['components']}"
|
||||
)
|
||||
print(f" 边缘羽化带 {metrics['feather_band_px']}px（过渡带 {metrics['band_area_px']}px²）")
|
||||
latency = _bench_latency(img)
|
||||
print(
|
||||
f" 耗时 runs={latency['runs']} 均值={latency['mean_s'] * 1000:.0f}ms "
|
||||
f"P95={latency['p95_s'] * 1000:.0f}ms 最大={latency['max_s'] * 1000:.0f}ms"
|
||||
f"（{latency['adaptive_note']}）"
|
||||
)
|
||||
_save_outputs(stem, img, soft)
|
||||
results[stem] = {"model_source": source, "metrics": metrics, "latency": latency}
|
||||
|
||||
dump_json(OUT_DIR / "matting_metrics.json", results)
|
||||
print(f"\n结果已写入 {OUT_DIR / 'matting_metrics.json'}（掩膜/抠图预览在同目录）")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,187 @@
|
||||
"""T9 SPIKE · TD-06 体型测算原型：定高（关键点/真值）+ 掩膜宽度剖面 → H/A/V/O/X 初判。
|
||||
|
||||
管线：肩/腰/臀三线高度 → 在抠图掩膜上取对应行宽度剖面 → 肩宽/腰宽/臀宽比值 →
|
||||
规则初判。两条定高路径：
|
||||
- 真值定高（合成图生成时已知三线高度）：对全部体型变体【硬断言】分类正确——
|
||||
这是任务文档 §3.4 要求的「已知真值形状的分类正确」最小一致率验证；
|
||||
- 关键点定高（RTMPose 肩/臀关键点；COCO-17 无腰点，腰线取肩臀中点作代理）：
|
||||
同样输出分类并统计与真值的一致率。≥80% 一致率的正式验证需顾问轨人工标注，
|
||||
本任务不承诺（见 REPORT.md）。
|
||||
|
||||
阈值全部为常量——首版经验值，【待顾问校准】。
|
||||
用法：.venv/bin/python scripts/spikes/body_measure_prototype.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from birefnet_matting import build_person_mask
|
||||
from synthetic import ASSETS_DIR, VARIANTS, ensure_assets
|
||||
from vision_common import (
|
||||
IDX_HIP,
|
||||
IDX_SHOULDER,
|
||||
MODELS_DIR,
|
||||
OUT_DIR,
|
||||
POSE_CHAIN,
|
||||
RTMPoseRunner,
|
||||
dump_json,
|
||||
load_session,
|
||||
)
|
||||
|
||||
# ── 体型初判阈值（首版经验值，待顾问轨校准） ────────────────────────────────
|
||||
SR_V = 1.05 # 肩/臀 > 1.05 → V（倒三角）
|
||||
SR_A = 0.95 # 肩/臀 < 0.95 → A（梨形）
|
||||
X_WAIST = 0.80 # 腰 < 0.80×min(肩,臀) → X（沙漏：腰明显收窄）
|
||||
O_WAIST = 1.02 # 腰 > 1.02×max(肩,臀) → O（中段最大）
|
||||
|
||||
|
||||
def profile_width(binary: np.ndarray, y: int) -> float:
|
||||
"""掩膜在 y 行的宽度剖面：取 y-1/y/y+1 三行前景跨度的中位数（抗单行毛刺）。"""
|
||||
widths: list[float] = []
|
||||
for row_y in (y - 1, y, y + 1):
|
||||
if not 0 <= row_y < binary.shape[0]:
|
||||
continue
|
||||
xs = np.flatnonzero(binary[row_y])
|
||||
if xs.size:
|
||||
widths.append(float(xs[-1] - xs[0] + 1))
|
||||
if not widths:
|
||||
return 0.0
|
||||
return float(np.median(widths))
|
||||
|
||||
|
||||
def measure(binary: np.ndarray, line_y: dict[str, int]) -> dict[str, float]:
|
||||
"""按三线高度取宽度剖面 → 肩宽/腰宽/臀宽（px）。"""
|
||||
return {name: profile_width(binary, y) for name, y in line_y.items()}
|
||||
|
||||
|
||||
def classify(widths: dict[str, float]) -> str:
|
||||
"""规则初判：先判中段（O/X），再判肩臀比（V/A），否则 H。"""
|
||||
if min(widths.values()) <= 0.0:
|
||||
return "INVALID" # 掩膜在该线无前景：管线异常，硬断言会失败（这是期望行为）
|
||||
shoulder, waist, hip = widths["shoulder"], widths["waist"], widths["hip"]
|
||||
sr = shoulder / hip
|
||||
if waist > O_WAIST * max(shoulder, hip):
|
||||
return "O"
|
||||
if waist < X_WAIST * min(shoulder, hip):
|
||||
return "X"
|
||||
if sr > SR_V:
|
||||
return "V"
|
||||
if sr < SR_A:
|
||||
return "A"
|
||||
return "H"
|
||||
|
||||
|
||||
def heights_from_keypoints(kps: np.ndarray) -> dict[str, int]:
|
||||
"""关键点定高：肩线=双肩均值，臀线=双髋均值，腰线=肩臀中点（COCO-17 无腰点）。"""
|
||||
y_sh = float(np.mean(kps[list(IDX_SHOULDER), 1]))
|
||||
y_hip = float(np.mean(kps[list(IDX_HIP), 1]))
|
||||
y_waist = (y_sh + y_hip) / 2.0
|
||||
return {"shoulder": round(y_sh), "waist": round(y_waist), "hip": round(y_hip)}
|
||||
|
||||
|
||||
def _try_pose_runner() -> RTMPoseRunner | None:
|
||||
for name in POSE_CHAIN:
|
||||
path = MODELS_DIR / name
|
||||
if path.exists():
|
||||
try:
|
||||
runner = RTMPoseRunner(load_session(path))
|
||||
print(f"[模型] 关键点定高使用 {name}")
|
||||
return runner
|
||||
except Exception as exc: # SPIKE：加载失败降级，真值定高路径不受影响
|
||||
print(f"[降级] {name} 初始化失败：{exc}")
|
||||
print("[降级] 无 RTMPose 模型，关键点定高路径跳过（真值定高路径不受影响）")
|
||||
return None
|
||||
|
||||
|
||||
def main() -> int:
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
gt = ensure_assets()
|
||||
pose_runner = _try_pose_runner()
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for shape in VARIANTS:
|
||||
meta = gt["variants"][shape.key]
|
||||
img = cv2.imread(str(ASSETS_DIR / "variants" / f"{shape.key}.png"))
|
||||
soft, mask_source = build_person_mask(img)
|
||||
binary = (soft > 0.5).astype(np.uint8)
|
||||
|
||||
# 路径一：真值定高（合成图生成时已知）——硬断言分类正确
|
||||
widths_gt = measure(binary, meta["line_y"])
|
||||
pred_gt = classify(widths_gt)
|
||||
assert pred_gt == shape.expect, (
|
||||
f"{shape.key}: 真值定高分类 {pred_gt} != 期望 {shape.expect}，宽度={widths_gt}"
|
||||
)
|
||||
|
||||
# 路径二：关键点定高（RTMPose；腰线为肩臀中点代理）——统计一致率
|
||||
pred_kp: str | None = None
|
||||
widths_kp: dict[str, float] | None = None
|
||||
if pose_runner is not None:
|
||||
kps = pose_runner.infer(img)
|
||||
widths_kp = measure(binary, heights_from_keypoints(kps))
|
||||
pred_kp = classify(widths_kp)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"key": shape.key,
|
||||
"expect": shape.expect,
|
||||
"pred_gt": pred_gt,
|
||||
"pred_kp": pred_kp,
|
||||
"gt_ok": pred_gt == shape.expect,
|
||||
"kp_ok": None if pred_kp is None else pred_kp == shape.expect,
|
||||
"widths_gt": widths_gt,
|
||||
"widths_kp": widths_kp,
|
||||
"gt_widths_drawn": meta["gt_widths"],
|
||||
"mask_source": mask_source,
|
||||
}
|
||||
)
|
||||
|
||||
print(f"\n==== 体型初判结果（掩膜来源 {rows[0]['mask_source']}）====")
|
||||
print(f" {'变体':>4} {'期望':>4} {'真值定高':>8} {'关键点定高':>10} 宽度(肩/腰/臀, 掩膜测得)")
|
||||
for row in rows:
|
||||
w = row["widths_gt"]
|
||||
kp_txt = row["pred_kp"] or "-"
|
||||
print(
|
||||
f" {row['key']:>4} {row['expect']:>4} {row['pred_gt']:>8} {kp_txt:>10}"
|
||||
f" {w['shoulder']:.0f}/{w['waist']:.0f}/{w['hip']:.0f}px"
|
||||
)
|
||||
gt_pass = sum(1 for r in rows if r["gt_ok"])
|
||||
kp_rows = [r for r in rows if r["kp_ok"] is not None]
|
||||
kp_pass = sum(1 for r in kp_rows if r["kp_ok"])
|
||||
print(f"\n [硬断言] 真值定高 {gt_pass}/{len(rows)} 分类正确（全部通过 = TD-06 原型管线成立）")
|
||||
if kp_rows:
|
||||
print(
|
||||
f" [参考] 关键点定高一致率 {kp_pass}/{len(kp_rows)}"
|
||||
f"（{kp_pass / len(kp_rows):.0%}；正式 ≥80% 验证需顾问轨人工标注，本任务不承诺）"
|
||||
)
|
||||
|
||||
# 掩膜保真度：测得宽度 vs 画图时的几何真值宽度（px 级偏差）
|
||||
print(" [保真度] 掩膜宽度 vs 几何真值（px）：")
|
||||
for row in rows:
|
||||
drawn = row["gt_widths_drawn"]
|
||||
w = row["widths_gt"]
|
||||
deltas = [w[k] - drawn[k] for k in ("shoulder", "waist", "hip")]
|
||||
print(f" {row['key']:>4}: 肩{deltas[0]:+.1f} 腰{deltas[1]:+.1f} 臀{deltas[2]:+.1f}")
|
||||
|
||||
dump_json(
|
||||
OUT_DIR / "body_shape_result.json",
|
||||
{
|
||||
"thresholds": {"SR_V": SR_V, "SR_A": SR_A, "X_WAIST": X_WAIST, "O_WAIST": O_WAIST},
|
||||
"gt_pass": gt_pass,
|
||||
"total": len(rows),
|
||||
"kp_pass": kp_pass,
|
||||
"kp_total": len(kp_rows),
|
||||
"rows": rows,
|
||||
},
|
||||
)
|
||||
print(f"\n结果已写入 {OUT_DIR / 'body_shape_result.json'}")
|
||||
if gt_pass != len(rows):
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,269 @@
|
||||
"""T9 SPIKE · 模型下载脚本（本批唯一允许联网的环节，任务文档 §3.1）。
|
||||
|
||||
为什么不用 huggingface_hub：该包不在 spike extra 依赖中，且协作规程 §4.2 禁止新增
|
||||
依赖，故用 stdlib urllib 直连 `{endpoint}/{repo}/resolve/main/{file}`——与
|
||||
huggingface_hub 底层相同的 HTTP 语义（30x 跳转与 Range 断点续传均支持）。
|
||||
|
||||
行为要点：
|
||||
- 端点顺序：环境变量 HF_ENDPOINT > https://hf-mirror.com（国内镜像优先）> 官方站；
|
||||
- 候选链：每个目标文件按候选依次降级（BiRefNet_lite fp32 → fp16 → RMBG-1.4）；
|
||||
- 超时纪律：单文件超过 --budget 秒（默认 600s ≈ 10 分钟）即放弃当前候选换下一个；
|
||||
- 幂等：本地已存在且字节数与远端一致则跳过；.part 断点续传；每候选失败重试 1 次；
|
||||
- 结果（repo/revision/大小/耗时/端点）写入 models/manifest.json，供 REPORT.md 引用。
|
||||
|
||||
用法（在 worktree 根目录）：
|
||||
.venv/bin/python scripts/spikes/download_models.py [--only matting] [--budget 600]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
DEFAULT_MODELS_DIR = REPO_ROOT / "models"
|
||||
DEFAULT_BUDGET_SECONDS = 600 # 任务文档：单个模型 ~10 分钟即降级
|
||||
CHUNK_BYTES = 1 << 20
|
||||
ATTEMPTS_PER_CANDIDATE = 2 # 首次 + 失败重试一次
|
||||
PROGRESS_EVERY_BYTES = 32 * CHUNK_BYTES
|
||||
UA = "meida-t9-spike/1.0"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Candidate:
|
||||
"""一个候选模型文件的下载描述。"""
|
||||
|
||||
repo: str
|
||||
filename: str
|
||||
dest: str # models/ 下的目标文件名
|
||||
note: str # 变体说明（写入 manifest，REPORT 引用）
|
||||
|
||||
|
||||
# 目标文件 → 候选链（顺序即优先级）。REQUIRED 目标缺失时退出码为 1，
|
||||
# 基准脚本随后走 HOG / Otsu 降级路线（见各脚本的降级分支）。
|
||||
TARGETS: dict[str, tuple[Candidate, ...]] = {
|
||||
"pose-s": (
|
||||
Candidate(
|
||||
repo="bukuroo/RTMPose-ONNX",
|
||||
filename="rtmpose-s.onnx",
|
||||
dest="rtmpose-s.onnx",
|
||||
note="RTMPose-s body(COCO-17, SimCC) 社区 ONNX 导出",
|
||||
),
|
||||
),
|
||||
"pose-m": (
|
||||
Candidate(
|
||||
repo="bukuroo/RTMPose-ONNX",
|
||||
filename="rtmpose-m.onnx",
|
||||
dest="rtmpose-m.onnx",
|
||||
note="RTMPose-m body(COCO-17, SimCC)，规模-时延对比用（可选）",
|
||||
),
|
||||
),
|
||||
"matting": (
|
||||
Candidate(
|
||||
repo="onnx-community/BiRefNet_lite-ONNX",
|
||||
filename="onnx/model.onnx",
|
||||
dest="birefnet_lite.onnx",
|
||||
note="BiRefNet_lite fp32 1024x1024（任务文档首选 lite 变体）",
|
||||
),
|
||||
Candidate(
|
||||
repo="onnx-community/BiRefNet_lite-ONNX",
|
||||
filename="onnx/model_fp16.onnx",
|
||||
dest="birefnet_lite_fp16.onnx",
|
||||
note="BiRefNet_lite fp16 降级档（体积约减半）",
|
||||
),
|
||||
Candidate(
|
||||
repo="briaai/RMBG-1.4",
|
||||
filename="onnx/model.onnx",
|
||||
dest="rmbg14.onnx",
|
||||
note="RMBG-1.4 fp32 末级降级（非 BiRefNet，需在报告注明）",
|
||||
),
|
||||
),
|
||||
}
|
||||
REQUIRED = ("pose-s", "matting")
|
||||
|
||||
|
||||
def _endpoints() -> list[str]:
|
||||
"""镜像优先：HF_ENDPOINT 可覆盖，默认 hf-mirror.com，再兜底官方站。"""
|
||||
eps = []
|
||||
for ep in (
|
||||
os.environ.get("HF_ENDPOINT", "").rstrip("/"),
|
||||
"https://hf-mirror.com",
|
||||
"https://huggingface.co",
|
||||
):
|
||||
if ep and ep not in eps:
|
||||
eps.append(ep)
|
||||
return eps
|
||||
|
||||
|
||||
def _url(endpoint: str, cand: Candidate) -> str:
|
||||
return f"{endpoint}/{cand.repo}/resolve/main/{cand.filename}"
|
||||
|
||||
|
||||
def _repo_revision(endpoint: str, repo: str) -> str:
|
||||
"""通过 model API 拿当前 revision（commit sha），失败记 unknown。"""
|
||||
req = urllib.request.Request(f"{endpoint}/api/models/{repo}", headers={"User-Agent": UA})
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=20) as resp:
|
||||
data = json.load(resp)
|
||||
return str(data.get("sha", "unknown"))
|
||||
except (urllib.error.URLError, json.JSONDecodeError, OSError):
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _download_once(
|
||||
endpoint: str,
|
||||
cand: Candidate,
|
||||
models_dir: Path,
|
||||
deadline: float,
|
||||
) -> bool:
|
||||
"""单次下载（支持 .part 续传）。超预算/网络错误返回 False，由外层重试或降级。"""
|
||||
partial = models_dir / f"{cand.dest}.part"
|
||||
offset = partial.stat().st_size if partial.exists() else 0
|
||||
headers = {"User-Agent": UA}
|
||||
if offset:
|
||||
headers["Range"] = f"bytes={offset}-"
|
||||
req = urllib.request.Request(_url(endpoint, cand), headers=headers)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30) as resp, open(partial, "ab") as fh:
|
||||
content_length = resp.headers.get("Content-Length")
|
||||
total = int(content_length) + offset if content_length else None
|
||||
last_bytes = 0
|
||||
last_time = time.perf_counter()
|
||||
while True:
|
||||
if time.perf_counter() > deadline:
|
||||
print(" [超时] 超过下载预算，放弃当前候选", flush=True)
|
||||
return False
|
||||
chunk = resp.read(CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
fh.write(chunk)
|
||||
offset += len(chunk)
|
||||
now = time.perf_counter()
|
||||
if offset - last_bytes >= PROGRESS_EVERY_BYTES or now - last_time > 10.0:
|
||||
total_str = f"/{total / 1048576:.1f}MB" if total else ""
|
||||
print(f" {offset / 1048576:.1f}MB{total_str} ...", flush=True)
|
||||
last_bytes = offset
|
||||
last_time = now
|
||||
partial.replace(models_dir / cand.dest)
|
||||
return True
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 416 and partial.exists(): # Range 起点已到文件尾 = 实际已下完
|
||||
partial.replace(models_dir / cand.dest)
|
||||
return True
|
||||
print(f" [HTTP {exc.code}] {exc.reason}", flush=True)
|
||||
return False
|
||||
except OSError as exc: # 含 socket.timeout / 连接重置：保留 .part 供续传
|
||||
print(f" [网络错误] {exc}", flush=True)
|
||||
return False
|
||||
|
||||
|
||||
def _mark_cached(cand: Candidate, dest: Path, manifest: dict[str, Any]) -> None:
|
||||
"""缓存命中也补全 manifest（repo/revision 供 REPORT 引用；离线时 revision 记 unknown）。"""
|
||||
entry = manifest.get(cand.dest, {})
|
||||
manifest[cand.dest] = {
|
||||
"repo": entry.get("repo", cand.repo),
|
||||
"revision": entry.get("revision") or _repo_revision(_endpoints()[0], cand.repo),
|
||||
"bytes": dest.stat().st_size,
|
||||
"cached": True,
|
||||
"note": cand.note,
|
||||
}
|
||||
|
||||
|
||||
def _try_candidate(
|
||||
cand: Candidate,
|
||||
models_dir: Path,
|
||||
budget_s: int,
|
||||
manifest: dict[str, Any],
|
||||
) -> bool:
|
||||
"""单候选全流程：双端点 × 重试一次，成功写 manifest。"""
|
||||
dest = models_dir / cand.dest
|
||||
for endpoint in _endpoints():
|
||||
deadline = time.perf_counter() + budget_s
|
||||
for attempt in range(1, ATTEMPTS_PER_CANDIDATE + 1):
|
||||
print(f"[下载] {cand.repo}/{cand.filename} <- {endpoint}（第 {attempt} 次）")
|
||||
t0 = time.perf_counter()
|
||||
if _download_once(endpoint, cand, models_dir, deadline):
|
||||
manifest[cand.dest] = {
|
||||
"repo": cand.repo,
|
||||
"revision": _repo_revision(endpoint, cand.repo),
|
||||
"endpoint": endpoint,
|
||||
"bytes": dest.stat().st_size,
|
||||
"elapsed_s": round(time.perf_counter() - t0, 1),
|
||||
"cached": False,
|
||||
"note": cand.note,
|
||||
}
|
||||
print(f"[完成] {cand.dest} {manifest[cand.dest]['bytes'] / 1048576:.1f}MB")
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def ensure_target(
|
||||
name: str,
|
||||
models_dir: Path,
|
||||
budget_s: int,
|
||||
manifest: dict[str, Any],
|
||||
) -> bool:
|
||||
"""目标文件级编排：按候选链降级。返回该目标是否本地可用。"""
|
||||
for cand in TARGETS[name]:
|
||||
# 候选链内任何一个目标文件已在本地（含上次跑过其它候选）都算满足
|
||||
dest = models_dir / cand.dest
|
||||
if dest.exists() and dest.stat().st_size > 0:
|
||||
print(f"[缓存] {name}: {cand.dest} 已在本地")
|
||||
_mark_cached(cand, dest, manifest)
|
||||
return True
|
||||
if _try_candidate(cand, models_dir, budget_s, manifest):
|
||||
return True
|
||||
print(f"[降级] {name} 的候选 {cand.repo}/{cand.filename} 失败，尝试下一候选")
|
||||
return False
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="T9 SPIKE 模型下载（唯一联网环节）")
|
||||
parser.add_argument("--models-dir", type=Path, default=DEFAULT_MODELS_DIR)
|
||||
parser.add_argument(
|
||||
"--budget", type=int, default=DEFAULT_BUDGET_SECONDS, help="单文件下载预算（秒）"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only", action="append", choices=sorted(TARGETS), help="只下载指定目标（可多次）"
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
models_dir: Path = args.models_dir
|
||||
models_dir.mkdir(parents=True, exist_ok=True)
|
||||
manifest_path = models_dir / "manifest.json"
|
||||
manifest: dict[str, Any] = {}
|
||||
if manifest_path.exists():
|
||||
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||
|
||||
selected = args.only or sorted(TARGETS)
|
||||
missing_required = False
|
||||
for name in selected:
|
||||
ok = ensure_target(name, models_dir, args.budget, manifest)
|
||||
required = name in REQUIRED
|
||||
if not ok:
|
||||
print(f"[失败] {name}：全部候选未取到（{'必需' if required else '可选'}）")
|
||||
missing_required |= required
|
||||
|
||||
dump = dict(sorted(manifest.items()))
|
||||
manifest_path.write_text(json.dumps(dump, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
print(f"\nmanifest 已写入 {manifest_path}")
|
||||
for key, entry in dump.items():
|
||||
repo = entry.get("repo", "?")
|
||||
revision = entry.get("revision", "?")
|
||||
print(f" {key}: {repo}@{revision} cached={entry.get('cached')}")
|
||||
if missing_required:
|
||||
print("\n有必需目标缺失：基准脚本将走降级路线（HOG/Otsu），详见 REPORT.md")
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,190 @@
|
||||
"""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())
|
||||
@@ -0,0 +1,139 @@
|
||||
"""T9 SPIKE · 合成人形图生成器（纯几何绘制，不使用任何真实人脸/人体数据——合规红线）。
|
||||
|
||||
rtmpose_benchmark / birefnet_matting 的输入图、body_measure_prototype 的体型变体
|
||||
全部由本模块确定性生成（同一参数永远产出同一张图），并把「生成时即已知」的真值
|
||||
（三线高度 / 三线宽度 / 期望体型）写入 ground_truth.json——这就是体型分类一致性
|
||||
验证的最小真值来源。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
ASSETS_DIR = Path(__file__).resolve().parent / "assets"
|
||||
|
||||
# ── 画布与几何（像素；画布 384x512，人形居中站立） ────────────────────────────
|
||||
CANVAS_W, CANVAS_H = 384, 512
|
||||
CX = CANVAS_W // 2
|
||||
HEAD_R = 34
|
||||
HEAD_TOP = 26
|
||||
HEAD_CY = HEAD_TOP + HEAD_R
|
||||
NECK_HW = 8
|
||||
Y_SHOULDER, Y_WAIST, Y_HIP = 112, 200, 280
|
||||
LEG_BOTTOM, LEG_GAP = 486, 8
|
||||
ARM_T = 10 # 手臂沿躯干边缘绘制，每侧超出边缘 ARM_T/2 → 三线宽度 ≈ 2×半宽 + ARM_T
|
||||
|
||||
BODY_COLOR: tuple[int, int, int] = (150, 168, 198) # BGR，米肤色调
|
||||
BG_FLAT: tuple[int, int, int] = (242, 242, 242) # 纯色背景
|
||||
GRAD_TOP = np.array([250, 250, 250], dtype=np.float32) # 渐变背景（上亮）
|
||||
GRAD_BOTTOM = np.array([158, 166, 182], dtype=np.float32) # 渐变背景（下暗）
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BodyShape:
|
||||
"""体型变体：三处剖面半宽（像素）+ 期望分类（生成时即已知，即一致性断言真值）。"""
|
||||
|
||||
key: str
|
||||
expect: str
|
||||
shoulder_hw: int
|
||||
waist_hw: int
|
||||
hip_hw: int
|
||||
|
||||
|
||||
# STD 是 §3.2/§3.3 的标准输入图（中性比例，分类为 H）；其余五种覆盖 H/A/V/O/X 全类别
|
||||
VARIANTS: tuple[BodyShape, ...] = (
|
||||
BodyShape("STD", "H", 60, 50, 60),
|
||||
BodyShape("H", "H", 58, 55, 57),
|
||||
BodyShape("A", "A", 44, 46, 68),
|
||||
BodyShape("V", "V", 72, 48, 54),
|
||||
BodyShape("X", "X", 60, 34, 62),
|
||||
BodyShape("O", "O", 52, 66, 56),
|
||||
)
|
||||
|
||||
|
||||
def _canvas(background: str) -> np.ndarray:
|
||||
if background == "gradient":
|
||||
t = np.linspace(0.0, 1.0, CANVAS_H, dtype=np.float32)[:, None]
|
||||
rows = GRAD_TOP[None, :] * (1.0 - t) + GRAD_BOTTOM[None, :] * t # (H,3) 逐行插值
|
||||
img = np.repeat(rows[:, None, :], CANVAS_W, axis=1).astype(np.uint8)
|
||||
return np.ascontiguousarray(img)
|
||||
return np.full((CANVAS_H, CANVAS_W, 3), BG_FLAT, dtype=np.uint8)
|
||||
|
||||
|
||||
def draw_figure(shape: BodyShape, background: str = "flat") -> np.ndarray:
|
||||
"""按体型参数绘制人形剪影：头圆 + 颈 + 躯干三段梯形 + 贴边手臂 + 双腿。"""
|
||||
img = _canvas(background)
|
||||
left_edge = [
|
||||
(CX - shape.shoulder_hw, Y_SHOULDER),
|
||||
(CX - shape.waist_hw, Y_WAIST),
|
||||
(CX - shape.hip_hw, Y_HIP),
|
||||
]
|
||||
right_edge = [
|
||||
(CX + shape.shoulder_hw, Y_SHOULDER),
|
||||
(CX + shape.waist_hw, Y_WAIST),
|
||||
(CX + shape.hip_hw, Y_HIP),
|
||||
]
|
||||
poly = np.array([*left_edge, right_edge[2], right_edge[1], right_edge[0]], dtype=np.int32)
|
||||
cv2.fillPoly(img, [poly], BODY_COLOR)
|
||||
# 手臂：沿躯干左右轮廓的粗折线，圆角过渡，对三条剖面线近似均匀加宽 ARM_T/2×2
|
||||
for edge in (left_edge, right_edge):
|
||||
cv2.polylines(
|
||||
img,
|
||||
[np.array(edge, dtype=np.int32)],
|
||||
isClosed=False,
|
||||
color=BODY_COLOR,
|
||||
thickness=ARM_T,
|
||||
lineType=cv2.LINE_8,
|
||||
)
|
||||
# 颈 + 头
|
||||
cv2.rectangle(img, (CX - NECK_HW, HEAD_CY), (CX + NECK_HW, Y_SHOULDER + 6), BODY_COLOR, -1)
|
||||
cv2.circle(img, (CX, HEAD_CY), HEAD_R, BODY_COLOR, -1)
|
||||
# 双腿：臀线以下的两个矩形，总宽 = 臀宽，中间留 LEG_GAP
|
||||
leg_w = shape.hip_hw - LEG_GAP // 2
|
||||
for x0 in (CX - shape.hip_hw, CX + LEG_GAP // 2):
|
||||
cv2.rectangle(img, (x0, Y_HIP - 2), (x0 + leg_w, LEG_BOTTOM), BODY_COLOR, -1)
|
||||
return img
|
||||
|
||||
|
||||
def _ground_truth(shape: BodyShape) -> dict[str, Any]:
|
||||
"""生成时即已知的真值：三线高度、三线宽度（含手臂加宽）、包围盒、期望分类。"""
|
||||
half = max(shape.shoulder_hw, shape.hip_hw)
|
||||
return {
|
||||
"key": shape.key,
|
||||
"expect": shape.expect,
|
||||
"half_widths": {
|
||||
"shoulder": shape.shoulder_hw,
|
||||
"waist": shape.waist_hw,
|
||||
"hip": shape.hip_hw,
|
||||
},
|
||||
"gt_widths": {
|
||||
"shoulder": 2 * shape.shoulder_hw + ARM_T,
|
||||
"waist": 2 * shape.waist_hw + ARM_T,
|
||||
"hip": 2 * shape.hip_hw + ARM_T,
|
||||
},
|
||||
"line_y": {"shoulder": Y_SHOULDER, "waist": Y_WAIST, "hip": Y_HIP},
|
||||
"bbox": [CX - half - ARM_T, HEAD_TOP, CX + half + ARM_T, LEG_BOTTOM],
|
||||
}
|
||||
|
||||
|
||||
def ensure_assets(assets_dir: Path = ASSETS_DIR) -> dict[str, Any]:
|
||||
"""（重新）生成全部合成图与真值 JSON。确定性输出，可反复执行（幂等）。"""
|
||||
variants_dir = assets_dir / "variants"
|
||||
variants_dir.mkdir(parents=True, exist_ok=True)
|
||||
gt: dict[str, Any] = {"canvas": [CANVAS_W, CANVAS_H], "arm_thickness": ARM_T, "variants": {}}
|
||||
for shape in VARIANTS:
|
||||
cv2.imwrite(str(variants_dir / f"{shape.key}.png"), draw_figure(shape, "flat"))
|
||||
gt["variants"][shape.key] = _ground_truth(shape)
|
||||
# §3.2/§3.3 的标准输入：纯色背景 + 渐变背景变体（各一张）
|
||||
cv2.imwrite(str(assets_dir / "synthetic_person.png"), draw_figure(VARIANTS[0], "flat"))
|
||||
cv2.imwrite(str(assets_dir / "synthetic_person_grad.png"), draw_figure(VARIANTS[0], "gradient"))
|
||||
(assets_dir / "ground_truth.json").write_text(
|
||||
json.dumps(gt, indent=2, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
return gt
|
||||
@@ -0,0 +1,169 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user