From 211341cbbb85eacc2b09bfce53020954f977ae15 Mon Sep 17 00:00:00 2001 From: linaro Date: Sat, 17 Oct 2026 20:33:00 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=8C=89=E7=AE=97=E6=B3=95=E7=94=9F?= =?UTF-8?q?=E6=88=90=E6=A8=A1=E6=9D=BF=E5=B9=B6=E6=94=AF=E6=8C=81=E4=BB=8E?= =?UTF-8?q?=20Git=20LFS=20=E4=B8=8B=E8=BD=BD=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - init 新增 yolo26/resnet/arcface/lprnet/face-recognition 模板（多模型人脸检测+识别） - 新增 embedding/text 输出类型与模板资源（COCO、字符表等） - 新增 algpkg download-models：从 Gitea LFS 下载 models/BM1688 并做 SHA256 校验 - 补充单元测试与 README --- README.md | 69 +- src/algpkg_sdk/api.py | 12 +- src/algpkg_sdk/cli.py | 212 +++--- src/algpkg_sdk/manifest.py | 2 +- src/algpkg_sdk/models.py | 258 +++++++ src/algpkg_sdk/templates.py | 1341 +++++++++++++++++++++++++++++++++++ tests/test_sdk.py | 436 ++++++++++++ 7 files changed, 2220 insertions(+), 110 deletions(-) create mode 100644 src/algpkg_sdk/models.py create mode 100644 src/algpkg_sdk/templates.py diff --git a/README.md b/README.md index 28b1178..3bfbe64 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,12 @@ worker 侧 `workers/video/pyproject.toml` 以**路径依赖**指向本项目（` # ① 生成脚手架（目录骨架 + manifest 模板 + 可运行示例 + 自测说明） algpkg init /tmp/smoke-detect --id demo.smoke-detect --name 抽烟检测 --version 1.0.0 +# 按算法生成（推荐）：YOLO26 模板已内置完整 SAIL 前后处理与 COCO 80 类标签 +algpkg init /tmp/yolo26-demo --id demo.yolo26 --name YOLO26检测 --algorithm yolo26 + # ② 实现 /tmp/smoke-detect/<包内代码目录>/__init__.py 的 load() / infer() -# 模型放 /tmp/smoke-detect/models/，或打包时用 --model 指定 +# yolo26 模板已实现，无需再写；模型放 /tmp/yolo26-demo/models/ +# 或打包时用 --model 指定 # ③ 打包 + 自检（不需要设备） algpkg build /tmp/smoke-detect --out demo.smoke-detect-1.0.0.mbank @@ -44,6 +48,64 @@ algpkg inspect demo.smoke-detect-1.0.0.mbank # 清单/模型/快照/ algpkg verify demo.smoke-detect-1.0.0.mbank # 结构与模型校验 ``` +### 2.1 算法模板（`algpkg init --algorithm/--template`） + +| 模板 | 生成内容 | +| ---- | -------- | +| `generic`（默认） | 目录骨架 + 占位 `load()/infer()`，自己按模型补前后处理 | +| `yolo26` | **可直接运行的 YOLO26 检测包**：`load()` 构造 SAIL 引擎；`infer()` 完成 letterbox 前处理 → 推理 → 置信度过滤/坐标反变换；manifest 预置 `yolo26s_fp16_1b` 模型、`conf/interval/draw` 参数与 COCO 80 类 labels | +| `resnet` | **可直接运行的 ResNet 图像分类包**：resize + ImageNet mean/std 归一化 → 推理 → softmax/top-1；分类结果用整帧框承载（score/类别），绘制与告警链路无需改动；manifest 预置 `resnet50_fp32_1b` 模型、`conf/interval/draw` 参数与 1000 类占位 labels | +| `arcface` | **ArcFace 人脸特征提取包**：RGB → resize 112×112 → 推理 → L2 归一化；输出 512 维 embedding 放在 `AlgoResult.meta["embedding"]`，`output.type=embedding`；manifest 预置 `arcface_resnet50_fp32_1b` 模型与 `interval` 参数 | +| `lprnet` | **LPRNet 车牌识别包**：BGR → resize 到模型输入 → `(x-127.5)/128` → CTC 贪心解码；车牌字符串放在 `AlgoResult.meta["text"]`，`output.type=text`；manifest 预置 `lprnet_fp32_1b` 模型、`interval` 参数与字符表 `assets/chars.txt` | +| `face-recognition` | **人脸检测 + 识别双模型包**：RetinaFace 检测人脸 → 按框外扩裁脸 → ArcFace 提取 512 维特征；`boxes` 为人脸框，`meta["embeddings"]` 与框逐项对应；manifest 同时声明 `face_detector` 与 `arcface` 两份模型 | + +`--algorithm` 也接受 `yolov26`、`yolo26-detect`、`resnet50`、`classifier`、`arc-face`、 +`face-embedding`、`lpr`、`license-plate`、`face`、`retinaface-arcface` 等别名； +`--kind` 缺省沿用模板（yolo26/face-recognition → `detector`，其余 → `classifier`）， +显式给出可覆盖。模型放到模板提示的路径即可打包，不在源码目录时用 +`--model <模型名>=<磁盘路径>` 指定实际文件： + +```bash +algpkg init /tmp/yolo26-demo --id demo.yolo26 --name YOLO26检测 --algorithm yolo26 +algpkg build /tmp/yolo26-demo --out demo.yolo26-1.0.0.mbank \ + --model yolo26s_fp16_1b=/path/to/yolo26s_fp16_1b.bmodel + +algpkg init /tmp/resnet-demo --id demo.resnet --name ResNet分类 --algorithm resnet +algpkg build /tmp/resnet-demo --out demo.resnet-1.0.0.mbank \ + --model resnet50_fp32_1b=/path/to/resnet50_fp32_1b.bmodel + +algpkg init /tmp/arcface-demo --id demo.arcface --name ArcFace特征 --algorithm arcface +algpkg build /tmp/arcface-demo --out demo.arcface-1.0.0.mbank \ + --model arcface_resnet50_fp32_1b=/path/to/arcface_resnet50_fp32_1b.bmodel + +algpkg init /tmp/lprnet-demo --id demo.lprnet --name 车牌识别 --algorithm lprnet +algpkg build /tmp/lprnet-demo --out demo.lprnet-1.0.0.mbank \ + --model lprnet_fp32_1b=/path/to/lprnet_fp32_1b.bmodel + +algpkg init /tmp/face-demo --id demo.face --name 人脸检测识别 --algorithm face-recognition +algpkg build /tmp/face-demo --out demo.face-1.0.0.mbank \ + --model face_detector=/path/to/retinaface_mobilenet0.25_fp32_1b.bmodel \ + --model arcface=/path/to/arcface_resnet50_fp32_1b.bmodel +``` + +### 2.2 下载演示模型（`algpkg download-models`） + +仓库里的 5 个 BM1688 演示模型（YOLO26 / ResNet50 / ArcFace / LPRNet / RetinaFace， +共约 299 MB）通过 **Git LFS** 存放在 `models/BM1688/`。没有 `git-lfs`、或不想拉取 +整个仓库时，可直接从 Gitea 的 LFS 端点下载（自动 SHA256 校验；已存在且正确的文件跳过）： + +```bash +algpkg download-models # 全部 5 个 → models/BM1688 +algpkg download-models --model yolo26 --model resnet +algpkg download-models --dest /tmp/models --force # 自定义目录 / 强制重下 +# 私有仓库或自定义地址： +algpkg download-models --url http:////.git --token +``` + +默认仓库地址可用 `ALGPKG_REPO_URL` 覆盖，token 也可用 `GITEA_TOKEN` 环境变量提供。 +模型名支持 `yolo26`、`resnet`、`arcface`、`lprnet`、`retinaface` 及常用别名 +（`yolov26`、`resnet50`、`lpr`、`face-detector` 等）。模型需已 push 到远端 LFS。 + ## 3. 本地测试（不需要 worker） ```bash @@ -144,10 +206,13 @@ algpkg-sdk/ ├── src/algpkg_sdk/ │ ├── api.py 接口契约（Algorithm / AlgoContext / AlgoResult）+ 别名注册 │ ├── manifest.py 清单字段与校验（models / params_schema / alarm / legacy…） +│ ├── models.py download-models：从 Git LFS 拉取 models/BM1688 演示模型 │ ├── loader.py AlgoPackage：读 svfs、挂代码（试加载）、构造实例 │ ├── builder.py 打包（module_bank 打代码 + svfs 放模型/资源 + 快照=版本 + 签名） │ ├── testkit.py 本地测试台：check / run / compare（无需 worker） -│ └── cli.py algpkg build|inspect|verify|init|list|check|run|compare +│ ├── templates.py init 算法模板（generic / yolo26 / resnet / arcface / lprnet / face-recognition） +│ └── cli.py algpkg build|inspect|verify|init|list|check|run|compare|download-models +├── models/BM1688/ Git LFS：5 个 BM1688 演示模型（可 download-models 拉取） ├── examples/person-count/ 完整示例：人员计数（自带实现，只依赖本 SDK） │ └── … └── tests/ pytest（无设备跑静态，有设备跑推理，自动跳过） diff --git a/src/algpkg_sdk/api.py b/src/algpkg_sdk/api.py index 70ca20b..85b14c0 100644 --- a/src/algpkg_sdk/api.py +++ b/src/algpkg_sdk/api.py @@ -36,8 +36,11 @@ SDK_MODULE_NAME = "video_algpkg_api" #: 算法包类型（manifest.kind）——决定 UI 呈现与结果渲染方式。 KINDS = ("detector", "classifier", "counter", "tracker") -#: 结果类型（manifest.output.type）。 -OUTPUT_TYPES = ("boxes", "count", "counter") +#: 结果类型（manifest.output.type）。`embedding` = 特征向量（如 ArcFace）， +#: 结果放在 :class:`AlgoResult` 的 ``meta["embedding"]``（列表）；`text` = 文本识别 +#: 结果（如 LPRNet 车牌号），放在 ``meta["text"]``；worker 的任务渲染仍以 boxes +#: 为主，embedding/text 供后续比对、检索或上层业务逻辑使用。 +OUTPUT_TYPES = ("boxes", "count", "counter", "embedding", "text") class AlgoError(RuntimeError): @@ -51,7 +54,10 @@ class AlgoResult: —— 绘制、告警、统计都基于它；无框的算法（纯计数）留 None； - ``counts``：计数类算法的分项计数（如 ``{"person": 3}``；无分项可用 ``{"": n}``）； - ``value``：单值结果（如"当前画面人数"）； - - ``labels``：类别名（下标与 boxes 的 cls 对应；缺省用包内 labels 或 COCO 表）。 + - ``labels``：类别名（下标与 boxes 的 cls 对应；缺省用包内 labels 或 COCO 表）； + - ``meta``：附加数据（如 ArcFace 的 ``{"embedding": [...], "dim": 512}``、 + LPRNet 的 ``{"text": "京A12345"}``；worker 当前不参与绘制/告警， + 供程序化调用与后续扩展）。 """ __slots__ = ("boxes", "counts", "value", "labels", "meta") diff --git a/src/algpkg_sdk/cli.py b/src/algpkg_sdk/cli.py index 88e10c1..491b787 100644 --- a/src/algpkg_sdk/cli.py +++ b/src/algpkg_sdk/cli.py @@ -2,11 +2,12 @@ 安装本项目后即可直接用（不需要 worker 仓库）:: - algpkg init --id my.algo --name 我的算法 # 脚手架 + algpkg init --id my.algo --name 我的算法 [--algorithm ...] algpkg build --out my.algo-1.0.0.mbank [--model 名=路径] [--encrypt KEY] [--sign KEY] algpkg inspect [--json] # 看清单/模型/快照/签名 algpkg verify [--signature-key KEY] # 结构+模型+（可选）验签 algpkg list + algpkg download-models [--dest models/BM1688] [--model yolo26 ...] # Git LFS 拉模型 本地测试（不需要 worker；`run`/`compare` 需要设备）:: @@ -25,6 +26,7 @@ from pathlib import Path from .builder import AlgoBuilder, BuildError, inspect_package, verify_package from .manifest import Manifest, code_package_for +from .templates import AlgorithmTemplate, get_template, template_summaries # ====================================================================== 子命令 @@ -115,7 +117,31 @@ def cmd_list(args: argparse.Namespace) -> int: return 0 +def cmd_download_models(args: argparse.Namespace) -> int: + from .models import DEFAULT_DEST, ModelDownloadError, download_models + + dest = args.dest or DEFAULT_DEST + try: + results = download_models(dest, names=args.model, repo_url=args.url, + token=args.token, username=args.username, + force=args.force) + except (ModelDownloadError, ValueError) as exc: + print(f"✗ 模型下载失败: {exc}", file=sys.stderr) + return 1 + downloaded = sum(1 for item in results if item["status"] == "downloaded") + skipped = len(results) - downloaded + print(f"✓ 模型目录: {Path(dest).resolve()}（下载 {downloaded}，跳过 {skipped}）") + return 0 + + def cmd_init(args: argparse.Namespace) -> int: + try: + template = get_template(args.algorithm) + except ValueError as exc: + print(f"✗ {exc}", file=sys.stderr) + for key, summary in template_summaries(): + print(f" - {key}: {summary}", file=sys.stderr) + return 2 root = Path(args.dir) manifest_path = root / "manifest.json" if manifest_path.exists() and not args.force: @@ -126,119 +152,78 @@ def cmd_init(args: argparse.Namespace) -> int: (root / code_package).mkdir(parents=True, exist_ok=True) (root / "assets").mkdir(parents=True, exist_ok=True) - manifest_path.write_text(json.dumps(_scaffold_manifest(args, code_package), + manifest_path.write_text(json.dumps(_scaffold_manifest(args, code_package, template), ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - (root / code_package / "__init__.py").write_text(_scaffold_code(args), encoding="utf-8") - labels = root / "assets" / "labels.txt" + (root / code_package / "__init__.py").write_text( + _scaffold_code(args, code_package, template), encoding="utf-8") + labels = root / template.labels_file if not labels.exists(): - labels.write_text("# 每行一个类别名（第 0 行 = 类别 0）；按需替换\nobject\n", encoding="utf-8") + labels.parent.mkdir(parents=True, exist_ok=True) + labels.write_text(template.labels, encoding="utf-8") readme = root / "README.md" if not readme.exists(): - readme.write_text(_scaffold_readme(args, code_package), encoding="utf-8") + readme.write_text(_scaffold_readme(args, code_package, template), encoding="utf-8") - print(f"✓ 已生成脚手架: {root}") - print(f" 1) 把模型放到 {root}/models/（或打包时用 --model 名=<路径> 指定）") - print(f" 2) 实现 {code_package}/__init__.py 的 load/infer") - print(f" 3) 打包: workers/video/algpkg build {root} --out {args.id}-{args.version}.mbank") + model_specs = template.model_specs(args.id) + print(f"✓ 已生成脚手架: {root}（模板 {template.key}: {template.label}）") + if len(model_specs) == 1: + print(f" 1) 把模型放到 {root}/{model_specs[0]['file']}" + f"（或打包时用 --model {model_specs[0]['name']}=<路径> 指定）") + else: + print(" 1) 准备以下模型（放源码目录下，或用 --model 名=<路径> 指定）：") + for spec in model_specs: + print(f" {root}/{spec['file']} （--model {spec['name']}=<路径>）") + if template.key == "generic": + print(f" 2) 实现 {code_package}/__init__.py 的 load/infer") + else: + print(f" 2) {code_package}/__init__.py 已按 {template.label} 初始化，按需微调") + print(f" 3) 打包: algpkg build {root} --out {args.id}-{args.version}.mbank") return 0 # ====================================================================== 脚手架内容 -def _scaffold_manifest(args: argparse.Namespace, code_package: str) -> dict: - return { - "manifest_format": 1, - "id": args.id, - "name": args.name, - "version": args.version, - "api_version": 1, - "kind": args.kind, - "entry": "Algorithm", - "code_package": code_package, - "vendor": args.vendor or "", - "description": f"{args.name}（由 algpkg init 生成）", - "models": [{ - "name": "default", - "file": f"models/{args.id}.bmodel", - "precision": "", - "resolution": [640, 640], - "batch": 1, - "default": True, - }], - "input": {"width": 640, "height": 640, "color": "rgb", "layout": "planar", - "mean": [0, 0, 0], "std": [255, 255, 255]}, - "output": {"type": "boxes", "format": "N6", "classes_file": "assets/labels.txt"}, - "params_schema": [ - {"key": "conf", "type": "float", "default": 0.25, "min": 0.01, "max": 0.99, - "step": 0.01, "label": "置信度阈值", "desc": "低于该分数的目标被丢弃"}, - {"key": "interval", "type": "float", "default": 0.2, "min": 0.02, "max": 60, - "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次推理之间的最小间隔"}, - ], - "defaults": {"conf": 0.25, "interval": 0.2}, - "alarm": {"type_template": "{class}", "cooldown_key": "algo", - "need_snapshot": True, "need_record": True}, - "requires": {"license_features": [], "devices": [], "frameworks": ["sail"]}, - "ui": {"order": 100, "tags": [], "icon": "mdi-puzzle"}, - "legacy": {}, - } +def _scaffold_manifest(args: argparse.Namespace, code_package: str, + template: AlgorithmTemplate | None = None) -> dict: + template = template or get_template(getattr(args, "algorithm", "generic")) + return template.render_manifest( + algo_id=args.id, name=args.name, version=args.version, + kind=args.kind or template.kind, vendor=args.vendor, + code_package=code_package, + ) -def _scaffold_code(args: argparse.Namespace) -> str: - return f'''"""算法包实现：{args.name}（{args.id}）。 - -只依赖 `video_algpkg_api`（worker 提供的稳定接口）与设备框架（sophon.sail / numpy）。 -本文件由 `algpkg init` 生成，改完即可打包： - workers/video/algpkg build <本目录> --out {args.id}-{args.version}.mbank -""" -from __future__ import annotations - -import numpy as np -import sophon.sail as sail - -from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult +def _scaffold_code(args: argparse.Namespace, code_package: str, + template: AlgorithmTemplate | None = None) -> str: + template = template or get_template(getattr(args, "algorithm", "generic")) + return template.render_code( + algo_id=args.id, name=args.name, version=args.version, + kind=args.kind or template.kind, code_package=code_package, + ) -class Algorithm(AlgorithmBase): - """{args.name}。""" - - kind = "{args.kind}" - - #: 引擎句柄：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 - #: （不必判空）；close() 置 None 与声明冲突，那处显式 `type: ignore`——运行时语义就是"丢引用"。 - _net: sail.Engine - - def load(self) -> None: - # 模型字节来自算法包自身（.mbank 内的 models/**），不落盘： - data = self.ctx.model_bytes() - size = len(data) - self.ctx.logger.info("加载模型 %s（%d 字节, conf=%s）", - self.ctx.model_name, size, self.ctx.param("conf")) - self._net = sail.Engine(data, size, self.ctx.dev_id, sail.IOMode.SYSO) - self._graph = self._net.get_graph_names()[0] - # TODO: 按模型实际输入做前处理（可参考 worker 的 YOLO26._preprocess） - self._handle = sail.Handle(self.ctx.dev_id) - self._bmcv = sail.Bmcv(self._handle) - - def infer(self, frame) -> AlgoResult: - """单帧推理。返回 AlgoResult（boxes 为 (N,6)：x1,y1,x2,y2,score,cls）。 - - TODO: 前处理 → self._net.process(...) → 后处理（NMS/反变换）， - 再 return AlgoResult(boxes=boxes, labels=[...])。 - 下面这行只做"链路可跑通"的占位（空结果），替换为真实实现。 - """ - _ = frame - boxes = np.zeros((0, 6), dtype=np.float32) - return AlgoResult(boxes=boxes) - - def close(self) -> None: - self._net = None # type: ignore[assignment] -''' - - -def _scaffold_readme(args: argparse.Namespace, code_package: str) -> str: +def _scaffold_readme(args: argparse.Namespace, code_package: str, + template: AlgorithmTemplate | None = None) -> str: + template = template or get_template(getattr(args, "algorithm", "generic")) + models = template.model_specs(args.id) + if len(models) == 1: + step1 = (f"1. 把模型放到 `{models[0]['file']}`（或打包时 " + f"`--model {models[0]['name']}=<磁盘路径>` 指定）；") + else: + lines = "\n".join(f" - `{m['file']}`（`--model {m['name']}=<磁盘路径>`）" + for m in models) + step1 = f"1. 把模型放到源码目录下：\n{lines}；" + if template.key == "generic": + step2 = (f"2. 实现 `{code_package}/__init__.py` 的 `load()`（加载模型字节）与 " + f"`infer(frame)`（单帧结果）；") + else: + step2 = (f"2. `{code_package}/__init__.py` 已按 **{template.label}** 完成初始化" + f"（模型加载 + 前/后处理），一般只需按现场调整参数；") + notes = f"\n## 模板说明\n\n{template.notes}\n" if template.notes else "" return f"""# {args.name}（算法包 {args.id}） -由 `algpkg init` 生成。目录结构： +由 `algpkg init` 生成，模板：**{template.key}**（{template.summary}）。 +目录结构： ``` manifest.json 清单（算法元数据 / 模型 / 参数表单 / 告警策略） @@ -246,13 +231,13 @@ manifest.json 清单（算法元数据 / 模型 / 参数表单 / models/ 随包模型（.bmodel；打包时进包，不单独分发） assets/ 可选资源（labels.txt 类别名等） ``` - +{notes} ## 开发步骤 -1. 把模型放到 `models/`（或打包时 `--model 名=<磁盘路径>` 指定）； -2. 实现 `{code_package}/__init__.py` 的 `load()`（加载模型字节）与 `infer(frame)`（单帧结果）； -3. 打包：`workers/video/algpkg build . --out {args.id}-{args.version}.mbank`； -4. 检查：`workers/video/algpkg inspect {args.id}-{args.version}.mbank`； +{step1} +{step2} +3. 打包：`algpkg build . --out {args.id}-{args.version}.mbank`； +4. 检查：`algpkg check {args.id}-{args.version}.mbank`（有设备时加 `--deep` 真加载）； 5. 安装到设备：`POST /api/v1/video/algorithm/install`（算法管理页上传亦可）， 随后在任务表单里选择本算法（参数表单由 `manifest.params_schema` 生成）。 @@ -387,6 +372,18 @@ def build_parser() -> argparse.ArgumentParser: listing.add_argument("--encrypt", metavar="KEY") listing.set_defaults(func=cmd_list) + download = sub.add_parser("download-models", + help="从 Git LFS 下载 models/BM1688 下的演示模型") + download.add_argument("--dest", default="models/BM1688", + help="下载目录（默认 models/BM1688）") + download.add_argument("--model", action="append", metavar="NAME", + help="模型名/别名/文件名（可多次；缺省 all）") + download.add_argument("--url", help="Git 仓库 URL（默认 Gitea / ALGPKG_REPO_URL）") + download.add_argument("--token", help="Gitea token（默认读 GITEA_TOKEN）") + download.add_argument("--username", default="git", help="Basic 认证用户名（默认 git）") + download.add_argument("--force", action="store_true", help="已存在也重新下载") + download.set_defaults(func=cmd_download_models) + check = sub.add_parser("check", help="静态检查一个包（清单/结构/试加载；不需 worker）") check.add_argument("mbank") check.add_argument("--deep", action="store_true", help="再实例化入口类并加载默认模型（需设备）") @@ -428,7 +425,14 @@ def build_parser() -> argparse.ArgumentParser: init.add_argument("--id", required=True, help="算法 id（小写字母/数字/._-）") init.add_argument("--name", required=True, help="显示名（如：抽烟检测）") init.add_argument("--version", default="1.0.0") - init.add_argument("--kind", default="detector", choices=["detector", "classifier", "counter", "tracker"]) + init.add_argument("--kind", default=None, + choices=["detector", "classifier", "counter", "tracker"], + help="算法类型（缺省由算法模板决定：yolo26/face-recognition → detector，其余 → classifier）") + init.add_argument("--algorithm", "--template", dest="algorithm", default="generic", + metavar="NAME", + help="算法模板：generic / yolo26 / resnet / arcface / lprnet / " + "face-recognition（RetinaFace+ArcFace 检测识别）；" + "均含完整 SAIL 前后处理") init.add_argument("--vendor", default="") init.add_argument("--code-package", default="", help="包内代码模块名（缺省由 id 推导）") init.add_argument("--force", action="store_true") diff --git a/src/algpkg_sdk/manifest.py b/src/algpkg_sdk/manifest.py index 504d322..d427bcc 100644 --- a/src/algpkg_sdk/manifest.py +++ b/src/algpkg_sdk/manifest.py @@ -17,7 +17,7 @@ | `code_package` | 包代码的模块名（`modules//…`；缺省由 id 推导 `alg_`） | | `models`* | `[{name, file, precision, resolution, batch, default}]`，含唯一 default | | `input` | `{width,height,color,layout,mean,std}`（前处理契约，展示用） | -| `output` | `{type: boxes|count|counter, format, classes_file}` | +| `output` | `{type: boxes|count|counter|embedding|text, format, classes_file}` | | `params_schema` | `[{key,type,default,min,max,step,label,desc,options}]`（任务参数表单来源） | | `defaults` | 默认参数（conf/nms/interval/draw 等，任务未显式给值时用） | | `alarm` | `{type_template,cooldown_key,need_snapshot,need_record}`（告警生成策略） | diff --git a/src/algpkg_sdk/models.py b/src/algpkg_sdk/models.py new file mode 100644 index 0000000..37ca97d --- /dev/null +++ b/src/algpkg_sdk/models.py @@ -0,0 +1,258 @@ +"""演示模型下载：按 Git LFS 批量接口拉取 `models/BM1688/**`。 + +模型随本仓库走 Git LFS（见 `.gitattributes`），但 LFS 文件只在克隆时按需拉取； +本模块提供 ``algpkg download-models``，不依赖 ``git-lfs``、也不需要克隆仓库， +直接从 Gitea 的 LFS 端点下载并做 SHA256 校验。 + +默认仓库地址可用 ``--url`` 或环境变量 ``ALGPKG_REPO_URL`` 覆盖； +私有仓库可用 ``--token``（也读 ``GITEA_TOKEN``）配 ``--username``（默认 ``git``）。 +""" +from __future__ import annotations + +import base64 +import hashlib +import json +import os +import urllib.error +import urllib.parse +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +#: 模型随本仓库发布（Gitea）；可用 --url / ALGPKG_REPO_URL 覆盖 +DEFAULT_REPO_URL = os.environ.get( + "ALGPKG_REPO_URL", "http://124.71.68.6:3000/sophon-edge/algpkg-sdk.git") + +#: 默认下载目录（相对当前工作目录） +DEFAULT_DEST = "models/BM1688" + + +class ModelDownloadError(RuntimeError): + """模型下载/校验失败（消息面向"怎么修"）。""" + + +@dataclass(frozen=True) +class ModelFile: + """一个随包发布的模型（oid/size 即 Git LFS 指针里的 SHA256 与字节数）。""" + + key: str + filename: str + oid: str + size: int + label: str + + @property + def size_mb(self) -> float: + return self.size / 1048576.0 + + +#: 本仓库 models/BM1688 下的演示模型（顺序 = 下载顺序） +MODEL_FILES: dict[str, ModelFile] = { + "yolo26": ModelFile( + "yolo26", "yolo26s_fp16_1b.bmodel", + "085a8e9fc8a72065cb8333ea86d27c61742755b0dd88d5766112e25f60256360", + 24928256, "YOLO26 目标检测"), + "resnet": ModelFile( + "resnet", "resnet50_fp32_1b.bmodel", + "3c5131b6bb753c7cfaccf1be42d2a770156470dc1dc86147231b38d1728b0322", + 102318696, "ResNet50 图像分类"), + "arcface": ModelFile( + "arcface", "arcface_resnet50_fp32_1b.bmodel", + "7427ad5cece4ac1ec21583a6a67a8aeb58b79e279f9126c027600dbca06282b6", + 179355648, "ArcFace 人脸特征"), + "lprnet": ModelFile( + "lprnet", "lprnet_fp32_1b.bmodel", + "f97bd651db73d9c8378adb8d1a1d7bfcfc91e139a83c77c240bbf6ca14394e02", + 4017640, "LPRNet 车牌识别"), + "retinaface": ModelFile( + "retinaface", "retinaface_mobilenet0.25_fp32_1b.bmodel", + "a7be883fc021f9c605fd356c9224a05aa5573ac530338cec72cbbf98722af03f", + 2082912, "RetinaFace 人脸检测"), +} + +#: 常用写法 → 模型 key（也接受完整文件名与 key 本身） +ALIASES: dict[str, str] = { + "yolov26": "yolo26", + "yolo": "yolo26", + "yolo26-detect": "yolo26", + "resnet50": "resnet", + "arc": "arcface", + "arc-face": "arcface", + "face-embedding": "arcface", + "lpr": "lprnet", + "license-plate": "lprnet", + "plate": "lprnet", + "face-detector": "retinaface", + "retina": "retinaface", +} + + +def resolve_models(names: Iterable[str] | None = None) -> list[ModelFile]: + """把 `--model` 的名字/别名/文件名解析为模型列表（空/`all` = 全部）。""" + items = [] + for raw in names or []: + for part in str(raw).replace(",", " ").split(): + if part and part not in items: + items.append(part) + if not items or any(item.lower() == "all" for item in items): + return list(MODEL_FILES.values()) + + selected: list[ModelFile] = [] + for name in items: + key = name.lower() + if key in MODEL_FILES: + spec = MODEL_FILES[key] + else: + by_file = next((m for m in MODEL_FILES.values() if m.filename == name), None) + alias = ALIASES.get(key) + spec = by_file or (MODEL_FILES.get(alias) if alias else None) + if spec is None: + available = "、".join(f"{m.key}（{m.filename}）" for m in MODEL_FILES.values()) + raise ValueError(f"未知模型 {name!r}；可选：{available}（或 all）") + if spec not in selected: + selected.append(spec) + return selected + + +def lfs_batch_url(repo_url: str) -> str: + """仓库 URL → Git LFS batch API 地址。""" + base = str(repo_url or "").strip().rstrip("/") + if not base: + raise ValueError("仓库 URL 不能为空") + if base.endswith("/info/lfs"): + return base + "/objects/batch" + if not base.endswith(".git"): + base += ".git" + return base + "/info/lfs/objects/batch" + + +def _auth_headers(token: str | None, username: str | None) -> dict[str, str]: + if not token: + return {} + raw = f"{username or 'git'}:{token}".encode("utf-8") + return {"Authorization": "Basic " + base64.b64encode(raw).decode("ascii")} + + +def _sha256_file(path: Path, chunk: int = 1 << 20) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def _is_valid(path: Path, spec: ModelFile) -> bool: + return (path.is_file() and path.stat().st_size == spec.size + and _sha256_file(path) == spec.oid) + + +def _post_json(url: str, payload: dict, headers: dict[str, str], + timeout: float = 60.0) -> dict: + request = urllib.request.Request( + url, data=json.dumps(payload).encode("utf-8"), method="POST", + headers={"Accept": "application/vnd.git-lfs+json", + "Content-Type": "application/vnd.git-lfs+json", **headers}) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", "replace")[:500] + raise ModelDownloadError( + f"Git LFS 接口返回 HTTP {exc.code}: {detail or exc.reason}") from exc + except urllib.error.URLError as exc: + raise ModelDownloadError(f"无法连接 Git LFS 服务 {url}: {exc.reason}") from exc + + +def _download_file(href: str, target: Path, spec: ModelFile, + headers: dict[str, str], timeout: float = 300.0) -> None: + target.parent.mkdir(parents=True, exist_ok=True) + tmp = target.with_name(target.name + ".part") + request = urllib.request.Request(href, headers=headers) + digest = hashlib.sha256() + total = 0 + try: + with urllib.request.urlopen(request, timeout=timeout) as response, \ + tmp.open("wb") as handle: + while block := response.read(1 << 20): + handle.write(block) + digest.update(block) + total += len(block) + if total != spec.size: + raise ModelDownloadError( + f"{spec.filename} 大小不符：期望 {spec.size} 字节，收到 {total}") + if digest.hexdigest() != spec.oid: + raise ModelDownloadError( + f"{spec.filename} SHA256 校验失败（期望 {spec.oid[:12]}…，" + f"收到 {digest.hexdigest()[:12]}…）") + tmp.replace(target) + except urllib.error.HTTPError as exc: + raise ModelDownloadError( + f"下载 {spec.filename} 返回 HTTP {exc.code}: {exc.reason}") from exc + except urllib.error.URLError as exc: + raise ModelDownloadError(f"下载 {spec.filename} 失败: {exc.reason}") from exc + finally: + if tmp.exists(): + tmp.unlink() + + +def download_models(dest: str | Path = DEFAULT_DEST, *, + names: Iterable[str] | None = None, + repo_url: str | None = None, + token: str | None = None, + username: str = "git", + force: bool = False, + log=print) -> list[dict]: + """下载选中的模型到 `dest`；返回每个模型的状态（downloaded/skipped）。 + + 已存在且 SHA256/大小正确的文件默认跳过；`force=True` 强制重新下载。 + """ + selected = resolve_models(names) + target_dir = Path(dest) + target_dir.mkdir(parents=True, exist_ok=True) + token = token or os.environ.get("GITEA_TOKEN") or os.environ.get("ALGPKG_TOKEN") + + results: list[dict] = [] + missing: list[ModelFile] = [] + for spec in selected: + target = target_dir / spec.filename + if not force and _is_valid(target, spec): + log(f"· {spec.filename} 已存在且校验通过，跳过") + results.append({"name": spec.key, "path": str(target), + "size": spec.size, "status": "skipped"}) + else: + missing.append(spec) + + if missing: + batch_url = lfs_batch_url(repo_url or DEFAULT_REPO_URL) + payload = {"operation": "download", "transfers": ["basic"], + "objects": [{"oid": spec.oid, "size": spec.size} for spec in missing]} + response = _post_json(batch_url, payload, _auth_headers(token, username)) + objects = {str(item.get("oid")): item for item in response.get("objects") or []} + for index, spec in enumerate(missing, 1): + item = objects.get(spec.oid) + if item is None: + raise ModelDownloadError( + f"LFS 批量接口未返回 {spec.filename}（模型可能尚未 push 到远端）") + if item.get("error"): + error = item["error"] + raise ModelDownloadError( + f"{spec.filename} 下载被拒绝: {error.get('message') or error}") + action = (item.get("actions") or {}).get("download") or {} + href = action.get("href") + if not href: + raise ModelDownloadError(f"{spec.filename} 没有 download 链接") + href = urllib.parse.urljoin( + (repo_url or DEFAULT_REPO_URL).rstrip("/") + "/", href) + log(f"[{index}/{len(missing)}] 下载 {spec.filename}（{spec.size_mb:.1f} MB）…") + _download_file(href, target_dir / spec.filename, spec, + action.get("header") or {}) + log(f" ✓ {spec.filename} 完成（SHA256 校验通过）") + results.append({"name": spec.key, "path": str(target_dir / spec.filename), + "size": spec.size, "status": "downloaded"}) + return results + + +__all__ = ["DEFAULT_REPO_URL", "DEFAULT_DEST", "MODEL_FILES", "ALIASES", + "ModelFile", "ModelDownloadError", "resolve_models", "lfs_batch_url", + "download_models"] diff --git a/src/algpkg_sdk/templates.py b/src/algpkg_sdk/templates.py new file mode 100644 index 0000000..00b0503 --- /dev/null +++ b/src/algpkg_sdk/templates.py @@ -0,0 +1,1341 @@ +"""`algpkg init` 的算法模板：按算法生成 manifest 与可直接运行的入口实现。 + +模板 = **清单片段**（模型/参数/输入输出契约）+ **包内代码** + **随包资源**（labels 等）。 +默认 `generic` 只生成占位实现；`--algorithm yolo26` / `resnet` / `arcface` / `lprnet` / +`face-recognition` 生成完整算法包（YOLO26 检测、ResNet 图像分类、ArcFace 人脸特征、 +LPRNet 车牌识别、RetinaFace+ArcFace 人脸检测识别），放好模型即可打包运行。 + +新增算法模板的步骤： +1. 在本模块定义 :class:`AlgorithmTemplate`（模型/参数/标签/代码模板）； +2. 登记到 :data:`TEMPLATES`（`aliases` 可给常用写法，如 `yolov26`）； +3. 在 `tests/test_sdk.py` 加一条 init → build → inspect 的用例。 + +代码模板里的占位符（`__ALGO_ID__` 等）用 :meth:`AlgorithmTemplate.render_code` 替换。 +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +#: 标准 YOLO 前处理契约（RGB_PLANAR、0-255 → 0-1、无均值） +_YOLO_INPUT = {"width": 640, "height": 640, "color": "rgb", "layout": "planar", + "mean": [0, 0, 0], "std": [255, 255, 255]} + +_YOLO_PARAMS = ( + {"key": "conf", "type": "float", "default": 0.25, "min": 0.01, "max": 0.99, + "step": 0.01, "label": "置信度阈值", "desc": "低于该分数的目标被丢弃"}, + {"key": "interval", "type": "float", "default": 0.2, "min": 0.02, "max": 60, + "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次推理之间的最小间隔"}, +) + +_RESNET_PARAMS = ( + {"key": "conf", "type": "float", "default": 0.1, "min": 0.01, "max": 1.0, + "step": 0.01, "label": "置信度阈值", "desc": "top-1 概率低于该值时不上报"}, + {"key": "interval", "type": "float", "default": 0.5, "min": 0.02, "max": 60, + "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次推理之间的最小间隔"}, + {"key": "draw", "type": "bool", "default": False, + "label": "绘制叠加", "desc": "在推理帧上画整帧框与类别/置信度"}, +) + +_ARCFACE_PARAMS = ( + {"key": "interval", "type": "float", "default": 0.5, "min": 0.02, "max": 60, + "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次特征提取之间的最小间隔"}, +) + +_LPRNET_PARAMS = ( + {"key": "interval", "type": "float", "default": 0.5, "min": 0.02, "max": 60, + "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次车牌识别之间的最小间隔"}, +) + +_FACE_RECOGNITION_PARAMS = ( + {"key": "conf", "type": "float", "default": 0.5, "min": 0.01, "max": 0.99, + "step": 0.01, "label": "人脸置信度阈值", "desc": "低于该分数的人脸被丢弃"}, + {"key": "nms", "type": "float", "default": 0.3, "min": 0.01, "max": 0.99, + "step": 0.01, "label": "NMS IoU 阈值", "desc": "重叠人脸框抑制阈值"}, + {"key": "face_margin", "type": "float", "default": 0.15, "min": 0.0, "max": 0.5, + "step": 0.01, "label": "人脸外扩比例", "desc": "裁脸时在检测框外扩的比例"}, + {"key": "min_face", "type": "int", "default": 16, "min": 4, "max": 200, + "step": 1, "label": "最小人脸尺寸（像素）", + "desc": "宽或高小于该值的人脸只保留框，不提取特征"}, + {"key": "max_faces", "type": "int", "default": 20, "min": 1, "max": 100, + "step": 1, "label": "最大人脸数", "desc": "每帧最多识别的人脸数"}, + {"key": "interval", "type": "float", "default": 0.2, "min": 0.02, "max": 60, + "step": 0.02, "label": "抽帧间隔（秒）", "desc": "两次推理之间的最小间隔"}, + {"key": "draw", "type": "bool", "default": False, + "label": "绘制叠加", "desc": "在推理帧上画人脸框（预览叠加流 / 告警快照）"}, +) + +#: COCO 80 类（第 0 行 = 类别 0；与 YOLO26 输出 cls 对齐） +COCO_LABELS = """person +bicycle +car +motorcycle +airplane +bus +train +truck +boat +traffic light +fire hydrant +stop sign +parking meter +bench +bird +cat +dog +horse +sheep +cow +elephant +bear +zebra +giraffe +backpack +umbrella +handbag +tie +suitcase +frisbee +skis +snowboard +sports ball +kite +baseball bat +baseball glove +skateboard +surfboard +tennis racket +bottle +wine glass +cup +fork +knife +spoon +bowl +banana +apple +sandwich +orange +broccoli +carrot +hot dog +pizza +donut +cake +chair +couch +potted plant +bed +dining table +toilet +tv +laptop +mouse +remote +keyboard +cell phone +microwave +oven +toaster +sink +refrigerator +book +clock +vase +scissors +teddy bear +hair drier +toothbrush +""" + + +# ====================================================================== 代码模板 + +_GENERIC_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +只依赖 `video_algpkg_api`（worker 提供的稳定接口）与设备框架（sophon.sail / numpy）。 +本文件由 `algpkg init` 生成，改完即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + + +class Algorithm(AlgorithmBase): + """__ALGO_NAME__。""" + + kind = "__ALGO_KIND__" + + #: 引擎句柄：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 + #: （不必判空）；close() 置 None 与声明冲突，那处显式 `type: ignore`——运行时语义就是"丢引用"。 + _net: sail.Engine + + def load(self) -> None: + # 模型字节来自算法包自身（.mbank 内的 models/**），不落盘： + data = self.ctx.model_bytes() + size = len(data) + self.ctx.logger.info("加载模型 %s（%d 字节, conf=%s）", + self.ctx.model_name, size, self.ctx.param("conf")) + self._net = sail.Engine(data, size, self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + # TODO: 按模型实际输入做前处理（可参考 `algpkg init --algorithm yolo26` 生成的完整实现） + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + + def infer(self, frame) -> AlgoResult: + """单帧推理。返回 AlgoResult（boxes 为 (N,6)：x1,y1,x2,y2,score,cls）。 + + TODO: 前处理 → self._net.process(...) → 后处理（NMS/反变换）， + 再 return AlgoResult(boxes=boxes, labels=[...])。 + 下面这行只做"链路可跑通"的占位（空结果），替换为真实实现。 + """ + _ = frame + boxes = np.zeros((0, 6), dtype=np.float32) + return AlgoResult(boxes=boxes) + + def close(self) -> None: + self._net = None # type: ignore[assignment] +''' + + +_YOLO26_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +算法模板：YOLO26 单模型检测（`algpkg init --algorithm yolo26` 生成）。 +模型：`models/__MODEL_FILE__`（端到端输出，无 NMS；每行 x1,y1,x2,y2,score,cls）。 +只依赖 `video_algpkg_api` + 设备框架（sophon.sail / numpy）；放好模型即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import time + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + + +class Algorithm(AlgorithmBase): + """YOLO26 单 bmodel 检测器（SAIL 引擎 + bmcv 硬件前处理）。""" + + kind = "__ALGO_KIND__" + + #: 引擎/输入张量：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 + #: （不必判空）；close() 置 None 与声明冲突，那两处显式 `type: ignore`。 + _net: sail.Engine + _input_tensor: sail.Tensor + + # ------------------------------------------------------------------ 生命周期 + + def load(self) -> None: + self._conf = float(self.ctx.param("conf", 0.25)) + data = self.ctx.model_bytes() + self._bytes = data # 引擎构造后仍持有引用（不依赖 SDK 拷贝语义） + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + self._net = sail.Engine(data, len(data), self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + + name = self._net.get_input_names(self._graph)[0] + self._input_name = name + self._input_dtype = self._net.get_input_dtype(self._graph, name) + self._img_dtype = self._bmcv.get_bm_image_data_format(self._input_dtype) + self._input_scale = self._net.get_input_scale(self._graph, name) + self._input_shape = self._net.get_input_shape(self._graph, name) + self._input_shapes = {name: self._input_shape} + if self._input_shape[0] != 1: + raise RuntimeError("本模板按 batch=1 单帧推理实现；多 batch 模型请改 infer()") + self._net_h, self._net_w = self._input_shape[2], self._input_shape[3] + + self._output_names = self._net.get_output_names(self._graph) + self._output_tensors = {} + for out in self._output_names: + shape = self._net.get_output_shape(self._graph, out) + dtype = self._net.get_output_dtype(self._graph, out) + self._output_tensors[out] = sail.Tensor(self._handle, shape, dtype, True, True) + + # scale/255，均值 0（与训练一致的 letterbox 归一化） + self._ab = [v * self._input_scale / 255.0 for v in (1, 0, 1, 0, 1, 0)] + self._input_tensor = sail.Tensor(self._handle, self._input_shape, + self._input_dtype, False, False) + self.ctx.logger.info("[%s] 模型就绪（%.1f MB, %.0fx%.0f, conf=%.2f）", + self.ctx.algo_id, len(data) / 1048576.0, + self._net_w, self._net_h, self._conf) + + def infer(self, frame) -> AlgoResult: + started = time.perf_counter() + pre, ratio, tx, ty = self._preprocess(frame) + self._bmcv.bm_image_to_tensor(pre, self._input_tensor) + self._net.process(self._graph, {self._input_name: self._input_tensor}, + self._input_shapes, self._output_tensors) + preds = self._output_tensors[self._output_names[0]].asnumpy() + batch = preds[0] if preds.ndim == 3 else preds + boxes = self._postprocess(batch, frame.width(), frame.height(), ratio, tx, ty) + return AlgoResult(boxes=boxes, + meta={"infer_ms": round((time.perf_counter() - started) * 1000, 1)}) + + def close(self) -> None: + self._net = None # type: ignore[assignment] + self._output_tensors = {} + self._input_tensor = None # type: ignore[assignment] + self._bytes = b"" + + # ------------------------------------------------------------------ 内部 + + def _preprocess(self, bmimg): + """BGRA/NV12 → RGB_PLANAR → letterbox → 归一化 到网络输入。""" + rgb = sail.BMImage(self._handle, bmimg.height(), bmimg.width(), + sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE) + self._bmcv.convert_format(bmimg, rgb) + img_w, img_h = rgb.width(), rgb.height() + r = min(self._net_w / img_w, self._net_h / img_h) + tw, th = int(round(r * img_w)), int(round(r * img_h)) + tx, ty = (self._net_w - tw) / 2, (self._net_h - th) / 2 + attr = sail.PaddingAtrr() + attr.set_stx(int(round(tx - 0.1))) + attr.set_sty(int(round(ty - 0.1))) + attr.set_w(tw) + attr.set_h(th) + attr.set_r(114) + attr.set_g(114) + attr.set_b(114) + resized = self._bmcv.crop_and_resize_padding( + rgb, 0, 0, img_w, img_h, self._net_w, self._net_h, attr, + sail.bmcv_resize_algorithm.BMCV_INTER_LINEAR, + ) + pre = sail.BMImage(self._handle, self._net_h, self._net_w, + sail.Format.FORMAT_RGB_PLANAR, self._img_dtype) + self._bmcv.convert_to( + resized, pre, + ((self._ab[0], self._ab[1]), (self._ab[2], self._ab[3]), + (self._ab[4], self._ab[5])), + ) + return pre, r, tx, ty + + def _postprocess(self, batch, org_w: int, org_h: int, ratio: float, + tx: float, ty: float) -> np.ndarray: + """置信度过滤（模型端到端输出，无 NMS）+ 反 letterbox 到原图坐标。""" + if batch is None or batch.size == 0: + return np.zeros((0, 6), dtype=np.float32) + rows = batch[batch[:, 4] > self._conf] + if not len(rows): + return np.zeros((0, 6), dtype=np.float32) + boxes = np.asarray(rows, dtype=np.float32).copy() + boxes[:, [0, 2]] -= tx + boxes[:, [1, 3]] -= ty + boxes[:, [0, 2]] /= ratio + boxes[:, [1, 3]] /= ratio + boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, org_w - 1) + boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, org_h - 1) + return boxes[:, :6] +''' + + +#: ImageNet-1k 类别名占位（第 0 行 = 类别 0）。分类模型换成自己的类别顺序即可。 +RESNET_LABELS = ( + "# ImageNet-1k 类别名（第 0 行 = 类别 0）\n" + "# 当前为占位名：请按训练/模型输出顺序替换为真实类别名\n" + + "".join(f"class_{i}\n" for i in range(1000)) +) + + +_ARCFACE_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +算法模板：ArcFace 人脸特征提取（`algpkg init --algorithm arcface` 生成）。 +预处理：任意格式 → RGB_PLANAR → resize 到 112×112 → 模型 dtype； +均值/缩放由 bmodel 内部完成（与 Sophon-DEMO 一致）。 +后处理：L2 归一化，输出 512 维 embedding。 + +embedding 放在 `AlgoResult.meta["embedding"]`（列表）；`counts={"face": 1}` +表示本帧产出了一条特征。当前 worker 的绘制/告警仍以框为主，特征供人脸比对/ +检索等上层逻辑使用。输入按**整帧** resize（不做人脸检测/对齐），实际使用时 +应先裁出人脸区域。 +只依赖 `video_algpkg_api` + 设备框架（sophon.sail / numpy）；放好模型即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import time + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + +#: bmcv convert_to 的 (a, b)：只做格式/精度转换；ArcFace bmodel 内部已含 mean/scale +_AB = ((1.0, 0.0), (1.0, 0.0), (1.0, 0.0)) + + +class Algorithm(AlgorithmBase): + """ArcFace 单模型人脸特征提取器（SAIL 引擎 + bmcv 前处理）。""" + + kind = "__ALGO_KIND__" + + #: 引擎/输入输出张量：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 + #: （不必判空）；close() 置 None 与声明冲突，那两处显式 `type: ignore`。 + _net: sail.Engine + _input_tensor: sail.Tensor + + # ------------------------------------------------------------------ 生命周期 + + def load(self) -> None: + data = self.ctx.model_bytes() + self._bytes = data # 引擎构造后仍持有引用（不依赖 SDK 拷贝语义） + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + self._net = sail.Engine(data, len(data), self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + + self._input_name = self._net.get_input_names(self._graph)[0] + self._input_dtype = self._net.get_input_dtype(self._graph, self._input_name) + self._img_dtype = self._bmcv.get_bm_image_data_format(self._input_dtype) + self._input_shape = self._net.get_input_shape(self._graph, self._input_name) + self._input_shapes = {self._input_name: self._input_shape} + if self._input_shape[0] != 1: + raise RuntimeError("本模板按 batch=1 单帧推理实现；多 batch 模型请改 infer()") + self._net_h, self._net_w = self._input_shape[2], self._input_shape[3] + + output_names = self._net.get_output_names(self._graph) + if len(output_names) != 1: + raise RuntimeError(f"ArcFace 模板要求单输出模型（收到 {len(output_names)} 个）") + self._output_name = output_names[0] + output_shape = self._net.get_output_shape(self._graph, self._output_name) + output_dtype = self._net.get_output_dtype(self._graph, self._output_name) + self._output_tensor = sail.Tensor(self._handle, output_shape, output_dtype, True, True) + + self._input_tensor = sail.Tensor(self._handle, self._input_shape, + self._input_dtype, False, False) + self.ctx.logger.info("[%s] 模型就绪（%.1f MB, %.0fx%.0f, output=%s）", + self.ctx.algo_id, len(data) / 1048576.0, + self._net_w, self._net_h, output_shape) + + def infer(self, frame) -> AlgoResult: + started = time.perf_counter() + pre = self._preprocess(frame) + self._bmcv.bm_image_to_tensor(pre, self._input_tensor) + self._net.process(self._graph, {self._input_name: self._input_tensor}, + self._input_shapes, {self._output_name: self._output_tensor}) + embedding = np.asarray(self._output_tensor.asnumpy(), dtype=np.float32).reshape(1, -1) + # 与 Sophon-DEMO 一致：按行做 L2 归一化 + norm = np.linalg.norm(embedding, axis=1, keepdims=True) + embedding = embedding / (norm + 1e-10) + vector = embedding.reshape(-1) + return AlgoResult( + counts={"face": 1}, value=1.0, + meta={"embedding": vector.tolist(), "dim": int(vector.shape[0]), + "norm": round(float(np.linalg.norm(vector)), 6), + "infer_ms": round((time.perf_counter() - started) * 1000, 1)}, + ) + + def close(self) -> None: + self._net = None # type: ignore[assignment] + self._output_tensor = None # type: ignore[assignment] + self._input_tensor = None # type: ignore[assignment] + self._bytes = b"" + + # ------------------------------------------------------------------ 内部 + + def _preprocess(self, bmimg): + """任意格式 → RGB_PLANAR → resize 到网络输入 → 模型 dtype。""" + rgb = sail.BMImage(self._handle, bmimg.height(), bmimg.width(), + sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE) + self._bmcv.convert_format(bmimg, rgb) + resized = self._bmcv.resize(rgb, self._net_w, self._net_h) + pre = sail.BMImage(self._handle, self._net_h, self._net_w, + sail.Format.FORMAT_RGB_PLANAR, self._img_dtype) + self._bmcv.convert_to(resized, pre, _AB) + return pre +''' + + +#: LPRNet 字符表（与 Sophon-DEMO 一致；最后一项是 CTC blank） +LPR_CHARS = ( + "京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", + "苏", "浙", "皖", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", + "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", + "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", + "A", "B", "C", "D", "E", "F", "G", "H", "J", "K", "L", "M", + "N", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "I", "O", + "-", +) + +LPR_CHARS_TXT = ( + "# LPRNet 字符表：顺序必须与模型输出一致；最后一行 '-' 是 CTC blank\n" + + "\n".join(LPR_CHARS) + "\n" +) + + +_LPRNET_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +算法模板：LPRNet 车牌识别（`algpkg init --algorithm lprnet` 生成）。 +预处理：任意格式 → BGR_PLANAR → resize 到模型输入 → (x - 127.5) / 128。 +后处理：argmax(axis=1) → CTC 贪心解码（去重、去 blank），得到车牌字符串。 + +结果放在 `AlgoResult.meta["text"]`（车牌号），`meta["chars"]` 是字符索引； +`counts={"plate": 1}` 表示本帧识别到一块车牌。当前 worker 的绘制/告警仍以框 +为主，文本供上层业务逻辑使用。输入应是**裁好的车牌图**（不要整帧车辆图）。 +只依赖 `video_algpkg_api` + 设备框架（sophon.sail / numpy）；放好模型即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import time + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + + +class Algorithm(AlgorithmBase): + """LPRNet 单模型车牌识别器（SAIL 引擎 + bmcv 前处理 + CTC 贪心解码）。""" + + kind = "__ALGO_KIND__" + + #: 引擎/输入输出张量：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 + #: （不必判空）；close() 置 None 与声明冲突，那两处显式 `type: ignore`。 + _net: sail.Engine + _input_tensor: sail.Tensor + + # ------------------------------------------------------------------ 生命周期 + + def load(self) -> None: + data = self.ctx.model_bytes() + self._bytes = data # 引擎构造后仍持有引用（不依赖 SDK 拷贝语义） + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + self._net = sail.Engine(data, len(data), self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + + self._input_name = self._net.get_input_names(self._graph)[0] + self._input_dtype = self._net.get_input_dtype(self._graph, self._input_name) + self._img_dtype = self._bmcv.get_bm_image_data_format(self._input_dtype) + self._input_scale = self._net.get_input_scale(self._graph, self._input_name) + self._input_shape = self._net.get_input_shape(self._graph, self._input_name) + self._input_shapes = {self._input_name: self._input_shape} + if self._input_shape[0] != 1: + raise RuntimeError("本模板按 batch=1 单帧推理实现；多 batch 模型请改 infer()") + self._net_h, self._net_w = self._input_shape[2], self._input_shape[3] + + self._output_name = self._net.get_output_names(self._graph)[0] + output_shape = self._net.get_output_shape(self._graph, self._output_name) + output_dtype = self._net.get_output_dtype(self._graph, self._output_name) + self._output_tensor = sail.Tensor(self._handle, output_shape, output_dtype, True, True) + + # 与 Sophon-DEMO 一致：BGR 像素 (x - 127.5) / 128 + self._ab = [x * self._input_scale * 0.0078125 + for x in (1, -127.5, 1, -127.5, 1, -127.5)] + + raw = self.ctx.read_asset("assets/chars.txt").decode("utf-8", errors="ignore") + self._chars = [line.strip() for line in raw.splitlines() + if line.strip() and not line.startswith("#")] + if not self._chars: + raise RuntimeError("assets/chars.txt 为空（每行一个字符，最后一行必须是 CTC blank '-'）") + self._blank = len(self._chars) - 1 + + self._input_tensor = sail.Tensor(self._handle, self._input_shape, + self._input_dtype, False, False) + self.ctx.logger.info("[%s] 模型就绪（%.1f MB, %.0fx%.0f, %d 字符）", + self.ctx.algo_id, len(data) / 1048576.0, + self._net_w, self._net_h, len(self._chars)) + + def infer(self, frame) -> AlgoResult: + started = time.perf_counter() + pre = self._preprocess(frame) + self._bmcv.bm_image_to_tensor(pre, self._input_tensor) + self._net.process(self._graph, {self._input_name: self._input_tensor}, + self._input_shapes, {self._output_name: self._output_tensor}) + text, indices = self._postprocess(self._output_tensor.asnumpy()) + return AlgoResult( + counts={"plate": 1} if text else {}, + value=1.0 if text else 0.0, + meta={"text": text, "chars": indices, + "infer_ms": round((time.perf_counter() - started) * 1000, 1)}, + ) + + def close(self) -> None: + self._net = None # type: ignore[assignment] + self._output_tensor = None # type: ignore[assignment] + self._input_tensor = None # type: ignore[assignment] + self._bytes = b"" + + # ------------------------------------------------------------------ 内部 + + def _preprocess(self, bmimg): + """任意格式 → BGR_PLANAR → resize → (x-127.5)/128 到模型输入。""" + bgr = sail.BMImage(self._handle, bmimg.height(), bmimg.width(), + sail.Format.FORMAT_BGR_PLANAR, bmimg.dtype()) + self._bmcv.convert_format(bmimg, bgr) + if bgr.width() != self._net_w or bgr.height() != self._net_h: + bgr = self._bmcv.resize(bgr, self._net_w, self._net_h) + pre = sail.BMImage(self._handle, self._net_h, self._net_w, + sail.Format.FORMAT_BGR_PLANAR, self._img_dtype) + self._bmcv.convert_to( + bgr, pre, + ((self._ab[0], self._ab[1]), (self._ab[2], self._ab[3]), + (self._ab[4], self._ab[5])), + ) + return pre + + def _postprocess(self, outputs) -> tuple[str, list[int]]: + """[N, C, T] logits → argmax(axis=1) → CTC 贪心解码。 + + BM1688 的 LPRNet 输出是 ``[1, 68, T]``（类别轴 = 1）；若模型是 + ``[N, T, C]`` 布局，把 axis 改成 2。 + """ + predictions = np.argmax(np.asarray(outputs), axis=1) + sequence = np.asarray(predictions).reshape(-1).tolist() + chars: list[str] = [] + indices: list[int] = [] + prev = None + for idx in sequence: + idx = int(idx) + if idx == self._blank: + prev = None + continue + if idx != prev and 0 <= idx < len(self._chars): + chars.append(self._chars[idx]) + indices.append(idx) + prev = idx + return "".join(chars), indices +''' + + +_FACE_RECOGNITION_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +算法模板：人脸检测 + 识别（`algpkg init --algorithm face-recognition` 生成）。 +- 检测：RetinaFace（640×640，BGR mean=[104,117,123]）； +- 识别：按检测框（外扩 `face_margin`）裁脸 → ArcFace → 512 维 L2 归一化 embedding。 + +结果：`AlgoResult.boxes` 是人脸框 (N,6)（cls=0=face）； +`meta["embeddings"]` 是与 boxes 逐项对应的 512 维特征列表（失败项为 None）。 +当前 worker 的绘制/告警使用 boxes；特征供上层人脸比对/检索逻辑使用。 +只依赖 `video_algpkg_api` + 设备框架（sophon.sail / numpy）；放好两个模型即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import time +from itertools import product +from math import ceil + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + +#: RetinaFace mobile0.25 先验框配置（与 Sophon-DEMO data/config.py 一致） +_MIN_SIZES = ((16.0, 32.0), (64.0, 128.0), (256.0, 512.0)) +_STEPS = (8.0, 16.0, 32.0) +_VARIANCE = (0.1, 0.2) + +#: ArcFace bmcv convert_to 系数：只做格式/精度转换（均值/缩放由 bmodel 内部完成） +_ARC_AB = ((1.0, 0.0), (1.0, 0.0), (1.0, 0.0)) + + +def _priorbox(net_h: int, net_w: int) -> np.ndarray: + """RetinaFace 先验框（与官方 PriorBox.forward 等价，clip=False）。""" + anchors = [] + feature_maps = [[ceil(net_h / step), ceil(net_w / step)] for step in _STEPS] + for k, (fh, fw) in enumerate(feature_maps): + for i, j in product(range(fh), range(fw)): + for min_size in _MIN_SIZES[k]: + s_kx = min_size / net_w + s_ky = min_size / net_h + cx = (j + 0.5) * _STEPS[k] / net_w + cy = (i + 0.5) * _STEPS[k] / net_h + anchors.append([cx, cy, s_kx, s_ky]) + return np.asarray(anchors, dtype=np.float32) + + +def _decode(loc: np.ndarray, priors: np.ndarray) -> np.ndarray: + boxes = np.concatenate(( + priors[:, :2] + loc[:, :2] * _VARIANCE[0] * priors[:, 2:], + priors[:, 2:] * np.exp(loc[:, 2:] * _VARIANCE[1]), + ), axis=1) + boxes[:, :2] -= boxes[:, 2:] / 2 + boxes[:, 2:] += boxes[:, :2] + return boxes + + +def _nms(dets: np.ndarray, thresh: float) -> list[int]: + """RetinaFace 官方 py_cpu_nms（+1 面积口径）。""" + x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4] + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + keep = [] + while order.size > 0: + i = order[0] + keep.append(int(i)) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + inds = np.where(ovr <= thresh)[0] + order = order[inds + 1] + return keep + + +class Algorithm(AlgorithmBase): + """RetinaFace 检测 + ArcFace 特征提取（双模型包）。""" + + kind = "__ALGO_KIND__" + + _net: sail.Engine + _arc_net: sail.Engine + _input_tensor: sail.Tensor + _arc_input_tensor: sail.Tensor + + # ------------------------------------------------------------------ 生命周期 + + def load(self) -> None: + self._conf = float(self.ctx.param("conf", 0.5)) + self._nms = float(self.ctx.param("nms", 0.3)) + self._margin = float(self.ctx.param("face_margin", 0.15)) + self._min_face = int(self.ctx.param("min_face", 16)) + self._max_faces = int(self.ctx.param("max_faces", 20)) + + det_data = self.ctx.model_bytes() # 默认模型：人脸检测 + arc_data = self.ctx.model_bytes("arcface") # 第二份模型：ArcFace + self._det_bytes = det_data + self._arc_bytes = arc_data + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + + # ---- 检测模型（RetinaFace） + self._net = sail.Engine(det_data, len(det_data), self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + self._input_name = self._net.get_input_names(self._graph)[0] + self._input_dtype = self._net.get_input_dtype(self._graph, self._input_name) + self._img_dtype = self._bmcv.get_bm_image_data_format(self._input_dtype) + self._input_scale = self._net.get_input_scale(self._graph, self._input_name) + self._input_shape = self._net.get_input_shape(self._graph, self._input_name) + self._input_shapes = {self._input_name: self._input_shape} + if self._input_shape[0] != 1: + raise RuntimeError("人脸检测模板按 batch=1 单帧推理实现") + self._det_h, self._det_w = self._input_shape[2], self._input_shape[3] + self._det_ab = [v * self._input_scale for v in (1, -104, 1, -117, 1, -123)] + self._output_names = self._net.get_output_names(self._graph) + self._output_tensors = {} + for name in self._output_names: + shape = self._net.get_output_shape(self._graph, name) + dtype = self._net.get_output_dtype(self._graph, name) + self._output_tensors[name] = sail.Tensor(self._handle, shape, dtype, True, True) + self._priors = _priorbox(self._det_h, self._det_w) + loc_rows = next( + (self._net.get_output_shape(self._graph, n)[1] for n in self._output_names + if self._net.get_output_shape(self._graph, n)[-1] == 4), None) + if loc_rows is not None and int(loc_rows) != int(self._priors.shape[0]): + raise RuntimeError( + f"先验框数量 {self._priors.shape[0]} 与模型 loc 输出 {loc_rows} 不一致") + self._input_tensor = sail.Tensor(self._handle, self._input_shape, + self._input_dtype, False, False) + + # ---- 识别模型（ArcFace） + self._arc_net = sail.Engine(arc_data, len(arc_data), self.ctx.dev_id, sail.IOMode.SYSO) + self._arc_graph = self._arc_net.get_graph_names()[0] + self._arc_input_name = self._arc_net.get_input_names(self._arc_graph)[0] + self._arc_input_dtype = self._arc_net.get_input_dtype( + self._arc_graph, self._arc_input_name) + self._arc_img_dtype = self._bmcv.get_bm_image_data_format(self._arc_input_dtype) + self._arc_input_shape = self._arc_net.get_input_shape( + self._arc_graph, self._arc_input_name) + self._arc_input_shapes = {self._arc_input_name: self._arc_input_shape} + if self._arc_input_shape[0] != 1: + raise RuntimeError("ArcFace 模板按 batch=1 单帧推理实现") + self._arc_h, self._arc_w = self._arc_input_shape[2], self._arc_input_shape[3] + self._arc_output_name = self._arc_net.get_output_names(self._arc_graph)[0] + arc_shape = self._arc_net.get_output_shape(self._arc_graph, self._arc_output_name) + arc_dtype = self._arc_net.get_output_dtype(self._arc_graph, self._arc_output_name) + self._arc_output_tensor = sail.Tensor(self._handle, arc_shape, arc_dtype, True, True) + self._arc_input_tensor = sail.Tensor(self._handle, self._arc_input_shape, + self._arc_input_dtype, False, False) + + self.ctx.logger.info( + "[%s] 检测模型 %.1f MB（%.0fx%.0f），识别模型 %.1f MB（%s）", + self.ctx.algo_id, len(det_data) / 1048576.0, self._det_w, self._det_h, + len(arc_data) / 1048576.0, arc_shape) + + def infer(self, frame) -> AlgoResult: + started = time.perf_counter() + bgr = self._to_bgr(frame) + det_input, ratio = self._det_preprocess(bgr) + self._bmcv.bm_image_to_tensor(det_input, self._input_tensor) + self._net.process(self._graph, {self._input_name: self._input_tensor}, + self._input_shapes, self._output_tensors) + outputs = [self._output_tensors[name].asnumpy() for name in self._output_names] + dets = self._detect(outputs, ratio, frame.width(), frame.height()) + + embeddings = [] + for det in dets: + fw, fh = float(det[2] - det[0]), float(det[3] - det[1]) + if fw < self._min_face or fh < self._min_face: + embeddings.append(None) # 过小人脸：保留框，不提取特征 + continue + try: + embeddings.append(self._embed(bgr, det, frame.width(), frame.height()).tolist()) + except Exception as exc: # 单张脸失败不拖垮整帧 + self.ctx.logger.warning("[%s] 人脸特征提取失败: %s", self.ctx.algo_id, exc) + embeddings.append(None) + + boxes = (np.asarray(dets[:, :6], dtype=np.float32) if len(dets) + else np.zeros((0, 6), dtype=np.float32)) + count = int(boxes.shape[0]) + return AlgoResult( + boxes=boxes, + counts={"face": count} if count else {}, + value=float(count), + labels=["face"], + meta={"embeddings": embeddings, "embedding_dim": 512 if count else 0, + "faces": count, + "embedded": sum(1 for e in embeddings if e is not None), + "infer_ms": round((time.perf_counter() - started) * 1000, 1)}, + ) + + def close(self) -> None: + self._net = None # type: ignore[assignment] + self._arc_net = None # type: ignore[assignment] + self._output_tensors = {} + self._input_tensor = None # type: ignore[assignment] + self._arc_input_tensor = None # type: ignore[assignment] + self._arc_output_tensor = None # type: ignore[assignment] + self._det_bytes = b"" + self._arc_bytes = b"" + + # ------------------------------------------------------------------ 前处理/检测 + + def _to_bgr(self, frame): + """原图 → BGR_PLANAR（检测与裁脸共用，避免重复转换）。""" + fmt = frame.format() + if fmt == sail.Format.FORMAT_BGR_PLANAR: + return frame + if fmt == sail.Format.FORMAT_YUV420P: + return self._bmcv.yuv2bgr(frame) + bgr = sail.BMImage(self._handle, frame.height(), frame.width(), + sail.Format.FORMAT_BGR_PLANAR, frame.dtype()) + self._bmcv.convert_format(frame, bgr) + return bgr + + def _det_preprocess(self, bgr): + """BGR → letterbox 到网络输入 → 减均值（与 Sophon-DEMO 一致）。""" + img_w, img_h = bgr.width(), bgr.height() + r_w = self._det_w / img_w + r_h = self._det_h / img_h + if r_h > r_w: + tw, th = self._det_w, int(r_w * img_h) + tx, ty = 0, 0 + else: + tw, th = int(r_h * img_w), self._det_h + tx, ty = 0, 0 + attr = sail.PaddingAtrr() + attr.set_stx(tx) + attr.set_sty(ty) + attr.set_w(tw) + attr.set_h(th) + attr.set_r(0) + attr.set_g(0) + attr.set_b(0) + resized = self._bmcv.vpp_crop_and_resize_padding( + bgr, 0, 0, img_w, img_h, self._det_w, self._det_h, attr) + pre = sail.BMImage(self._handle, self._det_h, self._det_w, + sail.Format.FORMAT_BGR_PLANAR, self._img_dtype) + self._bmcv.convert_to( + resized, pre, + ((self._det_ab[0], self._det_ab[1]), (self._det_ab[2], self._det_ab[3]), + (self._det_ab[4], self._det_ab[5])), + ) + return pre, (r_w if r_h > r_w else r_h) + + def _detect(self, outputs, resize, org_w, org_h) -> np.ndarray: + """3 个输出按最后一维识别 loc/conf → decode → 阈值 → NMS。""" + loc = conf = None + for out in outputs: + if out.shape[-1] == 4: + loc = out + elif out.shape[-1] == 2: + conf = out + if loc is None or conf is None: + raise RuntimeError("RetinaFace 输出缺少 loc(4) 或 conf(2)") + + boxes = _decode(loc.reshape(-1, 4), self._priors) + boxes = boxes * np.asarray([self._det_w, self._det_h, + self._det_w, self._det_h]) / resize + scores = conf.reshape(-1, 2)[:, 1] + mask = scores >= self._conf + boxes, scores = boxes[mask], scores[mask] + if not len(boxes): + return np.zeros((0, 6), dtype=np.float32) + order = scores.argsort()[::-1] + boxes, scores = boxes[order], scores[order] + dets = np.hstack((boxes, scores[:, None])).astype(np.float32) + dets = dets[_nms(dets, self._nms)] + if len(dets) > self._max_faces: + dets = dets[:self._max_faces] + dets[:, [0, 2]] = dets[:, [0, 2]].clip(0, org_w - 1) + dets[:, [1, 3]] = dets[:, [1, 3]].clip(0, org_h - 1) + cls = np.zeros((len(dets), 1), dtype=np.float32) + return np.hstack((dets, cls)) + + # ------------------------------------------------------------------ 识别 + + def _embed(self, bgr, det, org_w, org_h): + """按检测框外扩 margin 裁脸 → ArcFace → L2 归一化。""" + x1, y1, x2, y2 = (float(v) for v in det[:4]) + w, h = x2 - x1, y2 - y1 + mx, my = int(w * self._margin), int(h * self._margin) + cx1 = max(0, int(x1) - mx) & ~1 # VPSS 要求裁剪宽高为偶数 + cy1 = max(0, int(y1) - my) & ~1 + cx2 = min(org_w, int(x2) + mx) + cy2 = min(org_h, int(y2) + my) + cw = (cx2 - cx1) & ~1 + ch = (cy2 - cy1) & ~1 + if cw < 2 or ch < 2: + raise RuntimeError(f"人脸框过小: {cw}x{ch}") + + crop = self._bmcv.crop(bgr, cx1, cy1, cw, ch) + rgb = sail.BMImage(self._handle, crop.height(), crop.width(), + sail.Format.FORMAT_RGB_PLANAR, sail.DATA_TYPE_EXT_1N_BYTE) + self._bmcv.convert_format(crop, rgb) + resized = self._bmcv.resize(rgb, self._arc_w, self._arc_h) + pre = sail.BMImage(self._handle, self._arc_h, self._arc_w, + sail.Format.FORMAT_RGB_PLANAR, self._arc_img_dtype) + self._bmcv.convert_to(resized, pre, _ARC_AB) + + self._bmcv.bm_image_to_tensor(pre, self._arc_input_tensor) + self._arc_net.process(self._arc_graph, {self._arc_input_name: self._arc_input_tensor}, + self._arc_input_shapes, + {self._arc_output_name: self._arc_output_tensor}) + emb = np.asarray(self._arc_output_tensor.asnumpy(), dtype=np.float32).reshape(1, -1) + norm = np.linalg.norm(emb, axis=1, keepdims=True) + return (emb / (norm + 1e-10)).reshape(-1) +''' + + +_RESNET_CODE = '''"""算法包实现：__ALGO_NAME__（__ALGO_ID__）。 + +算法模板：ResNet 图像分类（`algpkg init --algorithm resnet` 生成）。 +预处理：resize 到网络输入 → RGB_PLANAR → ImageNet mean/std 归一化； +后处理：softmax → top-1（类别索引 + 置信度）。 + +worker 的结果通道是 (N,6) 框：分类结果用**整帧框**承载 +（score=置信度，cls=类别索引），绘制/告警/统计无需特殊处理。 +只依赖 `video_algpkg_api` + 设备框架（sophon.sail / numpy）；放好模型即可打包： + algpkg build <本目录> --out __ALGO_ID__-__ALGO_VERSION__.mbank +""" +from __future__ import annotations + +import time + +import numpy as np +import sophon.sail as sail + +from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult + +#: ImageNet 归一化（与训练一致） +_MEAN = (0.485, 0.456, 0.406) +_STD = (0.229, 0.224, 0.225) + + +class Algorithm(AlgorithmBase): + """ResNet 单模型图像分类器（SAIL 引擎 + bmcv 前处理）。""" + + kind = "__ALGO_KIND__" + + #: 引擎/输入输出张量：load() 赋值、close() 释放。类级声明为非 Optional，infer() 里可直接用 + #: （不必判空）；close() 置 None 与声明冲突，那两处显式 `type: ignore`。 + _net: sail.Engine + _input_tensor: sail.Tensor + + # ------------------------------------------------------------------ 生命周期 + + def load(self) -> None: + self._conf = float(self.ctx.param("conf", 0.1)) + data = self.ctx.model_bytes() + self._bytes = data # 引擎构造后仍持有引用（不依赖 SDK 拷贝语义） + self._handle = sail.Handle(self.ctx.dev_id) + self._bmcv = sail.Bmcv(self._handle) + self._net = sail.Engine(data, len(data), self.ctx.dev_id, sail.IOMode.SYSO) + self._graph = self._net.get_graph_names()[0] + + self._input_name = self._net.get_input_names(self._graph)[0] + self._input_dtype = self._net.get_input_dtype(self._graph, self._input_name) + self._img_dtype = self._bmcv.get_bm_image_data_format(self._input_dtype) + self._input_scale = self._net.get_input_scale(self._graph, self._input_name) + self._input_shape = self._net.get_input_shape(self._graph, self._input_name) + if self._input_shape[0] != 1: + raise RuntimeError("本模板按 batch=1 单帧推理实现；多 batch 模型请改 infer()") + self._net_h, self._net_w = self._input_shape[2], self._input_shape[3] + + self._output_name = self._net.get_output_names(self._graph)[0] + output_shape = self._net.get_output_shape(self._graph, self._output_name) + output_dtype = self._net.get_output_dtype(self._graph, self._output_name) + self._output_scale = self._net.get_output_scale(self._graph, self._output_name) + self._output_tensor = sail.Tensor(self._handle, output_shape, output_dtype, True, True) + self._output_tensors = {self._output_name: self._output_tensor} + + # ImageNet 归一化：x/255 → (x/255 - mean)/std，合并成 bmcv 的 a*x+b + a = [1.0 / (255.0 * v) for v in _STD] + b = [-m / s for m, s in zip(_MEAN, _STD)] + self._ab = [] + for i in range(3): + self._ab.append(a[i] * self._input_scale) + self._ab.append(b[i] * self._input_scale) + + self._labels = self._read_labels() + self._input_tensor = sail.Tensor(self._handle, self._input_shape, + self._input_dtype, False, False) + self.ctx.logger.info("[%s] 模型就绪（%.1f MB, %.0fx%.0f, conf=%.2f）", + self.ctx.algo_id, len(data) / 1048576.0, + self._net_w, self._net_h, self._conf) + + def infer(self, frame) -> AlgoResult: + started = time.perf_counter() + pre = self._preprocess(frame) + self._bmcv.bm_image_to_tensor(pre, self._input_tensor) + self._net.process(self._graph, {self._input_name: self._input_tensor}, + self._output_tensors) + cls, score = self._postprocess(self._output_tensor.asnumpy()) + + # worker 用 (N,6) 框渲染/告警：整帧框承载分类结果（score, cls） + if score >= self._conf: + boxes = np.array([[0.0, 0.0, float(frame.width() - 1), + float(frame.height() - 1), score, float(cls)]], + dtype=np.float32) + label = self._labels[cls] if 0 <= cls < len(self._labels) else str(cls) + counts = {label: 1} + else: + boxes = np.zeros((0, 6), dtype=np.float32) + counts = {} + return AlgoResult( + boxes=boxes, counts=counts, value=float(cls), + labels=self._labels or None, + meta={"class_id": cls, "score": round(score, 4), + "infer_ms": round((time.perf_counter() - started) * 1000, 1)}, + ) + + def close(self) -> None: + self._net = None # type: ignore[assignment] + self._output_tensors = {} + self._input_tensor = None # type: ignore[assignment] + self._bytes = b"" + + # ------------------------------------------------------------------ 内部 + + def _read_labels(self) -> list[str]: + """读包内类别表（缺省/缺失时用类别索引显示）。""" + try: + raw = self.ctx.read_asset("assets/labels.txt").decode("utf-8", errors="ignore") + except Exception: + return [] + return [line.strip() for line in raw.splitlines() + if line.strip() and not line.strip().startswith("#")] + + def _preprocess(self, bmimg): + """YUV420P/BGR → resize → RGB_PLANAR → ImageNet mean/std 归一化。""" + if bmimg.format() == sail.Format.FORMAT_YUV420P: + bgr = self._bmcv.yuv2bgr(bmimg) + else: + bgr = bmimg + resized = self._bmcv.resize(bgr, self._net_w, self._net_h, + sail.bmcv_resize_algorithm.BMCV_INTER_LINEAR) + rgb = sail.BMImage(self._handle, self._net_h, self._net_w, + sail.Format.FORMAT_RGB_PLANAR, resized.dtype()) + self._bmcv.convert_format(resized, rgb) + pre = sail.BMImage(self._handle, self._net_h, self._net_w, + sail.Format.FORMAT_RGB_PLANAR, self._img_dtype) + self._bmcv.convert_to( + rgb, pre, + ((self._ab[0], self._ab[1]), (self._ab[2], self._ab[3]), + (self._ab[4], self._ab[5])), + ) + return pre + + def _postprocess(self, outputs) -> tuple[int, float]: + """logits → softmax → top-1（类别索引, 置信度）。 + + 与 Sophon-DEMO 的 ResNet 一致，模型输出按 logits 处理；若 bmodel 已带 + softmax，把这里改成直接 `argmax` 即可。 + """ + logits = np.asarray(outputs, dtype=np.float32).reshape(-1) * self._output_scale + logits = logits - logits.max() + exp = np.exp(logits) + probs = exp / exp.sum() + cls = int(np.argmax(probs)) + return cls, float(probs[cls]) +''' + + +# ====================================================================== 模板注册 + +@dataclass(frozen=True) +class AlgorithmTemplate: + """一个算法模板：生成清单 + 入口代码 + 资源。""" + + key: str + label: str + kind: str + summary: str + model_name: str + model_file: str + precision: str = "" + resolution: tuple[int, int] = (640, 640) + batch: int = 1 + #: 额外模型（多模型包，如 人脸检测 + 特征提取）；每项含 name/file/precision/resolution/batch + extra_models: tuple[dict, ...] = () + input: dict = field(default_factory=lambda: dict(_YOLO_INPUT)) + output: dict = field(default_factory=lambda: { + "type": "boxes", "format": "N6", "classes_file": "assets/labels.txt"}) + params_schema: tuple[dict, ...] = _YOLO_PARAMS + defaults: dict = field(default_factory=dict) + labels: str = "# 每行一个类别名（第 0 行 = 类别 0）；按需替换\nobject\n" + #: `labels` 写到的包内路径（类别表/字符表；LPRNet 用 assets/chars.txt） + labels_file: str = "assets/labels.txt" + code: str = _GENERIC_CODE + aliases: tuple[str, ...] = () + devices: tuple[str, ...] = () + tags: tuple[str, ...] = () + icon: str = "mdi-puzzle" + #: 生成 README 里的模板注意事项（模型输出形态/类别表等） + notes: str = "" + + # ------------------------------------------------------------------ 渲染 + + def matches(self, name: str) -> bool: + wanted = str(name or "").strip().lower() + return wanted in (self.key, *self.aliases) + + def model_specs(self, algo_id: str) -> list[dict]: + """包内模型列表（第一项为默认模型；多模型包的其余项在 `extra_models`）。""" + def build(item: dict, default: bool) -> dict: + return { + "name": str(item["name"]), + "file": str(item["file"]).replace("__ALGO_ID__", algo_id), + "precision": str(item.get("precision") or ""), + "resolution": list(item.get("resolution") or self.resolution), + "batch": int(item.get("batch") or 1), + "default": default, + } + + primary = build({"name": self.model_name, "file": self.model_file, + "precision": self.precision, "resolution": self.resolution, + "batch": self.batch}, True) + return [primary, *[build(item, False) for item in self.extra_models]] + + def render_code(self, *, algo_id: str, name: str, version: str, kind: str, + code_package: str) -> str: + _ = code_package # 预留给"多文件模板"（入口拆分到子模块） + return (self.code + .replace("__ALGO_ID__", algo_id) + .replace("__ALGO_NAME__", name) + .replace("__ALGO_VERSION__", version) + .replace("__ALGO_KIND__", kind) + .replace("__MODEL_FILE__", self.model_file.rsplit("/", 1)[-1])) + + def render_manifest(self, *, algo_id: str, name: str, version: str, kind: str, + vendor: str, code_package: str) -> dict: + return { + "manifest_format": 1, + "id": algo_id, + "name": name, + "version": version, + "api_version": 1, + "kind": kind or self.kind, + "entry": "Algorithm", + "code_package": code_package, + "vendor": vendor or "", + "description": f"{name}（{self.label}，由 algpkg init --algorithm {self.key} 生成）", + "models": self.model_specs(algo_id), + "input": dict(self.input), + "output": dict(self.output), + "params_schema": [dict(item) for item in self.params_schema], + "defaults": dict(self.defaults), + "alarm": {"type_template": "{class}", "cooldown_key": "algo", + "need_snapshot": True, "need_record": True}, + "requires": {"license_features": [], "devices": list(self.devices), + "frameworks": ["sail"]}, + "ui": {"order": 100, "tags": list(self.tags), "icon": self.icon}, + "legacy": {}, + } + + +GENERIC = AlgorithmTemplate( + key="generic", + label="通用占位模板", + kind="detector", + summary="目录骨架 + 占位 load()/infer()，按自己的模型补前后处理", + model_name="default", + model_file="models/__ALGO_ID__.bmodel", + precision="", + params_schema=_YOLO_PARAMS, + defaults={"conf": 0.25, "interval": 0.2}, + aliases=("default", "basic"), +) + +YOLO26 = AlgorithmTemplate( + key="yolo26", + label="YOLO26 目标检测", + kind="detector", + summary="YOLO26 端到端检测（SAIL 硬件前处理 + N6 后处理），放好模型即可运行", + model_name="yolo26s_fp16_1b", + model_file="models/yolo26s_fp16_1b.bmodel", + precision="fp16", + resolution=(640, 640), + batch=1, + params_schema=(*_YOLO_PARAMS, { + "key": "draw", "type": "bool", "default": False, + "label": "绘制叠加", "desc": "在推理帧上画框（预览叠加流 / 告警快照）"}), + defaults={"conf": 0.25, "interval": 0.2, "draw": False}, + labels=COCO_LABELS, + code=_YOLO26_CODE, + aliases=("yolov26", "yolo26-detect", "yolov26-detect"), + devices=("BM1688", "BM1684X", "CV186X"), + tags=("YOLO26", "检测"), + icon="mdi-cctv", + notes="`infer()` 按 YOLO26 端到端输出 `(N,6)` 实现（模型端无 NMS）。" + "若换成原始输出（需要 NMS/解码）的模型，改 `_postprocess()`。", +) + +RESNET = AlgorithmTemplate( + key="resnet", + label="ResNet 图像分类", + kind="classifier", + summary="ResNet top-1 分类（resize + ImageNet 归一化 + softmax），结果以整帧框输出", + model_name="resnet50_fp32_1b", + model_file="models/resnet50_fp32_1b.bmodel", + precision="fp32", + resolution=(224, 224), + batch=1, + input={"width": 224, "height": 224, "color": "rgb", "layout": "planar", + "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]}, + params_schema=_RESNET_PARAMS, + defaults={"conf": 0.1, "interval": 0.5, "draw": False}, + labels=RESNET_LABELS, + code=_RESNET_CODE, + aliases=("resnet50", "resnet50-classify", "classifier"), + devices=("BM1688", "BM1684X", "CV186X"), + tags=("ResNet", "分类"), + icon="mdi-tag-multiple", + notes="`assets/labels.txt` 当前是 `class_0` … `class_999` 占位名：请按模型/训练" + "输出顺序替换为真实类别名（第 0 行 = 类别 0）。分类结果以**整帧框**返回" + "（score=置信度、cls=类别索引），便于现有绘制/告警链路直接使用。", +) + +ARCFACE = AlgorithmTemplate( + key="arcface", + label="ArcFace 人脸特征提取", + kind="classifier", + summary="ArcFace 512 维人脸特征（resize 112×112 + L2 归一化），embedding 放在结果 meta", + model_name="arcface_resnet50_fp32_1b", + model_file="models/arcface_resnet50_fp32_1b.bmodel", + precision="fp32", + resolution=(112, 112), + batch=1, + input={"width": 112, "height": 112, "color": "rgb", "layout": "planar", + "mean": [0, 0, 0], "std": [1, 1, 1]}, + output={"type": "embedding", "format": "vector", "embedding_size": 512}, + params_schema=_ARCFACE_PARAMS, + defaults={"interval": 0.5}, + labels="# ArcFace 输出 512 维特征向量，不需要类别表\n", + code=_ARCFACE_CODE, + aliases=("arc-face", "arcface-resnet50", "face-embedding", "embedding"), + devices=("BM1688", "BM1684X", "CV186X"), + tags=("ArcFace", "人脸"), + icon="mdi-face-recognition", + notes="输出为 512 维 **L2 归一化 embedding**，放在 `AlgoResult.meta[\"embedding\"]`" + "（`counts={\"face\": 1}` 表示本帧产出一条特征）。当前 worker 的绘制/告警" + "仍以框为主，embedding 供人脸比对/检索等上层逻辑使用。输入按整帧 resize 到 " + "112×112（与 Sophon-DEMO 一致，不做人脸检测/对齐）；实际视频场景应先裁脸。", +) + +LPRNET = AlgorithmTemplate( + key="lprnet", + label="LPRNet 车牌识别", + kind="classifier", + summary="LPRNet 车牌字符串识别（BGR resize + CTC 贪心解码），结果放在 meta['text']", + model_name="lprnet_fp32_1b", + model_file="models/lprnet_fp32_1b.bmodel", + precision="fp32", + resolution=(94, 24), + batch=1, + input={"width": 94, "height": 24, "color": "bgr", "layout": "planar", + "mean": [127.5, 127.5, 127.5], "std": [128, 128, 128]}, + output={"type": "text", "format": "string", "chars_file": "assets/chars.txt"}, + params_schema=_LPRNET_PARAMS, + defaults={"interval": 0.5}, + labels=LPR_CHARS_TXT, + labels_file="assets/chars.txt", + code=_LPRNET_CODE, + aliases=("lpr", "license-plate", "plate-recognition", "plate"), + devices=("BM1688", "BM1684X", "CV186X"), + tags=("LPRNet", "车牌"), + icon="mdi-car", + notes="结果放在 `AlgoResult.meta[\"text\"]`（车牌字符串），`meta[\"chars\"]` 为字符索引；" + "`counts={\"plate\": 1}` 表示识别到车牌。当前 worker 的绘制/告警仍以框为主，" + "文本供上层业务逻辑使用。输入应是**裁好的车牌图**（BM1688 模型为 24×94，" + "代码按引擎实际输入 resize）；`assets/chars.txt` 的顺序必须与模型一致，" + "最后一行 '-' 是 CTC blank。", +) + +FACE_RECOGNITION = AlgorithmTemplate( + key="face-recognition", + label="人脸检测 + 识别", + kind="detector", + summary="RetinaFace 检测 + ArcFace 512 维特征（双模型包，box 与 embedding 逐项对应）", + model_name="face_detector", + model_file="models/retinaface_mobilenet0.25_fp32_1b.bmodel", + precision="fp32", + resolution=(640, 640), + batch=1, + extra_models=( + {"name": "arcface", "file": "models/arcface_resnet50_fp32_1b.bmodel", + "precision": "fp32", "resolution": (112, 112), "batch": 1}, + ), + input={"width": 640, "height": 640, "color": "bgr", "layout": "planar", + "mean": [104, 117, 123], "std": [1, 1, 1]}, + output={"type": "boxes", "format": "N6", "classes_file": "assets/labels.txt", + "embedding_size": 512}, + params_schema=_FACE_RECOGNITION_PARAMS, + defaults={"conf": 0.5, "nms": 0.3, "face_margin": 0.15, "min_face": 16, + "max_faces": 20, "interval": 0.2, "draw": False}, + labels="face\n", + code=_FACE_RECOGNITION_CODE, + aliases=("face", "face-pipeline", "retinaface-arcface", "face-detect-recognize"), + devices=("BM1688", "BM1684X", "CV186X"), + tags=("RetinaFace", "ArcFace", "人脸"), + icon="mdi-face-recognition", + notes="双模型包：RetinaFace 检测 + ArcFace 识别。`boxes` 是人脸框，" + "`meta[\"embeddings\"]` 与 boxes 逐项对应（512 维 L2 归一化；失败或过小为 None）。" + "裁剪按检测框外扩 `face_margin`；宽/高小于 `min_face`（默认 16px，VPSS 裁剪下限）" + "的人脸只保留框、不提取特征。未做人脸关键点对齐；当前 worker 只绘制 boxes，" + "embedding 供上层比对/检索使用。", +) + +TEMPLATES: tuple[AlgorithmTemplate, ...] = ( + GENERIC, YOLO26, RESNET, ARCFACE, LPRNET, FACE_RECOGNITION) + + +def get_template(name: str = "") -> AlgorithmTemplate: + """按名字/别名查模板（空 = generic）；未知名字抛 :class:`ValueError`。""" + wanted = str(name or "").strip() or GENERIC.key + for template in TEMPLATES: + if template.matches(wanted): + return template + available = ", ".join(t.key for t in TEMPLATES) + raise ValueError(f"未知算法模板 {wanted!r}（可选: {available}）") + + +def template_summaries() -> list[tuple[str, str]]: + """`[(key, 说明)]`，给 CLI 帮助/错误提示用。""" + return [(t.key, t.summary) for t in TEMPLATES] + + +__all__ = ["AlgorithmTemplate", "COCO_LABELS", "RESNET_LABELS", "LPR_CHARS", + "LPR_CHARS_TXT", "GENERIC", "YOLO26", "RESNET", "ARCFACE", "LPRNET", + "FACE_RECOGNITION", "TEMPLATES", "get_template", "template_summaries"] diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 5fb56e2..b2dddbf 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -7,12 +7,18 @@ """ from __future__ import annotations +import hashlib +import http.server +import importlib.util import json import subprocess import sys import tempfile +import threading +import types from pathlib import Path +import numpy as np import pytest from algpkg_sdk import (AlgoBuilder, AlgoError, AlgoPackage, AlgoPackageError, Manifest, @@ -318,6 +324,436 @@ def test_cli_init_build_inspect_check(tmp_path: Path) -> None: assert cli_main(["list", str(tmp_path)]) == 0 +def test_cli_init_yolo26_template(tmp_path: Path) -> None: + """`--algorithm yolo26` 生成可直接运行的检测包（不是 TODO 占位）。""" + src = tmp_path / "yolo.demo" + assert cli_main(["init", str(src), "--id", "yolo.demo", "--name", "YOLO26 示例", + "--algorithm", "yolo26"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.id == "yolo.demo" and manifest.kind == "detector" + assert manifest.default_model().name == "yolo26s_fp16_1b" + assert manifest.default_model().file == "models/yolo26s_fp16_1b.bmodel" + assert manifest.output["format"] == "N6" + + code = (src / manifest.code_package / "__init__.py").read_text() + compile(code, "alg_yolo_demo/__init__.py", "exec") # 语法可编译（不 import 设备库） + assert "class Algorithm" in code and "_preprocess" in code and "_postprocess" in code + assert "TODO" not in code + + labels = [line for line in (src / "assets" / "labels.txt").read_text().splitlines() + if line.strip()] + assert len(labels) == 80 and labels[0] == "person" and labels[-1] == "toothbrush" + + model = tmp_path / "fake.bmodel" + model.write_bytes(b"FAKE-BMODEL" * 512) + out = tmp_path / "yolo.demo-1.0.0.mbank" + assert cli_main(["build", str(src), "--out", str(out), + "--model", f"yolo26s_fp16_1b={model}", "--quiet"]) == 0 + assert cli_main(["verify", str(out), "--json"]) == 0 + + +def test_cli_init_yolo26_alias_and_unknown_template(tmp_path: Path) -> None: + src = tmp_path / "alias.demo" + assert cli_main(["init", str(src), "--id", "alias.demo", "--name", "别名", + "--template", "yolov26"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.default_model().name == "yolo26s_fp16_1b" + + bad = tmp_path / "bad.demo" + assert cli_main(["init", str(bad), "--id", "bad.demo", "--name", "未知", + "--algorithm", "nope"]) == 2 + assert not (bad / "manifest.json").exists() + + +def test_cli_init_resnet_template(tmp_path: Path) -> None: + """`--algorithm resnet` 生成 ResNet 分类包（ImageNet 归一化 + softmax + 整帧框）。""" + src = tmp_path / "resnet.demo" + assert cli_main(["init", str(src), "--id", "resnet.demo", "--name", "ResNet 示例", + "--algorithm", "resnet"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.id == "resnet.demo" and manifest.kind == "classifier" + assert manifest.default_model().name == "resnet50_fp32_1b" + assert manifest.default_model().file == "models/resnet50_fp32_1b.bmodel" + assert manifest.input["mean"] == [0.485, 0.456, 0.406] + assert manifest.schema_defaults()["conf"] == 0.1 + + code = (src / manifest.code_package / "__init__.py").read_text() + compile(code, "alg_resnet_demo/__init__.py", "exec") + assert "class Algorithm" in code and "_preprocess" in code + assert "_postprocess" in code and "yuv2bgr" in code + + labels = [line for line in (src / "assets" / "labels.txt").read_text().splitlines() + if line.strip() and not line.startswith("#")] + assert len(labels) == 1000 and labels[0] == "class_0" and labels[-1] == "class_999" + + model = tmp_path / "fake.bmodel" + model.write_bytes(b"FAKE-BMODEL" * 512) + out = tmp_path / "resnet.demo-1.0.0.mbank" + assert cli_main(["build", str(src), "--out", str(out), + "--model", f"resnet50_fp32_1b={model}", "--quiet"]) == 0 + assert cli_main(["verify", str(out), "--json"]) == 0 + + +def test_resnet_template_infer_with_fake_sail(tmp_path: Path, monkeypatch) -> None: + """不依赖设备验证分类后处理：top-1 → 整帧框（score/cls）与阈值过滤。""" + src = tmp_path / "resnet.infer" + cli_main(["init", str(src), "--id", "resnet.infer", "--name", "推理示例", + "--algorithm", "resnet"]) + code_path = src / "alg_resnet_infer" / "__init__.py" + + # 生成代码只在方法体内使用 sail；空模块足以导入（不需要真实设备） + monkeypatch.setitem(sys.modules, "sophon", types.ModuleType("sophon")) + monkeypatch.setitem(sys.modules, "sophon.sail", types.ModuleType("sophon.sail")) + spec = importlib.util.spec_from_file_location("alg_resnet_infer_under_test", code_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + class Frame: + def width(self) -> int: + return 640 + + def height(self) -> int: + return 480 + + class Tensor: + def asnumpy(self): + return [[0.1, 3.0, 0.2]] # logits: top-1 = 类别 1 + + class Net: + def process(self, *args, **kwargs) -> None: + return None + + class Bmcv: + def bm_image_to_tensor(self, *args, **kwargs) -> None: + return None + + algo = module.Algorithm(ctx=None) + algo._conf = 0.5 + algo._labels = ["cat", "dog", "bird"] + algo._output_scale = 1.0 + algo._output_tensor = Tensor() + algo._output_tensors = {0: algo._output_tensor} + algo._input_tensor = object() + algo._input_name, algo._graph = "input", "graph" + algo._net = Net() + algo._bmcv = Bmcv() + algo._preprocess = lambda frame: object() + + result = algo.infer(Frame()) + assert result.boxes.shape == (1, 6) + assert (float(result.boxes[0, 0]), float(result.boxes[0, 2])) == (0.0, 639.0) + assert int(result.boxes[0, 5]) == 1 + assert float(result.boxes[0, 4]) > 0.8 + assert result.counts == {"dog": 1} and result.value == 1.0 + assert result.meta["class_id"] == 1 + + algo._conf = 0.99 # 低于阈值 → 空结果，不上报告警 + assert algo.infer(Frame()).boxes.shape == (0, 6) + + +def test_cli_init_arcface_template(tmp_path: Path) -> None: + """`--algorithm arcface` 生成 ArcFace 特征包（embedding 输出、无类别表）。""" + src = tmp_path / "arcface.demo" + assert cli_main(["init", str(src), "--id", "arcface.demo", "--name", "ArcFace 示例", + "--algorithm", "arcface"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.id == "arcface.demo" and manifest.kind == "classifier" + assert manifest.default_model().name == "arcface_resnet50_fp32_1b" + assert manifest.default_model().file == "models/arcface_resnet50_fp32_1b.bmodel" + assert manifest.output["type"] == "embedding" + assert manifest.output["embedding_size"] == 512 + assert manifest.schema_defaults()["interval"] == 0.5 + + code = (src / manifest.code_package / "__init__.py").read_text() + compile(code, "alg_arcface_demo/__init__.py", "exec") + assert "class Algorithm" in code and "_preprocess" in code + assert "meta" in code and "linalg.norm" in code + + labels = [line for line in (src / "assets" / "labels.txt").read_text().splitlines() + if line.strip() and not line.startswith("#")] + assert labels == [] # 特征包不需要类别表 + + model = tmp_path / "fake.bmodel" + model.write_bytes(b"FAKE-BMODEL" * 512) + out = tmp_path / "arcface.demo-1.0.0.mbank" + assert cli_main(["build", str(src), "--out", str(out), + "--model", f"arcface_resnet50_fp32_1b={model}", "--quiet"]) == 0 + assert cli_main(["verify", str(out), "--json"]) == 0 + + +def test_arcface_template_infer_with_fake_sail(tmp_path: Path, monkeypatch) -> None: + """不依赖设备验证特征提取：L2 归一化、512 维 embedding、face 计数。""" + src = tmp_path / "arcface.infer" + cli_main(["init", str(src), "--id", "arcface.infer", "--name", "特征示例", + "--algorithm", "arcface"]) + code_path = src / "alg_arcface_infer" / "__init__.py" + + monkeypatch.setitem(sys.modules, "sophon", types.ModuleType("sophon")) + monkeypatch.setitem(sys.modules, "sophon.sail", types.ModuleType("sophon.sail")) + spec = importlib.util.spec_from_file_location("alg_arcface_infer_under_test", code_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + class Frame: + def width(self) -> int: + return 112 + + def height(self) -> int: + return 112 + + class Tensor: + def asnumpy(self): + return [[3.0, 4.0] + [0.0] * 510] # 归一化后前两维 = 0.6, 0.8 + + class Net: + def process(self, *args, **kwargs) -> None: + return None + + class Bmcv: + def bm_image_to_tensor(self, *args, **kwargs) -> None: + return None + + algo = module.Algorithm(ctx=None) + algo._output_tensor = Tensor() + algo._input_tensor = object() + algo._input_name, algo._input_shapes = "input", {} + algo._output_name, algo._graph = "output", "graph" + algo._net, algo._bmcv = Net(), Bmcv() + algo._preprocess = lambda frame: object() + + result = algo.infer(Frame()) + assert result.boxes is None + assert result.counts == {"face": 1} and result.value == 1.0 + assert result.meta["dim"] == 512 + assert result.meta["embedding"][0] == pytest.approx(0.6) + assert result.meta["embedding"][1] == pytest.approx(0.8) + assert result.meta["norm"] == pytest.approx(1.0) + + +def test_cli_init_lprnet_template(tmp_path: Path) -> None: + """`--algorithm lprnet` 生成车牌识别包（text 输出 + 字符表）。""" + src = tmp_path / "lprnet.demo" + assert cli_main(["init", str(src), "--id", "lprnet.demo", "--name", "车牌示例", + "--algorithm", "lprnet"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.id == "lprnet.demo" and manifest.kind == "classifier" + assert manifest.default_model().name == "lprnet_fp32_1b" + assert manifest.default_model().file == "models/lprnet_fp32_1b.bmodel" + assert manifest.default_model().resolution == [94, 24] + assert manifest.output["type"] == "text" + assert manifest.output["chars_file"] == "assets/chars.txt" + assert manifest.schema_defaults()["interval"] == 0.5 + + code = (src / manifest.code_package / "__init__.py").read_text() + compile(code, "alg_lprnet_demo/__init__.py", "exec") + assert "_preprocess" in code and "_postprocess" in code and "chars.txt" in code + + chars = [line for line in (src / "assets" / "chars.txt").read_text().splitlines() + if line.strip() and not line.startswith("#")] + assert len(chars) == 68 and chars[0] == "京" and chars[-1] == "-" + + model = tmp_path / "fake.bmodel" + model.write_bytes(b"FAKE-BMODEL" * 512) + out = tmp_path / "lprnet.demo-1.0.0.mbank" + assert cli_main(["build", str(src), "--out", str(out), + "--model", f"lprnet_fp32_1b={model}", "--quiet"]) == 0 + assert cli_main(["verify", str(out), "--json"]) == 0 + + +def test_lprnet_template_infer_with_fake_sail(tmp_path: Path, monkeypatch) -> None: + """不依赖设备验证 CTC 贪心解码：去重、去 blank、空结果。""" + src = tmp_path / "lprnet.infer" + cli_main(["init", str(src), "--id", "lprnet.infer", "--name", "车牌推理", + "--algorithm", "lprnet"]) + code_path = src / "alg_lprnet_infer" / "__init__.py" + + monkeypatch.setitem(sys.modules, "sophon", types.ModuleType("sophon")) + monkeypatch.setitem(sys.modules, "sophon.sail", types.ModuleType("sophon.sail")) + spec = importlib.util.spec_from_file_location("alg_lprnet_infer_under_test", code_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + class Frame: + def width(self) -> int: + return 94 + + def height(self) -> int: + return 24 + + class Tensor: + def __init__(self) -> None: + self.logits = np.full((1, 4, 6), -10.0, dtype=np.float32) + + def asnumpy(self): + return self.logits + + class Net: + def process(self, *args, **kwargs) -> None: + return None + + class Bmcv: + def bm_image_to_tensor(self, *args, **kwargs) -> None: + return None + + tensor = Tensor() + algo = module.Algorithm(ctx=None) + algo._chars = ["京", "A", "1", "-"] + algo._blank = 3 + algo._output_tensor = tensor + algo._input_tensor = object() + algo._input_name, algo._input_shapes = "input", {} + algo._output_name, algo._graph = "output", "graph" + algo._net, algo._bmcv = Net(), Bmcv() + algo._preprocess = lambda frame: object() + + # 模型输出 [1, C=4, T=6]：0,0,3(blank),1,2,2 → CTC 去重后 "京A1" + for t, cls in enumerate([0, 0, 3, 1, 2, 2]): + tensor.logits[0, cls, t] = 10.0 + result = algo.infer(Frame()) + assert result.boxes is None + assert result.meta["text"] == "京A1" + assert result.meta["chars"] == [0, 1, 2] + assert result.counts == {"plate": 1} and result.value == 1.0 + + tensor.logits = np.full((1, 4, 6), -10.0, dtype=np.float32) # 全是 blank + tensor.logits[0, 3, :] = 10.0 + empty = algo.infer(Frame()) + assert empty.meta["text"] == "" and empty.meta["chars"] == [] + assert empty.counts == {} and empty.value == 0.0 + + +def test_cli_init_face_recognition_template(tmp_path: Path, monkeypatch) -> None: + """`--algorithm face-recognition` 生成双模型包（检测框 + 512 维 embedding）。""" + src = tmp_path / "face.demo" + assert cli_main(["init", str(src), "--id", "face.demo", "--name", "人脸示例", + "--algorithm", "face-recognition"]) == 0 + manifest = Manifest.parse(json.loads((src / "manifest.json").read_text())) + assert manifest.id == "face.demo" and manifest.kind == "detector" + assert [m.name for m in manifest.models] == ["face_detector", "arcface"] + assert manifest.default_model().name == "face_detector" + assert manifest.model("arcface").file == "models/arcface_resnet50_fp32_1b.bmodel" + assert manifest.output["embedding_size"] == 512 + assert manifest.schema_defaults()["nms"] == 0.3 + assert manifest.schema_defaults()["min_face"] == 16 + + code = (src / manifest.code_package / "__init__.py").read_text() + compile(code, "alg_face_demo/__init__.py", "exec") + assert "_priorbox" in code and "_detect" in code and "_embed" in code + assert 'model_bytes("arcface")' in code + + labels = [line for line in (src / "assets" / "labels.txt").read_text().splitlines() + if line.strip() and not line.startswith("#")] + assert labels == ["face"] + + det = tmp_path / "det.bmodel" + det.write_bytes(b"DET" * 2048) + arc = tmp_path / "arc.bmodel" + arc.write_bytes(b"ARC" * 2048) + out = tmp_path / "face.demo-1.0.0.mbank" + assert cli_main(["build", str(src), "--out", str(out), + "--model", f"face_detector={det}", + "--model", f"arcface={arc}", "--quiet"]) == 0 + assert cli_main(["verify", str(out), "--json"]) == 0 + + # 不依赖设备验证先验框/解码/NMS 的集成（官方 RetinaFace 口径） + monkeypatch.setitem(sys.modules, "sophon", types.ModuleType("sophon")) + monkeypatch.setitem(sys.modules, "sophon.sail", types.ModuleType("sophon.sail")) + spec = importlib.util.spec_from_file_location( + "alg_face_detect_under_test", src / "alg_face_demo" / "__init__.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + priors = module._priorbox(640, 640) + assert priors.shape == (16800, 4) + assert np.allclose(priors[0], [4 / 640, 4 / 640, 16 / 640, 16 / 640]) + + algo = module.Algorithm(ctx=None) + algo._conf, algo._nms, algo._max_faces = 0.5, 0.3, 20 + algo._priors = priors + algo._det_w = algo._det_h = 640 + loc = np.zeros((1, 16800, 4), dtype=np.float32) + conf = np.zeros((1, 16800, 2), dtype=np.float32) + conf[0, :, 1] = 0.01 + conf[0, 0, 1] = 0.9 + dets = algo._detect([loc, conf], 1.0, 640, 640) + assert dets.shape == (1, 6) + assert np.allclose(dets[0, :4], [0, 0, 12, 12]) + assert dets[0, 4] == pytest.approx(0.9) and dets[0, 5] == 0 + + +def test_model_download_specs_and_aliases() -> None: + from algpkg_sdk import models as model_store + + assert [m.key for m in model_store.resolve_models(["all"])] == list(model_store.MODEL_FILES) + assert model_store.resolve_models(["yolov26"])[0].key == "yolo26" + assert model_store.resolve_models(["lpr"])[0].key == "lprnet" + assert model_store.lfs_batch_url("http://host/repo.git").endswith( + "/repo.git/info/lfs/objects/batch") + with pytest.raises(ValueError): + model_store.resolve_models(["nope"]) + + +def test_cli_download_models_via_lfs(tmp_path: Path, monkeypatch, capsys) -> None: + """本地假 LFS 服务验证下载/校验/跳过（不依赖外网与 git-lfs）。""" + from algpkg_sdk import models as model_store + + content = b"FAKE-BMODEL-DOWNLOAD" * 4096 + oid = hashlib.sha256(content).hexdigest() + spec = model_store.ModelFile("demo", "demo_model.bmodel", oid, len(content), "测试模型") + monkeypatch.setattr(model_store, "MODEL_FILES", {"demo": spec}) + monkeypatch.setattr(model_store, "ALIASES", {"demo": "demo"}) + monkeypatch.setenv("NO_PROXY", "127.0.0.1,localhost") + + class Handler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: # LFS batch + length = int(self.headers.get("Content-Length") or 0) + payload = json.loads(self.rfile.read(length) or b"{}") + port = self.server.server_port + objects = [{ + "oid": item["oid"], "size": item["size"], + "actions": {"download": { + "href": f"http://127.0.0.1:{port}/objects/{item['oid']}"}}, + } for item in payload.get("objects", [])] + body = json.dumps({"objects": objects}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/vnd.git-lfs+json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self) -> None: # LFS object + if self.path.startswith("/objects/"): + self.send_response(200) + self.send_header("Content-Length", str(len(content))) + self.end_headers() + self.wfile.write(content) + else: + self.send_error(404) + + def log_message(self, *args) -> None: + return None + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + dest = tmp_path / "models" / "BM1688" + url = f"http://127.0.0.1:{server.server_port}/sophon-edge/algpkg-sdk.git" + assert cli_main(["download-models", "--dest", str(dest), + "--url", url, "--model", "demo"]) == 0 + assert (dest / "demo_model.bmodel").read_bytes() == content + assert "下载 1" in capsys.readouterr().out + + assert cli_main(["download-models", "--dest", str(dest), + "--url", url, "--model", "demo"]) == 0 + assert "跳过 1" in capsys.readouterr().out + assert (dest / "demo_model.bmodel").read_bytes() == content + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + def test_cli_check_reports_static_results(tmp_path: Path) -> None: """check 在无设备时也要给出"清单/结构 OK、试加载跳过"的诚实结论。""" src = tmp_path / "chk.demo"