203 lines
6.8 KiB
Python
203 lines
6.8 KiB
Python
|
"""OpenVINO 引擎实现 —— YOLO 5/8/11/26 + 全任务 + 设备支持
|
|||
|
|
|
|||
|
|
依赖：openvino, opencv-python, numpy
|
|||
|
|
推理设备：
|
|||
|
|
cpu -> CPU
|
|||
|
|
gpu -> GPU
|
|||
|
|
cuda -> 不适用，回退 CPU（OpenVINO 不走 CUDA）
|
|||
|
|
"""
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from app.analysis.engines.base import BaseEngine, DetectionResult
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("analysis.engines.openvino")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import numpy as np
|
|||
|
|
_NP_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
np = None
|
|||
|
|
_NP_AVAILABLE = False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import cv2
|
|||
|
|
_CV2_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
cv2 = None
|
|||
|
|
_CV2_AVAILABLE = False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
from openvino import Core
|
|||
|
|
_OV_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
try:
|
|||
|
|
from openvino.runtime import Core
|
|||
|
|
_OV_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
Core = None
|
|||
|
|
_OV_AVAILABLE = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _device_for(device):
|
|||
|
|
d = (device or "cpu").lower()
|
|||
|
|
if d in ("gpu", "cuda", "0"):
|
|||
|
|
return "GPU"
|
|||
|
|
return "CPU"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OpenVinoEngine(BaseEngine):
|
|||
|
|
ENGINE_NAME = "openvino"
|
|||
|
|
|
|||
|
|
def __init__(self, **kwargs):
|
|||
|
|
super(OpenVinoEngine, self).__init__(**kwargs)
|
|||
|
|
self._compiled = None
|
|||
|
|
self._infer_req = None
|
|||
|
|
self._input_key = None
|
|||
|
|
self._output_keys = None
|
|||
|
|
self.task_type = (kwargs.get("task_type") or "detect").lower()
|
|||
|
|
self.device = kwargs.get("device") or "cpu"
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def is_available():
|
|||
|
|
return _OV_AVAILABLE and _CV2_AVAILABLE and _NP_AVAILABLE
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def version():
|
|||
|
|
if not _OV_AVAILABLE:
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
from openvino.runtime import get_version
|
|||
|
|
return get_version()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
import openvino
|
|||
|
|
return getattr(openvino, "__version__", "unknown")
|
|||
|
|
except Exception:
|
|||
|
|
return "unknown"
|
|||
|
|
|
|||
|
|
def _resolve_model_path(self):
|
|||
|
|
if not self.model_file:
|
|||
|
|
return None
|
|||
|
|
if self.model_file.lower().endswith(".onnx"):
|
|||
|
|
return self.model_file
|
|||
|
|
base, ext = os.path.splitext(self.model_file)
|
|||
|
|
if ext.lower() == ".xml":
|
|||
|
|
return self.model_file
|
|||
|
|
xml_p = base + ".xml"
|
|||
|
|
if os.path.exists(xml_p):
|
|||
|
|
return xml_p
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def load(self):
|
|||
|
|
if not self.is_available():
|
|||
|
|
logger.warning("OpenVinoEngine: 依赖未安装")
|
|||
|
|
return False
|
|||
|
|
path = self._resolve_model_path()
|
|||
|
|
if not path or not os.path.exists(path):
|
|||
|
|
logger.warning("OpenVinoEngine: 模型文件不可用: %s", self.model_file)
|
|||
|
|
return False
|
|||
|
|
if not self.labels:
|
|||
|
|
self.labels = self._resolve_labels(self.model_file)
|
|||
|
|
try:
|
|||
|
|
core = Core()
|
|||
|
|
model = core.read_model(path)
|
|||
|
|
dev = _device_for(self.device)
|
|||
|
|
# GPU 不可用时回退 CPU
|
|||
|
|
try:
|
|||
|
|
self._compiled = core.compile_model(model, dev)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("OpenVinoEngine: 设备 %s 编译失败，回退 CPU: %s", dev, e)
|
|||
|
|
self._compiled = core.compile_model(model, "CPU")
|
|||
|
|
dev = "CPU"
|
|||
|
|
self._infer_req = self._compiled.create_infer_request()
|
|||
|
|
self._input_key = list(self._compiled.inputs)[0]
|
|||
|
|
self._output_keys = list(self._compiled.outputs)
|
|||
|
|
self._loaded = True
|
|||
|
|
logger.info("OpenVinoEngine: 已加载 %s, task=%s, device=%s, labels=%d",
|
|||
|
|
path, self.task_type, dev, len(self.labels))
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("OpenVinoEngine: 加载失败: %s", e)
|
|||
|
|
self._loaded = False
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def _preprocess(self, frame_bgr):
|
|||
|
|
iw, ih = self.input_size
|
|||
|
|
resized = cv2.resize(frame_bgr, (iw, ih))
|
|||
|
|
rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
|
|||
|
|
blob = rgb.astype(np.float32) / 255.0
|
|||
|
|
blob = np.transpose(blob, (2, 0, 1))[None, ...]
|
|||
|
|
return blob
|
|||
|
|
|
|||
|
|
def detect(self, frame_bgr):
|
|||
|
|
if not self.ready() or frame_bgr is None:
|
|||
|
|
return []
|
|||
|
|
try:
|
|||
|
|
h, w = frame_bgr.shape[:2]
|
|||
|
|
blob = self._preprocess(frame_bgr)
|
|||
|
|
self._infer_req.infer({self._input_key: blob})
|
|||
|
|
outputs = []
|
|||
|
|
out_count = len(self._output_keys)
|
|||
|
|
for i in range(out_count):
|
|||
|
|
try:
|
|||
|
|
outputs.append(self._infer_req.get_output_tensor(i).data)
|
|||
|
|
except Exception:
|
|||
|
|
try:
|
|||
|
|
outputs.append(self._infer_req.get_output_tensor().data)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("OpenVinoEngine: 读取输出 tensor[%d] 失败: %s", i, e)
|
|||
|
|
if not outputs:
|
|||
|
|
return []
|
|||
|
|
from app.analysis.engines.yolo_postprocess import decode_outputs
|
|||
|
|
results = decode_outputs(
|
|||
|
|
outputs=outputs,
|
|||
|
|
algorithm_type=self.algorithm_type,
|
|||
|
|
task_type=self.task_type,
|
|||
|
|
labels=self.labels,
|
|||
|
|
input_size=self.input_size,
|
|||
|
|
conf_threshold=self.conf_threshold,
|
|||
|
|
iou_threshold=self.iou_threshold,
|
|||
|
|
orig_size=(w, h),
|
|||
|
|
)
|
|||
|
|
return [DetectionResult(**r) for r in results]
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("OpenVinoEngine.detect() err: %s", e)
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def info(self):
|
|||
|
|
d = super(OpenVinoEngine, self).info()
|
|||
|
|
d["version"] = self.version()
|
|||
|
|
d["task_type"] = self.task_type
|
|||
|
|
d["device"] = _device_for(self.device)
|
|||
|
|
return d
|
|||
|
|
|
|||
|
|
def probe(self):
|
|||
|
|
info = {"engine": self.ENGINE_NAME, "available": self.is_available(),
|
|||
|
|
"version": self.version(), "input_shape": None, "output_shape": None,
|
|||
|
|
"labels": self.labels, "model_file": self.model_file,
|
|||
|
|
"task_type": self.task_type, "device": _device_for(self.device)}
|
|||
|
|
if not self.is_available() or not self.model_file:
|
|||
|
|
return info
|
|||
|
|
path = self._resolve_model_path()
|
|||
|
|
if not path or not os.path.exists(path):
|
|||
|
|
info["error"] = "model file not resolvable"
|
|||
|
|
return info
|
|||
|
|
try:
|
|||
|
|
core = Core()
|
|||
|
|
model = core.read_model(path)
|
|||
|
|
if model.inputs:
|
|||
|
|
shape = list(model.inputs[0].shape)
|
|||
|
|
info["input_shape"] = shape
|
|||
|
|
if len(shape) >= 4:
|
|||
|
|
info["input_size_inferred"] = (int(shape[-1]), int(shape[-2]))
|
|||
|
|
if model.outputs:
|
|||
|
|
info["output_shape"] = [list(o.shape) for o in model.outputs]
|
|||
|
|
if not self.labels:
|
|||
|
|
self.labels = self._resolve_labels(self.model_file)
|
|||
|
|
info["labels"] = self.labels
|
|||
|
|
except Exception as e:
|
|||
|
|
info["error"] = str(e)
|
|||
|
|
return info
|