67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
# coding=utf-8
|
||
"""推理加速：自动选择 CUDA / OpenVINO / CPU，并解析已导出的 OpenVINO 模型目录。"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from pathlib import Path
|
||
|
||
_ROOT = Path(__file__).resolve().parents[1]
|
||
|
||
|
||
def project_root() -> Path:
|
||
return _ROOT
|
||
|
||
|
||
def resolve_yolo_weights(pt_path: str | Path | None = None) -> str:
|
||
"""优先使用已导出的 OpenVINO 模型，否则回退到 .pt。
|
||
|
||
查找顺序:
|
||
1) 环境变量 DISH_YOLO_MODEL
|
||
2) best_openvino_model/（ultralytics export 默认目录）
|
||
3) 传入的 pt_path / best.pt
|
||
"""
|
||
env = (os.environ.get("DISH_YOLO_MODEL") or "").strip()
|
||
if env and Path(env).exists():
|
||
return str(Path(env).resolve())
|
||
|
||
ov_dir = _ROOT / "best_openvino_model"
|
||
if ov_dir.is_dir() and any(ov_dir.glob("*.xml")):
|
||
return str(ov_dir.resolve())
|
||
|
||
if pt_path:
|
||
p = Path(pt_path)
|
||
if p.exists():
|
||
return str(p.resolve())
|
||
|
||
fallback = _ROOT / "best.pt"
|
||
return str(fallback.resolve())
|
||
|
||
|
||
def torch_device():
|
||
import torch
|
||
|
||
if torch.cuda.is_available():
|
||
return torch.device("cuda")
|
||
return torch.device("cpu")
|
||
|
||
|
||
def openvino_available() -> bool:
|
||
try:
|
||
import openvino # noqa: F401
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def resolve_feature_openvino() -> Path | None:
|
||
"""ConvNeXt 特征提取 OpenVINO 模型目录（可选）。"""
|
||
env = (os.environ.get("DISH_FEATURE_OPENVINO") or "").strip()
|
||
if env:
|
||
p = Path(env)
|
||
if p.is_dir() and any(p.glob("*.xml")):
|
||
return p.resolve()
|
||
default = _ROOT / "feature_openvino_model"
|
||
if default.is_dir() and any(default.glob("*.xml")):
|
||
return default.resolve()
|
||
return None
|