linaro
170e4a10a4
本仓库原先内嵌在 sophon-edge 的 `algpkg-sdk/`（该仓库历史上的三次相关提交： e8101ba「算法开发能力独立成项目」、090c406「契约名落地为真实模块 + 多模型 ctx」、 b666161「sail 存根 + 编写风格修正」）。现拆出独立建库，便于算法作者单独 clone/迭代、 也便于单独发版（打成 wheel 分发给第三方或进 worker 交付包的 vendor/）。 - 内容：`src/algpkg_sdk/`（api/manifest/loader/builder/testkit/cli）+ 契约名模块 `src/video_algpkg_api/`、`examples/person-count`（完整示例）、`tests/`（pytest，无设备自动跳过）、 `README.md`（已改成独立仓库口径：命令在项目根执行，配套 worker 用路径依赖/`ALGPKG_SDK_DIR`）、 `pyproject.toml`（console script `algpkg`）、`uv.lock`； - 不含 `.venv/`（`.gitignore` 已覆盖，clone 后 `uv sync` 重建）。 配套 worker 仓库（sophon-edge）： `workers/video/pyproject.toml` 以路径依赖指向本仓库（`path = "../../../algpkg-sdk"`，可用 `ALGPKG_SDK_DIR` 覆盖）；交付打包 `deploy/pack_worker.py` 会把本仓库打成 wheel 放进 worker 包 `vendor/`。
147 lines
7.1 KiB
Python
147 lines
7.1 KiB
Python
"""人员计数示例算法包（M8.2/M8.4）——**自带推理实现**，不依赖 worker 内部模块。
|
||
|
||
这是"第三方算法包"的样板：只依赖
|
||
- `video_algpkg_api`（worker 提供的稳定接口：Algorithm / AlgoContext / AlgoResult）；
|
||
- 设备框架 `sophon.sail` 与 `numpy`；
|
||
- 包内模型（`ctx.model_bytes()` 从 `.mbank` 读字节，**不落盘**）。
|
||
|
||
算法语义：YOLOv26s 检测 → 只保留 person（COCO 类 0）→ 置信度/最小人高过滤 →
|
||
人数作为 `AlgoResult(counts={"person": n}, value=n)` 返回，并把人员框一并给出（供绘制/告警）。
|
||
|
||
与内置包（`builtin.yolov26-detect`）的区别：kind=counter、默认抽帧更慢（0.5s）、
|
||
只报人员、带最小人高参数 —— 用来验证一个任务同时挂"检测 + 计数"两个算法。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from typing import List
|
||
|
||
import numpy as np
|
||
import sophon.sail as sail
|
||
|
||
from video_algpkg_api import Algorithm as AlgorithmBase, AlgoResult
|
||
|
||
PERSON_CLASS = 0 # COCO 类别索引（person）
|
||
_BATCH_SUPPORTED = (1, 2, 3, 4, 8, 16, 32, 64, 128, 256)
|
||
|
||
|
||
class Algorithm(AlgorithmBase):
|
||
"""画面人员计数（单 bmodel，SAIL 引擎 + bmcv 前处理）。"""
|
||
|
||
kind = "counter"
|
||
|
||
#: 引擎与输入张量：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.3))
|
||
self._min_height = int(self.ctx.param("min_height", 0))
|
||
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] not in _BATCH_SUPPORTED:
|
||
raise RuntimeError(f"batch={self._input_shape[0]} 不受支持")
|
||
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, conf=%.2f, min_height=%d）",
|
||
self.ctx.algo_id, len(data) / 1048576.0, self._conf, self._min_height)
|
||
|
||
# ------------------------------------------------------------------ 推理
|
||
|
||
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()[:1]
|
||
persons = self._postprocess(preds[0], frame.width(), frame.height(), ratio, tx, ty)
|
||
boxes = np.zeros((len(persons), 6), dtype=np.float32) if not persons else \
|
||
np.stack(persons).astype(np.float32)
|
||
count = int(boxes.shape[0])
|
||
return AlgoResult(boxes=boxes if count else None,
|
||
counts={"person": count} if count else {},
|
||
value=float(count),
|
||
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) -> List[list]:
|
||
"""(N,6) 端到端输出 → 人员框（原图坐标；conf / 最小人高过滤）。"""
|
||
rows = batch[batch[:, 4] > self._conf] if batch.size else batch
|
||
out: List[list] = []
|
||
for row in rows:
|
||
if int(row[5]) != PERSON_CLASS:
|
||
continue
|
||
x1 = float((row[0] - tx) / ratio)
|
||
y1 = float((row[1] - ty) / ratio)
|
||
x2 = float((row[2] - tx) / ratio)
|
||
y2 = float((row[3] - ty) / ratio)
|
||
x1, x2 = max(0.0, min(x1, org_w - 1)), max(0.0, min(x2, org_w - 1))
|
||
y1, y2 = max(0.0, min(y1, org_h - 1)), max(0.0, min(y2, org_h - 1))
|
||
if self._min_height and (y2 - y1) < self._min_height:
|
||
continue
|
||
out.append([x1, y1, x2, y2, float(row[4]), PERSON_CLASS])
|
||
return out
|