105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
"""Ultralytics YOLO detector.
|
||
|
||
Wraps an Ultralytics YOLO model (detection or segmentation) behind the
|
||
:class:`Detector` interface. Designed for the LADA mosaic-detection weights
|
||
(https://huggingface.co/ladaapp/lada), which are YOLO *segmentation* models with
|
||
a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose class
|
||
names map onto :class:`CensorType`.
|
||
|
||
Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the
|
||
rest of the app — and the classic-CV detector — never pull them in.
|
||
|
||
Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
import numpy as np
|
||
|
||
from ...config import DetectionConfig
|
||
from ..video.frame import Frame
|
||
from .base import Detector
|
||
from .types import CensorType, Detection
|
||
|
||
|
||
def _name_to_type(name: str) -> CensorType:
|
||
n = name.lower()
|
||
if "mosaic" in n or "pixel" in n:
|
||
return CensorType.MOSAIC
|
||
if "blur" in n:
|
||
return CensorType.BLUR
|
||
if "bar" in n or "black" in n:
|
||
return CensorType.BLACK_BAR
|
||
return CensorType.UNKNOWN
|
||
|
||
|
||
class YoloDetector(Detector):
|
||
def __init__(self, model_path: str, config: DetectionConfig | None = None) -> None:
|
||
self.cfg = config or DetectionConfig()
|
||
if not os.path.isfile(model_path):
|
||
raise FileNotFoundError(
|
||
f"Файл весов не найден: {model_path}\n"
|
||
"Скачайте модель детекции мозаики LADA, например:\n"
|
||
" curl.exe -L -o models\\lada_mosaic_detection_model_v4_accurate.pt "
|
||
'"https://huggingface.co/ladaapp/lada/resolve/main/'
|
||
'lada_mosaic_detection_model_v4_accurate.pt?download=true"'
|
||
)
|
||
try:
|
||
from ultralytics import YOLO
|
||
except ImportError as exc: # pragma: no cover - environment dependent
|
||
raise ImportError(
|
||
"Не установлен ultralytics. Установите: pip install -e \".[yolo]\" "
|
||
"(и PyTorch с CUDA отдельно — см. README)."
|
||
) from exc
|
||
|
||
# Resolve the device: explicit override, else CUDA when available.
|
||
device = self.cfg.yolo_device
|
||
if device is None:
|
||
try:
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
except Exception: # noqa: BLE001
|
||
device = "cpu"
|
||
self._device = device
|
||
self._model = YOLO(model_path)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return f"YoloDetector(device={self._device})"
|
||
|
||
def detect(self, frame: Frame) -> list[Detection]:
|
||
results = self._model.predict(
|
||
source=frame.image, # BGR ndarray; ultralytics handles it
|
||
conf=self.cfg.yolo_conf,
|
||
imgsz=self.cfg.yolo_imgsz,
|
||
device=self._device,
|
||
verbose=False,
|
||
)
|
||
if not results:
|
||
return []
|
||
res = results[0]
|
||
boxes = getattr(res, "boxes", None)
|
||
if boxes is None or len(boxes) == 0:
|
||
return []
|
||
|
||
names = res.names # {class_index: class_name}
|
||
xyxy = boxes.xyxy.cpu().numpy()
|
||
confs = boxes.conf.cpu().numpy()
|
||
classes = boxes.cls.cpu().numpy().astype(int)
|
||
# Segmentation polygons in source-pixel coords, one per detection (if any).
|
||
polygons = res.masks.xy if getattr(res, "masks", None) is not None else None
|
||
|
||
out: list[Detection] = []
|
||
for i in range(len(xyxy)):
|
||
x1, y1, x2, y2 = xyxy[i]
|
||
bbox = (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
|
||
poly: list[tuple[int, int]] = []
|
||
if polygons is not None and i < len(polygons):
|
||
poly = [(int(px), int(py)) for px, py in polygons[i]]
|
||
ctype = _name_to_type(names.get(int(classes[i]), ""))
|
||
out.append(Detection(type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly))
|
||
return out
|