Implement multi-model detection in HVideoTool: updated the detection system to support multiple YOLO models simultaneously, enhancing detection capabilities. Reflected changes in the UI with a new model selection menu and updated documentation in README and CLAUDE.md to guide users on model management and configuration.
This commit is contained in:
Leonid Pershin
@@ -24,11 +24,13 @@ from .types import Detection
|
||||
_VERSION = 1
|
||||
|
||||
|
||||
def make_key(detector: str, model_path: str | None, yolo_conf: float, yolo_imgsz: int) -> dict:
|
||||
"""Identity of the detector that produced a cache; cache is only valid for a match."""
|
||||
def make_key(models: list[str], yolo_conf: float, yolo_imgsz: int) -> dict:
|
||||
"""Identity of the detector set that produced a cache; cache is only valid for a match.
|
||||
|
||||
Keyed by the (sorted) model **basenames** so it's portable across machines/paths.
|
||||
"""
|
||||
return {
|
||||
"detector": detector,
|
||||
"model_path": model_path or "",
|
||||
"models": sorted(Path(m).name for m in models),
|
||||
"yolo_conf": round(float(yolo_conf), 4),
|
||||
"yolo_imgsz": int(yolo_imgsz),
|
||||
}
|
||||
|
||||
@@ -4,26 +4,33 @@ Kept separate from ``app.py`` so both the app bootstrap and the UI can build
|
||||
detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on
|
||||
bad config so the GUI can show the message instead of exiting.
|
||||
|
||||
Only the YOLO detector is supported — the classic-CV heuristic (and the composite
|
||||
mode that combined them) were removed: they were noisy/approximate on real footage.
|
||||
YOLO-only, but multi-model: every path in ``config.detector_models`` (ticked under
|
||||
models/yolo/<category>/) becomes a YoloDetector tagged with its category, and they
|
||||
run together via :class:`~.multi.MultiYoloDetector`. The classic-CV detector and the
|
||||
``combined`` composite were removed (noisy/approximate on real footage).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from ...config import AppConfig
|
||||
from .base import Detector
|
||||
|
||||
|
||||
def _require_model(config: AppConfig) -> str:
|
||||
if not config.model_path:
|
||||
raise ValueError(
|
||||
"Для детектора YOLO укажите путь к весам (.pt) в Параметрах "
|
||||
"или скачайте модель LADA — см. README."
|
||||
)
|
||||
return config.model_path
|
||||
from .registry import category_of
|
||||
|
||||
|
||||
def build_detector(config: AppConfig) -> Detector:
|
||||
from .yolo import YoloDetector # lazy: pulls torch/ultralytics
|
||||
models = [m for m in config.detector_models if Path(m).is_file()]
|
||||
if not models:
|
||||
raise ValueError(
|
||||
"Не выбрана ни одна модель детекции.\n"
|
||||
"Положите веса YOLO в models/yolo/<категория>/ (например models/yolo/mosaic/) "
|
||||
"и отметьте их галочкой в меню «Модели». См. README."
|
||||
)
|
||||
# lazy imports: YoloDetector pulls in torch/ultralytics only when a detect runs
|
||||
from .multi import MultiYoloDetector
|
||||
from .yolo import YoloDetector
|
||||
|
||||
return YoloDetector(_require_model(config), config.detection)
|
||||
return MultiYoloDetector([
|
||||
YoloDetector(m, config.detection, label=category_of(m)) for m in models
|
||||
])
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Run several detectors over a frame and merge their detections.
|
||||
|
||||
Used for the multi-model (ADetailer-style) setup: each ticked ``models/yolo/<cat>/*.pt``
|
||||
becomes a :class:`~.yolo.YoloDetector` (tagged with its category), and this detector
|
||||
concatenates all their results. Detections keep their own ``label`` (category), so the
|
||||
overlay/table show every model's output together — no cross-model dedup (different
|
||||
categories are meant to coexist).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ..video.frame import Frame
|
||||
from .base import Detector
|
||||
from .types import Detection
|
||||
|
||||
|
||||
class MultiYoloDetector(Detector):
|
||||
def __init__(self, detectors: list[Detector]) -> None:
|
||||
if not detectors:
|
||||
raise ValueError("MultiYoloDetector requires at least one detector")
|
||||
self._detectors = detectors
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "Multi(" + " + ".join(d.name for d in self._detectors) + ")"
|
||||
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
out: list[Detection] = []
|
||||
for d in self._detectors:
|
||||
out.extend(d.detect(frame))
|
||||
return out
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Discover the YOLO models available for detection.
|
||||
|
||||
ADetailer-style layout: drop weights under ``models/yolo/<category>/*.pt``. The
|
||||
*category* (the sub-folder) becomes the detection label and its overlay colour, so
|
||||
e.g. ``models/yolo/mosaic/lada.pt`` tags its boxes "mosaic" and ``models/yolo/face/
|
||||
yolov8n-face.pt`` tags "face". The user ticks which discovered models are active; a
|
||||
detect runs every ticked model and merges the results (see ``multi.MultiYoloDetector``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
YOLO_SUBDIR = ("models", "yolo") # relative to the working directory
|
||||
_UNCATEGORIZED = "misc" # category for a .pt sitting directly under models/yolo
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelEntry:
|
||||
path: str # absolute path to the .pt
|
||||
category: str # sub-folder under models/yolo (the detection label / colour key)
|
||||
name: str # filename stem (display name)
|
||||
|
||||
|
||||
def yolo_root(root: Path | None = None) -> Path:
|
||||
return (root or Path.cwd()).joinpath(*YOLO_SUBDIR)
|
||||
|
||||
|
||||
def _category_of(pt: Path, root: Path) -> str:
|
||||
rel = pt.relative_to(root).parts
|
||||
return rel[0] if len(rel) > 1 else _UNCATEGORIZED
|
||||
|
||||
|
||||
def discover_models(root: Path | None = None) -> list[ModelEntry]:
|
||||
"""All ``*.pt`` under ``models/yolo/**``, sorted by (category, name)."""
|
||||
base = yolo_root(root)
|
||||
if not base.is_dir():
|
||||
return []
|
||||
out = [
|
||||
ModelEntry(path=str(p), category=_category_of(p, base), name=p.stem)
|
||||
for p in base.rglob("*.pt")
|
||||
]
|
||||
out.sort(key=lambda e: (e.category.lower(), e.name.lower()))
|
||||
return out
|
||||
|
||||
|
||||
def category_of(model_path: str, root: Path | None = None) -> str:
|
||||
"""Category (label) for a model path, derived from its folder under models/yolo."""
|
||||
base = yolo_root(root)
|
||||
p = Path(model_path)
|
||||
try:
|
||||
return _category_of(p, base)
|
||||
except ValueError: # outside models/yolo — fall back to the parent folder name
|
||||
return p.parent.name or _UNCATEGORIZED
|
||||
@@ -17,12 +17,18 @@ class CensorType(StrEnum):
|
||||
|
||||
@dataclass
|
||||
class Detection:
|
||||
"""A single detected censored region, in source-frame pixel coordinates."""
|
||||
"""A single detected region, in source-frame pixel coordinates."""
|
||||
|
||||
type: CensorType
|
||||
score: float # confidence, 0..1
|
||||
bbox: tuple[int, int, int, int] # x, y, w, h
|
||||
polygon: list[tuple[int, int]] = field(default_factory=list) # contour points
|
||||
label: str = "" # model category (models/yolo/<label>); drives colour/grouping
|
||||
|
||||
@property
|
||||
def display(self) -> str:
|
||||
"""Label shown in the UI — the category if set, else the CensorType."""
|
||||
return self.label or self.type.value
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
@@ -30,6 +36,7 @@ class Detection:
|
||||
"score": self.score,
|
||||
"bbox": list(self.bbox),
|
||||
"polygon": [list(p) for p in self.polygon],
|
||||
"label": self.label,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
@@ -39,4 +46,5 @@ class Detection:
|
||||
score=float(data["score"]),
|
||||
bbox=tuple(data["bbox"]), # type: ignore[arg-type]
|
||||
polygon=[tuple(p) for p in data.get("polygon", [])],
|
||||
label=data.get("label", ""),
|
||||
)
|
||||
|
||||
@@ -34,8 +34,11 @@ def _name_to_type(name: str) -> CensorType:
|
||||
|
||||
|
||||
class YoloDetector(Detector):
|
||||
def __init__(self, model_path: str, config: DetectionConfig | None = None) -> None:
|
||||
def __init__(
|
||||
self, model_path: str, config: DetectionConfig | None = None, label: str = ""
|
||||
) -> None:
|
||||
self.cfg = config or DetectionConfig()
|
||||
self._label = label # category (models/yolo/<label>) tagged onto every detection
|
||||
if not os.path.isfile(model_path):
|
||||
raise FileNotFoundError(
|
||||
f"Файл весов не найден: {model_path}\n"
|
||||
@@ -71,7 +74,8 @@ class YoloDetector(Detector):
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"YoloDetector(device={self._device})"
|
||||
tag = f", {self._label}" if self._label else ""
|
||||
return f"YoloDetector(device={self._device}{tag})"
|
||||
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
results = self._model.predict(
|
||||
@@ -103,5 +107,7 @@ class YoloDetector(Detector):
|
||||
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))
|
||||
out.append(Detection(
|
||||
type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly, label=self._label
|
||||
))
|
||||
return out
|
||||
|
||||
@@ -36,7 +36,7 @@ _VERSION = 1
|
||||
# The subset of AppConfig fields a project remembers (mirrored to/from project.json).
|
||||
_SETTING_KEYS = (
|
||||
"detector",
|
||||
"model_path",
|
||||
"detector_models",
|
||||
"default_threshold",
|
||||
"restorer",
|
||||
"dm_dir",
|
||||
|
||||
Reference in New Issue
Block a user