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
|