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.

hvideotool/ui/image_view.py

@@ -8,6 +8,8 @@ what the detector found.
 
 from __future__ import annotations
 
+import zlib
+
 import cv2
 import numpy as np
 from PySide6.QtCore import QPointF, QRectF, Qt
@@ -15,7 +17,7 @@ from PySide6.QtGui import QBrush, QColor, QFont, QImage, QPainter, QPen, QPolygo
 from PySide6.QtWidgets import QWidget
 
 from ..config import OverlayConfig
-from ..core.detection.types import CensorType, Detection
+from ..core.detection.types import Detection
 
 
 class ImageView(QWidget):
@@ -50,9 +52,13 @@ class ImageView(QWidget):
         self.update()
 
     # ------------------------------------------------------------------ paint
-    def _color(self, ctype: CensorType) -> QColor:
-        r, g, b = self._cfg.colors.get(ctype.value, (255, 0, 0))
-        return QColor(r, g, b)
+    def _color(self, key: str) -> QColor:
+        """Colour for a detection category — fixed if configured, else a stable palette pick."""
+        rgb = self._cfg.colors.get(key)
+        if rgb is None:
+            palette = self._cfg.palette
+            rgb = palette[zlib.crc32(key.encode("utf-8")) % len(palette)]
+        return QColor(*rgb)
 
     def paintEvent(self, event) -> None:
         painter = QPainter(self)
@@ -87,7 +93,7 @@ class ImageView(QWidget):
         self, painter: QPainter, d: Detection, ox: float, oy: float,
         scale: float, highlighted: bool, dim: bool,
     ) -> None:
-        color = self._color(d.type)
+        color = self._color(d.display)
         width = self._cfg.line_width * (2 if highlighted else 1)
         pen_color = QColor(color)
         if dim:
@@ -103,7 +109,7 @@ class ImageView(QWidget):
 
         if self._cfg.show_labels and not dim:
             x, y, _w, _h = d.bbox
-            self._draw_label(painter, f"{d.type.value} {d.score:.2f}", ox + x * scale, oy + y * scale, color)
+            self._draw_label(painter, f"{d.display} {d.score:.2f}", ox + x * scale, oy + y * scale, color)
 
     @staticmethod
     def _bbox_points(bbox: tuple[int, int, int, int]) -> list[tuple[int, int]]:

hvideotool/ui/main_window.py

@@ -23,6 +23,7 @@ project; switching the model clears the cache.
 from __future__ import annotations
 
 import contextlib
+import os
 import shutil
 from pathlib import Path
 
@@ -40,6 +41,7 @@ from PySide6.QtWidgets import (
     QListWidget,
     QListWidgetItem,
     QMainWindow,
+    QMenu,
     QMessageBox,
     QPlainTextEdit,
     QProgressBar,
@@ -47,13 +49,15 @@ from PySide6.QtWidgets import (
     QSplitter,
     QTableWidget,
     QTableWidgetItem,
+    QToolButton,
     QVBoxLayout,
     QWidget,
 )
 
 from .. import settings_store
-from ..config import AppConfig, normalize_config
+from ..config import AppConfig
 from ..core.detection import cache as detection_cache
+from ..core.detection import registry as model_registry
 from ..core.detection.factory import build_detector
 from ..core.detection.types import Detection
 from ..core.imageio import imread_unicode, imwrite_unicode
@@ -137,10 +141,14 @@ class MainWindow(QMainWindow):
         tb.addAction(from_video)
         tb.addSeparator()
 
-        tb.addWidget(QLabel(" Детектор: YOLO "))
-        self.model_action = QAction("Модель…", self, triggered=self._choose_model)
-        self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики")
-        tb.addAction(self.model_action)
+        tb.addWidget(QLabel(" Детекторы: "))
+        self._models_menu = QMenu(self)
+        self.models_button = QToolButton()
+        self.models_button.setPopupMode(QToolButton.InstantPopup)
+        self.models_button.setMenu(self._models_menu)
+        self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)")
+        tb.addWidget(self.models_button)
+        self._rebuild_models_menu()
 
         tb.addSeparator()
         calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
@@ -218,7 +226,7 @@ class MainWindow(QMainWindow):
         self.detail_header.setWordWrap(True)
         rlayout.addWidget(self.detail_header)
         self.detail_table = QTableWidget(0, 4)
