188 lines
6.9 KiB
Python
188 lines
6.9 KiB
Python
|
"""PyTorch 引擎实现
|
|||
|
|
|
|||
|
|
支持加载 ultralytics YOLOv8/v5/v7 .pt 模型（依赖 ultralytics 包），
|
|||
|
|
或原生 torch.hub YOLOv5 custom 加载。
|
|||
|
|
依赖：torch, ultralytics（推荐）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from app.analysis.engines.base import BaseEngine, DetectionResult
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("analysis.engines.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
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PyTorchEngine(BaseEngine):
|
|||
|
|
ENGINE_NAME = "pytorch"
|
|||
|
|
|
|||
|
|
def __init__(self, **kwargs):
|
|||
|
|
super(PyTorchEngine, self).__init__(**kwargs)
|
|||
|
|
self._model = None
|
|||
|
|
self._kind = None # "ultralytics" / "torchhub"
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def is_available():
|
|||
|
|
return _TORCH_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_available():
|
|||
|
|
return _ULTRALYTICS_AVAILABLE
|
|||
|
|
|
|||
|
|
def load(self):
|
|||
|
|
if not self.is_available():
|
|||
|
|
logger.warning("PyTorchEngine: 依赖未安装 (torch=%s)", _TORCH_AVAILABLE)
|
|||
|
|
return False
|
|||
|
|
if not self.model_file or not os.path.exists(self.model_file):
|
|||
|
|
logger.warning("PyTorchEngine: 模型文件不存在: %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)
|
|||
|
|
|
|||
|
|
# 优先用 ultralytics 加载 .pt / .engine / .onnx
|
|||
|
|
if _ULTRALYTICS_AVAILABLE:
|
|||
|
|
try:
|
|||
|
|
self._model = _UltralyticsYOLO(self.model_file)
|
|||
|
|
self._kind = "ultralytics"
|
|||
|
|
self._loaded = True
|
|||
|
|
# 推断 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))
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
logger.info("PyTorchEngine(ultralytics): 已加载 %s, labels=%d", self.model_file, len(self.labels))
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("PyTorchEngine(ultralytics) 加载失败，尝试 torch.hub: %s", e)
|
|||
|
|
|
|||
|
|
# 退化：torch.hub YOLOv5 custom
|
|||
|
|
try:
|
|||
|
|
self._model = torch.hub.load("ultralytics/yolov5", "custom", path=self.model_file, trust_repo=True)
|
|||
|
|
self._kind = "torchhub"
|
|||
|
|
self._loaded = True
|
|||
|
|
logger.info("PyTorchEngine(torchhub): 已加载 %s", self.model_file)
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("PyTorchEngine: 加载失败: %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:
|
|||
|
|
if self._kind == "ultralytics":
|
|||
|
|
iw, ih = self.input_size
|
|||
|
|
res = self._model.predict(frame_bgr, imgsz=max(iw, ih), conf=self.conf_threshold,
|
|||
|
|
iou=self.iou_threshold, verbose=False)
|
|||
|
|
return self._parse_ultralytics(res, frame_bgr.shape[:2])
|
|||
|
|
else:
|
|||
|
|
# torchhub yolov5
|
|||
|
|
res = self._model(frame_bgr, size=max(self.input_size))
|
|||
|
|
return self._parse_torchhub(res, frame_bgr.shape[:2])
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("PyTorchEngine.detect() err: %s", e)
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
def _parse_ultralytics(self, results, orig_shape):
|
|||
|
|
out = []
|
|||
|
|
try:
|
|||
|
|
r = results[0]
|
|||
|
|
if hasattr(r, "boxes") and r.boxes is not None:
|
|||
|
|
boxes = r.boxes.xyxy.cpu().numpy().astype(int)
|
|||
|
|
confs = r.boxes.conf.cpu().numpy()
|
|||
|
|
cls_ids = r.boxes.cls.cpu().numpy().astype(int)
|
|||
|
|
names = getattr(r, "names", {}) or {}
|
|||
|
|
for (x1, y1, x2, y2), s, cid in zip(boxes, 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)))
|
|||
|
|
out.append(DetectionResult(box=[int(x1), int(y1), int(x2), int(y2)],
|
|||
|
|
label=str(label), score=float(s)))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("PyTorchEngine._parse_ultralytics err: %s", e)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
def _parse_torchhub(self, results, orig_shape):
|
|||
|
|
out = []
|
|||
|
|
try:
|
|||
|
|
df = results.pandas().xyxy[0]
|
|||
|
|
for _, row in df.iterrows():
|
|||
|
|
s = float(row["confidence"])
|
|||
|
|
if s < self.conf_threshold:
|
|||
|
|
continue
|
|||
|
|
out.append(DetectionResult(box=[int(row["xmin"]), int(row["ymin"]), int(row["xmax"]), int(row["ymax"])],
|
|||
|
|
label=str(row["name"]), score=s))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("PyTorchEngine._parse_torchhub err: %s", e)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
def info(self):
|
|||
|
|
d = super(PyTorchEngine, self).info()
|
|||
|
|
d["version"] = self.version()
|
|||
|
|
d["ultralytics_available"] = _ULTRALYTICS_AVAILABLE
|
|||
|
|
d["kind"] = self._kind
|
|||
|
|
d["device"] = "cuda" if (_TORCH_AVAILABLE and torch.cuda.is_available()) else "cpu"
|
|||
|
|
return d
|
|||
|
|
|
|||
|
|
def probe(self):
|
|||
|
|
info = {"engine": self.ENGINE_NAME, "available": self.is_available(),
|
|||
|
|
"version": self.version(), "ultralytics_available": _ULTRALYTICS_AVAILABLE,
|
|||
|
|
"input_shape": None, "output_shape": None,
|
|||
|
|
"labels": self.labels, "model_file": self.model_file}
|
|||
|
|
if not self.is_available() or not self.model_file or not os.path.exists(self.model_file):
|
|||
|
|
return info
|
|||
|
|
try:
|
|||
|
|
# 仅读文件大小，不实际加载（torch 模型加载慢且占内存）
|
|||
|
|
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
|
|||
|
|
except Exception as e:
|
|||
|
|
info["error"] = str(e)
|
|||
|
|
return info
|