111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
|
"""引擎工厂：按 inference_engine 名分发实例
|
|||
|
|
|
|||
|
|
注册引擎：
|
|||
|
|
- yolo_pytorch : Yolo-PyTorch（ultralytics 原生，全版本全任务，主引擎）
|
|||
|
|
- onnxruntime : OnnxRuntime
|
|||
|
|
- openvino : OpenVINO
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
from app.analysis.engines.base import BaseEngine, EngineNotAvailableError
|
|||
|
|
from app.analysis.engines.yolo_pytorch_engine import YoloPytorchEngine
|
|||
|
|
from app.analysis.engines.onnx_engine import OnnxEngine
|
|||
|
|
from app.analysis.engines.openvino_engine import OpenVinoEngine
|
|||
|
|
from app.analysis.engines.reid_onnx_engine import ReidOnnxEngine
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("analysis.engines.factory")
|
|||
|
|
|
|||
|
|
_ENGINE_REGISTRY = {
|
|||
|
|
"yolo_pytorch": YoloPytorchEngine,
|
|||
|
|
"yolopytorch": YoloPytorchEngine, # 别名
|
|||
|
|
"pytorch": YoloPytorchEngine, # 兼容旧名
|
|||
|
|
"onnxruntime": OnnxEngine,
|
|||
|
|
"onnx": OnnxEngine,
|
|||
|
|
"openvino": OpenVinoEngine,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 设备选项按引擎分组（供前端动态下拉使用）
|
|||
|
|
_DEVICE_OPTIONS = {
|
|||
|
|
"yolo_pytorch": [
|
|||
|
|
{"value": "cpu", "label": "CPU"},
|
|||
|
|
{"value": "cuda", "label": "CUDA (GPU)"},
|
|||
|
|
],
|
|||
|
|
"onnxruntime": [
|
|||
|
|
{"value": "cpu", "label": "CPU"},
|
|||
|
|
{"value": "cuda", "label": "CUDA (GPU)"},
|
|||
|
|
],
|
|||
|
|
"openvino": [
|
|||
|
|
{"value": "cpu", "label": "CPU"},
|
|||
|
|
{"value": "gpu", "label": "GPU (Intel iGPU/dGPU)"},
|
|||
|
|
],
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def list_engines():
|
|||
|
|
"""返回所有注册引擎的可用性信息"""
|
|||
|
|
out = []
|
|||
|
|
seen = set()
|
|||
|
|
for name, cls in _ENGINE_REGISTRY.items():
|
|||
|
|
if name in ("yolopytorch", "pytorch", "onnx"):
|
|||
|
|
continue
|
|||
|
|
if name in seen:
|
|||
|
|
continue
|
|||
|
|
seen.add(name)
|
|||
|
|
try:
|
|||
|
|
available = cls.is_available()
|
|||
|
|
version = cls.version() if available else None
|
|||
|
|
except Exception as e:
|
|||
|
|
available = False
|
|||
|
|
version = None
|
|||
|
|
logger.warning("list_engines %s err: %s", name, e)
|
|||
|
|
# 主引擎额外提供 ultralytics 版本与 CUDA 可用性
|
|||
|
|
item = {"name": name, "available": available, "version": version,
|
|||
|
|
"devices": _DEVICE_OPTIONS.get(name, _DEVICE_OPTIONS["yolo_pytorch"])}
|
|||
|
|
if name == "yolo_pytorch" and available:
|
|||
|
|
try:
|
|||
|
|
item["ultralytics_version"] = YoloPytorchEngine.ultralytics_version()
|
|||
|
|
import torch
|
|||
|
|
item["cuda_available"] = bool(torch.cuda.is_available())
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
out.append(item)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
def device_options(engine_name):
|
|||
|
|
return _DEVICE_OPTIONS.get((engine_name or "").lower(), _DEVICE_OPTIONS["yolo_pytorch"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
class EngineFactory(object):
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def create(engine_name, **kwargs):
|
|||
|
|
task_type = (kwargs.get("task_type") or "detect").lower()
|
|||
|
|
if task_type == "reid":
|
|||
|
|
eng = (engine_name or "").lower()
|
|||
|
|
if eng not in ("onnxruntime", "onnx"):
|
|||
|
|
raise EngineNotAvailableError("ReID models only support onnxruntime")
|
|||
|
|
if not ReidOnnxEngine.is_available():
|
|||
|
|
raise EngineNotAvailableError("reid onnxruntime not installed")
|
|||
|
|
return ReidOnnxEngine(**kwargs)
|
|||
|
|
cls = _ENGINE_REGISTRY.get((engine_name or "").lower())
|
|||
|
|
if cls is None:
|
|||
|
|
raise EngineNotAvailableError("unknown engine: %s" % engine_name)
|
|||
|
|
if not cls.is_available():
|
|||
|
|
raise EngineNotAvailableError("engine %s not installed" % engine_name)
|
|||
|
|
return cls(**kwargs)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def is_available(engine_name):
|
|||
|
|
cls = _ENGINE_REGISTRY.get((engine_name or "").lower())
|
|||
|
|
return bool(cls and cls.is_available())
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def list_engines():
|
|||
|
|
return list_engines()
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def device_options(engine_name):
|
|||
|
|
return device_options(engine_name)
|