-        self.detail_table.setHorizontalHeaderLabels(["Тип", "Увер.", "BBox (x,y,w,h)", "Полигон"])
+        self.detail_table.setHorizontalHeaderLabels(["Категория", "Увер.", "BBox (x,y,w,h)", "Полигон"])
         self.detail_table.verticalHeader().setVisible(False)
         self.detail_table.setSelectionBehavior(QTableWidget.SelectRows)
         self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers)
@@ -421,7 +429,7 @@ class MainWindow(QMainWindow):
         self._cancel = False
         self.stop_action.setEnabled(True)
         # Disable inputs that would race a running job (they clear cache / rebuild engines).
-        self.model_action.setEnabled(False)
+        self.models_button.setEnabled(False)
         if total is None:
             self.progress.setRange(0, 0)  # indeterminate
         else:
@@ -432,7 +440,7 @@ class MainWindow(QMainWindow):
     def _end_busy(self) -> None:
         self._busy = False
         self.stop_action.setEnabled(False)
-        self.model_action.setEnabled(True)
+        self.models_button.setEnabled(True)
         self.progress.setVisible(False)
         self.progress.setRange(0, 100)  # leave it determinate for the next user
 
@@ -484,50 +492,97 @@ class MainWindow(QMainWindow):
     # --------------------------------------------------------------- detector
     def _make_detector(self):
         d = self._cfg.detection
-        key = (self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz)
+        key = (tuple(sorted(self._cfg.detector_models)), d.yolo_conf, d.yolo_imgsz)
         if key != self._detector_key:
             self._detector = build_detector(self._cfg)  # may raise ValueError / import / file errors
             self._detector_key = key
         return self._detector
 
-    def _ensure_model(self) -> None:
-        """Make sure the YOLO detector has weights — auto-pick from ./models silently.
+    # --------------------------------------------------------- model selection
+    def _ensure_models(self) -> None:
+        """Drop selected models that vanished; default to all discovered if none picked.
 
-        Called on project open. Doesn't prompt (the user can pick via "Модель…"); the
-        detector factory raises a clear message if a detect is attempted without one.
+        Called on project open. The factory raises a clear message if a detect runs with
+        nothing selected, so we don't prompt here.
         """
-        if self._cfg.model_path and Path(self._cfg.model_path).is_file():
+        entries = model_registry.discover_models()
+        available = {e.path for e in entries}
+        kept = [m for m in self._cfg.detector_models if m in available]
+        if not kept and entries:  # first run / fresh project — enable everything found
+            kept = [e.path for e in entries]
+        if kept != self._cfg.detector_models:
+            self._cfg.detector_models = kept
+            self._persist_settings()
+
+    def _rebuild_models_menu(self) -> None:
+        """Repopulate the toolbar "Модели" menu with a checkable item per discovered model."""
+        self._models_menu.clear()
+        selected = set(self._cfg.detector_models)
+        entries = model_registry.discover_models()
+        if not entries:
+            self._models_menu.addAction("(нет моделей в models/yolo/<категория>)").setEnabled(False)
+        else:
+            last_cat = None
+            for e in entries:
+                if e.category != last_cat:
+                    self._models_menu.addSection(e.category)
+                    last_cat = e.category
+                act = self._models_menu.addAction(e.name)
+                act.setCheckable(True)
+                act.setChecked(e.path in selected)
+                act.toggled.connect(lambda on, p=e.path: self._on_model_toggled(p, on))
+        self._models_menu.addSeparator()
+        self._models_menu.addAction("Добавить модель…", self._add_model)
+        self._models_menu.addAction("Открыть папку моделей", self._open_models_dir)
+        self._models_menu.addAction("Обновить список", self._rebuild_models_menu)
+        self._update_models_button()
+
+    def _update_models_button(self) -> None:
+        n = len([m for m in self._cfg.detector_models if Path(m).is_file()])
+        self.models_button.setText(f"Модели ({n}) ▾")
+
+    def _on_model_toggled(self, path: str, on: bool) -> None:
+        sel = [m for m in self._cfg.detector_models if m != path]
+        if on:
+            sel.append(path)
+        self._cfg.detector_models = sel
+        self._persist_settings()
+        self._invalidate_results()  # different model set => recompute
+        self._update_models_button()
+
+    def _add_model(self) -> None:
+        """Copy a chosen .pt into models/yolo/<category>/ and tick it."""
+        path, _ = QFileDialog.getOpenFileName(
+            self, "Выберите веса YOLO (.pt)", str(Path.cwd()), "Веса YOLO (*.pt)"
+        )
+        if not path:
             return
-        found = self._auto_find_model()
-        if found:
-            self._cfg.model_path = found
-            self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}")
-            self._persist_settings()
+        category, ok = QInputDialog.getText(
+            self, "Категория модели",
+            "Категория (папка под models/yolo, напр. mosaic, face):", text="misc"
+        )
+        if not ok:
+            return
+        category = (category.strip() or "misc")
+        dest_dir = model_registry.yolo_root() / category
+        dest_dir.mkdir(parents=True, exist_ok=True)
+        dest = self._unique_dest(dest_dir, Path(path).name)
+        try:
+            shutil.copy2(path, dest)
+        except OSError as exc:
+            QMessageBox.warning(self, "Ошибка", f"Не удалось скопировать модель:\n{exc}")
+            return
+        self._cfg.detector_models = [*self._cfg.detector_models, str(dest)]
+        self._persist_settings()
+        self._rebuild_models_menu()
+        self._invalidate_results()
+        self.statusBar().showMessage(f"Модель добавлена: {dest.name} → {category}")
 
-    @staticmethod
-    def _auto_find_model() -> str | None:
-        """Find a censorship YOLO model under ./models without prompting.
-
-        Matches LADA/mosaic weights by filename; deliberately ignores generic COCO
-        models (e.g. yolo11n-seg.pt) that would map objects to purple "noise".
-        """
-        models_dir = Path.cwd() / "models"
-        if not models_dir.is_dir():
-            return None
-        for p in sorted(models_dir.rglob("*.pt")):
-            name = p.name.lower()
-            if "lada" in name or "mosaic" in name:
-                return str(p)
-        return None
-
-    def _choose_model(self) -> None:
-        start = self._cfg.model_path or str(Path.cwd() / "models")
-        path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)")
-        if path:
-            self._cfg.model_path = path
-            self._persist_settings()
-            self.statusBar().showMessage(f"Модель: {path}")
-            self._invalidate_results()
+    def _open_models_dir(self) -> None:
+        root = model_registry.yolo_root()
+        root.mkdir(parents=True, exist_ok=True)
+        with contextlib.suppress(OSError, AttributeError):
+            os.startfile(str(root))  # noqa: S606 - Windows: open in Explorer
 
     def _invalidate_results(self) -> None:
         """Detector changed — drop the in-memory cache and refresh the current image.
