199 lines
7.0 KiB
Python
199 lines
7.0 KiB
Python
# 作者：北小菜
|
||
"""ReID 特征提取引擎 — OSNet 系列 ONNX（仅 OnnxRuntime）
|
||
|
||
输入：人体 crop BGR → resize → ImageNet 归一化 → embedding 向量
|
||
输出：512 维 L2 归一化特征（默认 OSNet-AIN x1.0）
|
||
"""
|
||
import logging
|
||
import os
|
||
|
||
from app.analysis.engines.base import BaseEngine, EngineNotAvailableError
|
||
|
||
logger = logging.getLogger("analysis.engines.reid_onnx")
|
||
|
||
try:
|
||
import numpy as np
|
||
_NP = True
|
||
except Exception:
|
||
np = None
|
||
_NP = False
|
||
|
||
try:
|
||
import cv2
|
||
_CV2 = True
|
||
except Exception:
|
||
cv2 = None
|
||
_CV2 = False
|
||
|
||
try:
|
||
import onnxruntime as ort
|
||
_ORT = True
|
||
except Exception:
|
||
ort = None
|
||
_ORT = False
|
||
|
||
REID_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) if _NP else None
|
||
REID_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) if _NP else None
|
||
|
||
|
||
def _providers_for_device(device):
|
||
d = (device or "cpu").lower()
|
||
if d in ("cuda", "gpu", "0") and _ORT:
|
||
avail = ort.get_available_providers()
|
||
if "CUDAExecutionProvider" in avail:
|
||
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
return ["CPUExecutionProvider"]
|
||
|
||
|
||
class ReidOnnxEngine(BaseEngine):
|
||
"""ReID embedding 引擎（batch=1 静态 ONNX）。"""
|
||
|
||
ENGINE_NAME = "reid_onnx"
|
||
|
||
def __init__(self, **kwargs):
|
||
super(ReidOnnxEngine, self).__init__(**kwargs)
|
||
self.task_type = "reid"
|
||
self._session = None
|
||
self._input_name = None
|
||
self._output_name = None
|
||
self._embedding_dim = 512
|
||
if not self.providers:
|
||
self.providers = _providers_for_device(self.device)
|
||
|
||
@staticmethod
|
||
def is_available():
|
||
return _ORT and _CV2 and _NP
|
||
|
||
@staticmethod
|
||
def version():
|
||
if not _ORT:
|
||
return None
|
||
try:
|
||
return getattr(ort, "__version__", "unknown")
|
||
except Exception:
|
||
return "unknown"
|
||
|
||
def load(self):
|
||
if not self.is_available():
|
||
logger.warning("ReidOnnxEngine: 依赖未安装")
|
||
return False
|
||
if not self.model_file or not os.path.exists(self.model_file):
|
||
logger.warning("ReidOnnxEngine: 模型不存在 %s", self.model_file)
|
||
return False
|
||
try:
|
||
so = ort.SessionOptions()
|
||
so.log_severity_level = 3
|
||
self._session = ort.InferenceSession(
|
||
self.model_file, sess_options=so, providers=self.providers)
|
||
self._input_name = self._session.get_inputs()[0].name
|
||
self._output_name = self._session.get_outputs()[0].name
|
||
out_shape = self._session.get_outputs()[0].shape
|
||
if out_shape and len(out_shape) >= 2 and out_shape[-1]:
|
||
try:
|
||
self._embedding_dim = int(out_shape[-1])
|
||
except Exception:
|
||
pass
|
||
self._loaded = True
|
||
logger.info("ReidOnnxEngine: loaded %s providers=%s dim=%d",
|
||
self.model_file, self._session.get_providers(), self._embedding_dim)
|
||
return True
|
||
except Exception as e:
|
||
logger.error("ReidOnnxEngine load failed: %s", e)
|
||
self._session = None
|
||
self._loaded = False
|
||
return False
|
||
|
||
@property
|
||
def session(self):
|
||
return self._session
|
||
|
||
def _preprocess_crop(self, frame_bgr, box):
|
||
iw = int(self.input_size[0] or 128)
|
||
ih = int(self.input_size[1] or 256)
|
||
x1, y1, x2, y2 = [int(v) for v in box]
|
||
h, w = frame_bgr.shape[:2]
|
||
x1 = max(0, min(x1, w - 1))
|
||
x2 = max(0, min(x2, w))
|
||
y1 = max(0, min(y1, h - 1))
|
||
y2 = max(0, min(y2, h))
|
||
if x2 <= x1 or y2 <= y1:
|
||
return None
|
||
crop = frame_bgr[y1:y2, x1:x2]
|
||
if crop.size == 0:
|
||
return None
|
||
rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
|
||
resized = cv2.resize(rgb, (iw, ih), interpolation=cv2.INTER_LINEAR)
|
||
arr = resized.astype(np.float32) / 255.0
|
||
arr = (arr - REID_MEAN) / REID_STD
|
||
return np.transpose(arr, (2, 0, 1))[None, ...].astype(np.float32)
|
||
|
||
def extract_embeddings(self, frame_bgr, boxes):
|
||
"""对多个 bbox 提取 embedding，返回与 boxes 对齐的 (valid_idx, embeddings)。"""
|
||
if not self.ready() or frame_bgr is None:
|
||
return [], np.zeros((0, self._embedding_dim), dtype=np.float32)
|
||
valid_idx = []
|
||
rows = []
|
||
for i, box in enumerate(boxes or []):
|
||
blob = self._preprocess_crop(frame_bgr, box)
|
||
if blob is None:
|
||
continue
|
||
out = self._session.run([self._output_name], {self._input_name: blob})[0]
|
||
vec = np.asarray(out, dtype=np.float32).reshape(-1)
|
||
norm = np.linalg.norm(vec)
|
||
if norm > 1e-12:
|
||
vec = vec / norm
|
||
valid_idx.append(i)
|
||
rows.append(vec)
|
||
if not rows:
|
||
return [], np.zeros((0, self._embedding_dim), dtype=np.float32)
|
||
return valid_idx, np.stack(rows, axis=0)
|
||
|
||
def detect(self, frame_bgr):
|
||
"""兼容 BaseEngine 接口；ReID 单模型无法对全图直接检测，返回空列表。"""
|
||
return []
|
||
|
||
def probe(self):
|
||
info = {
|
||
"engine": self.ENGINE_NAME,
|
||
"available": self.is_available(),
|
||
"version": self.version(),
|
||
"input_shape": None,
|
||
"output_shape": None,
|
||
"labels": [],
|
||
"model_file": self.model_file,
|
||
"task_type": "reid",
|
||
"algorithm_type": self.algorithm_type,
|
||
"device": self.device,
|
||
"embedding_dim": self._embedding_dim,
|
||
}
|
||
if not self.is_available() or not self.model_file or not os.path.exists(self.model_file):
|
||
return info
|
||
try:
|
||
so = ort.SessionOptions()
|
||
so.log_severity_level = 3
|
||
sess = ort.InferenceSession(self.model_file, sess_options=so, providers=["CPUExecutionProvider"])
|
||
inputs = sess.get_inputs()
|
||
outputs = sess.get_outputs()
|
||
info["input_shape"] = list(inputs[0].shape) if inputs else None
|
||
info["output_shape"] = [list(o.shape) for o in outputs] if outputs else None
|
||
if inputs and len(inputs[0].shape) >= 4:
|
||
# ONNX NCHW: [N,C,H,W] → width=shape[-1], height=shape[-2]
|
||
info["input_size_inferred"] = (int(inputs[0].shape[-1]), int(inputs[0].shape[-2]))
|
||
if outputs and outputs[0].shape:
|
||
sh = outputs[0].shape
|
||
if len(sh) >= 2 and sh[-1]:
|
||
info["embedding_dim"] = int(sh[-1])
|
||
except Exception as e:
|
||
info["error"] = str(e)
|
||
if self.model_file and os.path.isfile(self.model_file):
|
||
info["model_file_size"] = os.path.getsize(self.model_file)
|
||
return info
|
||
|
||
def info(self):
|
||
d = super(ReidOnnxEngine, self).info()
|
||
d["version"] = self.version()
|
||
d["task_type"] = "reid"
|
||
d["embedding_dim"] = self._embedding_dim
|
||
d["providers"] = self._session.get_providers() if self._session else []
|
||
return d
|