277 lines
12 KiB
Python
277 lines
12 KiB
Python
|
"""Yolo-PyTorch 引擎（主引擎，最高准确率）
|
|||
|
|
|
|||
|
|
基于 ultralytics 包，原生支持 YOLOv5/v8/v11/YOLO26 全部任务：
|
|||
|
|
detect / segment / classify / pose / obb
|
|||
|
|
|
|||
|
|
依赖：torch, ultralytics, opencv-python, numpy
|
|||
|
|
推理设备：cpu / cuda / cuda:0 / 0 / gpu
|
|||
|
|
"""
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from app.analysis.engines.base import BaseEngine, DetectionResult
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("analysis.engines.yolo_pytorch")
|
|||
|
|
|
|||
|
|
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:
|
|||
|
|
import torch
|
|||
|
|
_TORCH_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
torch = None
|
|||
|
|
_TORCH_AVAILABLE = False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
from ultralytics import YOLO as _UltralyticsYOLO
|
|||
|
|
_ULTRALYTICS_AVAILABLE = True
|
|||
|
|
except Exception:
|
|||
|
|
_UltralyticsYOLO = None
|
|||
|
|
_ULTRALYTICS_AVAILABLE = False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _normalize_device(device, algorithm_type):
|
|||
|
|
"""把模型字段 device 归一化为 ultralytics 可接受的 device 字符串。
|
|||
|
|
|
|||
|
|
cpu -> 'cpu'
|
|||
|
|
cuda -> 'cuda'（如可用）否则回退 'cpu'
|
|||
|
|
gpu -> 'cuda'（OpenVINO 风格命名兼容）
|
|||
|
|
"""
|
|||
|
|
if not device:
|
|||
|
|
return "cpu"
|
|||
|
|
d = str(device).lower().strip()
|
|||
|
|
if d in ("cpu", ""):
|
|||
|
|
return "cpu"
|
|||
|
|
if d in ("cuda", "gpu", "cuda:0", "0"):
|
|||
|
|
if _TORCH_AVAILABLE and torch.cuda.is_available():
|
|||
|
|
return "cuda:0" if d in ("gpu", "0") else "cuda"
|
|||
|
|
logger.warning("YoloPytorchEngine: CUDA 不可用，回退 CPU")
|
|||
|
|
return "cpu"
|
|||
|
|
return d
|
|||
|
|
|
|||
|
|
|
|||
|
|
class YoloPytorchEngine(BaseEngine):
|
|||
|
|
ENGINE_NAME = "yolo_pytorch"
|
|||
|
|
|
|||
|
|
def __init__(self, **kwargs):
|
|||
|
|
super(YoloPytorchEngine, self).__init__(**kwargs)
|
|||
|
|
self._model = None
|
|||
|
|
self.task_type = (kwargs.get("task_type") or "detect").lower()
|
|||
|
|
self.device = kwargs.get("device") or "cpu"
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def is_available():
|
|||
|
|
return _TORCH_AVAILABLE and _ULTRALYTICS_AVAILABLE and _CV2_AVAILABLE and _NP_AVAILABLE
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def version():
|
|||
|
|
if not _TORCH_AVAILABLE:
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
return getattr(torch, "__version__", "unknown")
|
|||
|
|
except Exception:
|
|||
|
|
return "unknown"
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def ultralytics_version():
|
|||
|
|
if not _ULTRALYTICS_AVAILABLE:
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
import ultralytics
|
|||
|
|
return getattr(ultralytics, "__version__", "unknown")
|
|||
|
|
except Exception:
|
|||
|
|
return "unknown"
|
|||
|
|
|
|||
|
|
def load(self):
|
|||
|
|
if not self.is_available():
|
|||
|
|
logger.warning("YoloPytorchEngine: 依赖未安装 (torch=%s ultralytics=%s cv2=%s np=%s)",
|
|||
|
|
_TORCH_AVAILABLE, _ULTRALYTICS_AVAILABLE, _CV2_AVAILABLE, _NP_AVAILABLE)
|
|||
|
|
return False
|
|||
|
|
if not self.model_file or not os.path.exists(self.model_file):
|
|||
|
|
logger.warning("YoloPytorchEngine: 模型文件不存在: %s", self.model_file)
|
|||
|
|
return False
|
|||
|
|
from app.utils.ModelTrust import require_trusted_model
|
|||
|
|
require_trusted_model(self.model_file)
|
|||
|
|
if not self.labels:
|
|||
|
|
self.labels = self._resolve_labels(self.model_file)
|
|||
|
|
try:
|
|||
|
|
self._model = _UltralyticsYOLO(self.model_file, task=self.task_type)
|
|||
|
|
# 推断 input_size
|
|||
|
|
try:
|
|||
|
|
cfg = getattr(self._model, "overrides", {}) or {}
|
|||
|
|
imgsz = cfg.get("imgsz", None)
|
|||
|
|
if isinstance(imgsz, int) and imgsz > 0:
|
|||
|
|
self.input_size = (int(imgsz), int(imgsz))
|
|||
|
|
elif isinstance(imgsz, (list, tuple)) and len(imgsz) >= 1:
|
|||
|
|
s = int(imgsz[0])
|
|||
|
|
self.input_size = (s, s)
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
# 预热（小尺寸 dummy），让模型迁移到目标设备
|
|||
|
|
try:
|
|||
|
|
dev = _normalize_device(self.device, self.algorithm_type)
|
|||
|
|
dummy = np.zeros((self.input_size[1], self.input_size[0], 3), dtype=np.uint8)
|
|||
|
|
self._model.predict(dummy, imgsz=max(self.input_size), device=dev,
|
|||
|
|
conf=self.conf_threshold, iou=self.iou_threshold,
|
|||
|
|
verbose=False, save=False)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine 预热失败（忽略）: %s", e)
|
|||
|
|
self._loaded = True
|
|||
|
|
logger.info("YoloPytorchEngine: 已加载 %s task=%s device=%s labels=%d",
|
|||
|
|
self.model_file, self.task_type, _normalize_device(self.device, self.algorithm_type),
|
|||
|
|
len(self.labels))
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("YoloPytorchEngine: 加载失败: %s", e)
|
|||
|
|
self._loaded = False
|
|||
|
|
self._model = None
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def detect(self, frame_bgr):
|
|||
|
|
if not self.ready() or frame_bgr is None or self._model is None:
|
|||
|
|
return []
|
|||
|
|
try:
|
|||
|
|
dev = _normalize_device(self.device, self.algorithm_type)
|
|||
|
|
iw, ih = self.input_size
|
|||
|
|
results = self._model.predict(frame_bgr, imgsz=max(iw, ih), device=dev,
|
|||
|
|
conf=self.conf_threshold, iou=self.iou_threshold,
|
|||
|
|
verbose=False, save=False)
|
|||
|
|
return self._parse_results(results)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine.detect() err: %s", e)
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def _parse_results(self, results):
|
|||
|
|
out = []
|
|||
|
|
try:
|
|||
|
|
r = results[0]
|
|||
|
|
task = self.task_type
|
|||
|
|
if task == "classify":
|
|||
|
|
probs = getattr(r, "probs", None)
|
|||
|
|
if probs is not None:
|
|||
|
|
cid = int(probs.top1)
|
|||
|
|
score = float(probs.top1conf)
|
|||
|
|
names = getattr(r, "names", {}) or {}
|
|||
|
|
label = names.get(cid, str(cid)) if isinstance(names, dict) else (
|
|||
|
|
self.labels[cid] if 0 <= cid < len(self.labels) else str(cid))
|
|||
|
|
out.append(DetectionResult(box=[0, 0, 0, 0], label=str(label), score=score, task="classify"))
|
|||
|
|
return out
|
|||
|
|
# 其余任务都有 boxes
|
|||
|
|
boxes_obj = getattr(r, "boxes", None)
|
|||
|
|
names = getattr(r, "names", {}) or {}
|
|||
|
|
if boxes_obj is not None:
|
|||
|
|
try:
|
|||
|
|
xyxy = boxes_obj.xyxy.cpu().numpy()
|
|||
|
|
confs = boxes_obj.conf.cpu().numpy()
|
|||
|
|
cls_ids = boxes_obj.cls.cpu().numpy().astype(int)
|
|||
|
|
for (b, s, cid) in zip(xyxy, confs, cls_ids):
|
|||
|
|
if s < self.conf_threshold:
|
|||
|
|
continue
|
|||
|
|
label = names.get(int(cid), str(int(cid))) if isinstance(names, dict) else (
|
|||
|
|
self.labels[int(cid)] if 0 <= int(cid) < len(self.labels) else str(int(cid)))
|
|||
|
|
item = DetectionResult(box=[int(b[0]), int(b[1]), int(b[2]), int(b[3])],
|
|||
|
|
label=str(label), score=float(s), task=task)
|
|||
|
|
out.append(item)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine parse boxes err: %s", e)
|
|||
|
|
|
|||
|
|
# segment: masks
|
|||
|
|
if task == "segment":
|
|||
|
|
masks_obj = getattr(r, "masks", None)
|
|||
|
|
if masks_obj is not None and out:
|
|||
|
|
try:
|
|||
|
|
# masks_obj.xy 是 list of (N,2) 多边形点（原图坐标）
|
|||
|
|
polys = masks_obj.xy
|
|||
|
|
for i, p in enumerate(polys):
|
|||
|
|
if i < len(out):
|
|||
|
|
out[i]["mask_polygon"] = [[float(x), float(y)] for x, y in p]
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine parse masks err: %s", e)
|
|||
|
|
|
|||
|
|
# pose: keypoints
|
|||
|
|
if task == "pose":
|
|||
|
|
kp_obj = getattr(r, "keypoints", None)
|
|||
|
|
if kp_obj is not None and out:
|
|||
|
|
try:
|
|||
|
|
kpts = kp_obj.xy.cpu().numpy() # [N, nk, 2]
|
|||
|
|
confs = kp_obj.conf.cpu().numpy() # [N, nk]
|
|||
|
|
for i in range(min(len(kpts), len(out))):
|
|||
|
|
kp = kpts[i]
|
|||
|
|
kc = confs[i]
|
|||
|
|
out[i]["keypoints"] = [
|
|||
|
|
[float(kp[j, 0]), float(kp[j, 1]), float(kc[j])]
|
|||
|
|
for j in range(len(kp))
|
|||
|
|
]
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine parse keypoints err: %s", e)
|
|||
|
|
|
|||
|
|
# obb: rotated boxes + angle
|
|||
|
|
if task == "obb":
|
|||
|
|
obb_obj = getattr(r, "obb", None)
|
|||
|
|
if obb_obj is not None and out:
|
|||
|
|
try:
|
|||
|
|
# obb.theta: [N] radians
|
|||
|
|
thetas = obb_obj.theta.cpu().numpy()
|
|||
|
|
for i in range(min(len(thetas), len(out))):
|
|||
|
|
out[i]["angle"] = float(thetas[i])
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine parse obb err: %s", e)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("YoloPytorchEngine._parse_results err: %s", e)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
def info(self):
|
|||
|
|
d = super(YoloPytorchEngine, self).info()
|
|||
|
|
d["version"] = self.version()
|
|||
|
|
d["ultralytics_version"] = self.ultralytics_version()
|
|||
|
|
d["task_type"] = self.task_type
|
|||
|
|
d["device"] = _normalize_device(self.device, self.algorithm_type)
|
|||
|
|
d["cuda_available"] = bool(_TORCH_AVAILABLE and torch.cuda.is_available())
|
|||
|
|
return d
|
|||
|
|
|
|||
|
|
def probe(self):
|
|||
|
|
info = {"engine": self.ENGINE_NAME, "available": self.is_available(),
|
|||
|
|
"version": self.version(), "ultralytics_version": self.ultralytics_version(),
|
|||
|
|
"cuda_available": bool(_TORCH_AVAILABLE and torch.cuda.is_available()),
|
|||
|
|
"input_shape": None, "output_shape": None,
|
|||
|
|
"labels": self.labels, "model_file": self.model_file,
|
|||
|
|
"task_type": self.task_type, "device": self.device}
|
|||
|
|
if not self.is_available() or not self.model_file or not os.path.exists(self.model_file):
|
|||
|
|
return info
|
|||
|
|
try:
|
|||
|
|
info["model_file_size"] = os.path.getsize(self.model_file)
|
|||
|
|
if not self.labels:
|
|||
|
|
self.labels = self._resolve_labels(self.model_file)
|
|||
|
|
info["labels"] = self.labels
|
|||
|
|
# 尝试读取 ultralytics yaml 的 imgsz 与 task
|
|||
|
|
base, _ = os.path.splitext(self.model_file)
|
|||
|
|
for p in (base + ".yaml", os.path.join(os.path.dirname(self.model_file), "model.yaml")):
|
|||
|
|
if os.path.exists(p):
|
|||
|
|
try:
|
|||
|
|
import yaml
|
|||
|
|
with open(p, "r", encoding="utf-8") as f:
|
|||
|
|
cfg = yaml.safe_load(f) or {}
|
|||
|
|
if "imgsz" in cfg:
|
|||
|
|
s = int(cfg["imgsz"]) if not isinstance(cfg["imgsz"], (list, tuple)) else int(cfg["imgsz"][0])
|
|||
|
|
info["input_size_inferred"] = (s, s)
|
|||
|
|
if "task" in cfg and not self.task_type:
|
|||
|
|
info["task_type"] = str(cfg["task"])
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
break
|
|||
|
|
except Exception as e:
|
|||
|
|
info["error"] = str(e)
|
|||
|
|
return info
|