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
|