"""人员计数示例算法包（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