Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.

hvideotool/core/detection/yolo.py

@@ -0,0 +1,104 @@
+"""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