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
|