@@ -694,8 +749,7 @@ class MainWindow(QMainWindow):
         self._project = project
         project.frames_dir.mkdir(parents=True, exist_ok=True)
         project.apply_to_config(self._cfg)  # per-project settings -> live config
-        normalize_config(self._cfg)         # coerce any legacy classic/inpaint values
-        self._ensure_model()                # YOLO needs weights — auto-pick if missing
+        self._ensure_models()               # prune missing / default-select discovered models
         self._sync_settings_ui()
         self._detector_key = None
         self._restorer_key = None
@@ -708,6 +762,7 @@ class MainWindow(QMainWindow):
 
     def _sync_settings_ui(self) -> None:
         """Reflect the (project's) config onto the toolbar widgets without signal loops."""
+        self._rebuild_models_menu()  # reflect this project's model selection
         self.threshold_spin.blockSignals(True)
         self.threshold_spin.setValue(self._cfg.default_threshold)
         self.threshold_spin.blockSignals(False)
@@ -1177,7 +1232,7 @@ class MainWindow(QMainWindow):
         """Detector identity used to tag/validate the on-disk detection cache."""
         d = self._cfg.detection
         return detection_cache.make_key(
-            self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz
+            self._cfg.detector_models, d.yolo_conf, d.yolo_imgsz
         )
 
     def _save_results(self) -> None:
@@ -1235,7 +1290,7 @@ class MainWindow(QMainWindow):
 
         by_type: dict[str, int] = {}
         for d in dets:
-            by_type[d.type.value] = by_type.get(d.type.value, 0) + 1
+            by_type[d.display] = by_type.get(d.display, 0) + 1
         summary = ", ".join(f"{k}: {v}" for k, v in sorted(by_type.items())) or "ничего не найдено"
         self.detail_header.setText(f"<b>{path.name}</b> · {w}×{h} · всего {len(dets)} ({summary})")
 
@@ -1243,7 +1298,7 @@ class MainWindow(QMainWindow):
         self.detail_table.setRowCount(len(dets))
         for row, d in enumerate(dets):
             x, y, bw, bh = d.bbox
-            cells = [d.type.value, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
+            cells = [d.display, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
             for col, text in enumerate(cells):
                 self.detail_table.setItem(row, col, QTableWidgetItem(text))
         self.detail_table.blockSignals(False)