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:
@@ -12,7 +12,7 @@ import sys
|
||||
|
||||
from . import settings_store
|
||||
from .app import run
|
||||
from .config import AppConfig, normalize_config
|
||||
from .config import AppConfig
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -21,15 +21,14 @@ def main() -> int:
|
||||
description="Инспектор детекции уже наложенной цензуры на картинках.",
|
||||
)
|
||||
parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)")
|
||||
parser.add_argument("--model", dest="model_path", default=None, help="путь к весам (YOLO)")
|
||||
parser.add_argument("--model", dest="model", default=None, help="путь к весам YOLO (.pt) — использовать только эту модель")
|
||||
args = parser.parse_args()
|
||||
|
||||
config = AppConfig()
|
||||
settings_store.apply(config) # persisted defaults first
|
||||
normalize_config(config) # drop any legacy classic/inpaint values
|
||||
|
||||
if args.model_path:
|
||||
config.model_path = args.model_path
|
||||
if args.model:
|
||||
config.detector_models = [args.model]
|
||||
|
||||
return run(config, target=args.target)
|
||||
|
||||
|
||||
+18
-21
@@ -8,9 +8,6 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
DETECTORS = ("yolo",)
|
||||
RESTORERS = ("deepmosaics", "deepmosaics_video")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DetectionConfig:
|
||||
@@ -25,15 +22,26 @@ class DetectionConfig:
|
||||
class OverlayConfig:
|
||||
"""How detections are drawn over the image."""
|
||||
|
||||
# RGB per CensorType value
|
||||
# RGB per detection category/label (the models/yolo/<category> folder name, or the
|
||||
# CensorType for legacy detections). Unknown categories get a stable palette colour.
|
||||
colors: dict[str, tuple[int, int, int]] = field(
|
||||
default_factory=lambda: {
|
||||
"mosaic": (231, 76, 60), # red
|
||||
"blur": (241, 196, 15), # yellow
|
||||
"mosaic": (231, 76, 60), # red
|
||||
"blur": (241, 196, 15), # yellow
|
||||
"black_bar": (26, 188, 156), # teal
|
||||
"unknown": (155, 89, 182), # purple
|
||||
"unknown": (155, 89, 182), # purple
|
||||
"face": (46, 204, 113), # green
|
||||
"hand": (52, 152, 219), # blue
|
||||
"person": (230, 126, 34), # orange
|
||||
"eyes": (155, 89, 182), # purple
|
||||
"text": (149, 165, 166), # grey
|
||||
}
|
||||
)
|
||||
# Fallback colours cycled (deterministically) for categories not listed above.
|
||||
palette: tuple[tuple[int, int, int], ...] = (
|
||||
(231, 76, 60), (46, 204, 113), (52, 152, 219), (241, 196, 15),
|
||||
(155, 89, 182), (26, 188, 156), (230, 126, 34), (149, 165, 166),
|
||||
)
|
||||
line_width: int = 2
|
||||
fill_alpha: int = 48 # 0..255 translucency of the region fill
|
||||
show_labels: bool = True
|
||||
@@ -44,7 +52,9 @@ class AppConfig:
|
||||
detection: DetectionConfig = field(default_factory=DetectionConfig)
|
||||
overlay: OverlayConfig = field(default_factory=OverlayConfig)
|
||||
detector: str = "yolo" # only "yolo"
|
||||
model_path: str | None = None # weights path, used by the YOLO detector
|
||||
# Active YOLO models (paths under models/yolo/<category>/). A detect runs every
|
||||
# selected model and merges results — see core/detection/multi.MultiYoloDetector.
|
||||
detector_models: list[str] = field(default_factory=list)
|
||||
default_threshold: float = 0.20 # initial overlay confidence threshold
|
||||
|
||||
# --- restoration ("расцензурить") ---
|
||||
@@ -52,16 +62,3 @@ class AppConfig:
|
||||
dm_dir: str | None = None # optional extra dir to search for mosaic_position.pth
|
||||
dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
|
||||
dm_gpu: str = "0" # CUDA device id, "-1" for CPU
|
||||
|
||||
|
||||
def normalize_config(cfg: AppConfig) -> None:
|
||||
"""Coerce legacy/removed settings to supported values (mutates ``cfg``).
|
||||
|
||||
Old projects / settings.json may carry the removed ``classic``/``combined``
|
||||
detectors or the ``inpaint`` restorer — map those onto the survivors so loading
|
||||
them doesn't blow up at build time.
|
||||
"""
|
||||
if cfg.detector not in DETECTORS:
|
||||
cfg.detector = "yolo"
|
||||
if cfg.restorer not in RESTORERS:
|
||||
cfg.restorer = "deepmosaics"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -33,8 +33,8 @@ def apply(config: AppConfig) -> None:
|
||||
data = _read()
|
||||
if data.get("detector"):
|
||||
config.detector = data["detector"]
|
||||
if "model_path" in data:
|
||||
config.model_path = data["model_path"]
|
||||
if isinstance(data.get("detector_models"), list):
|
||||
config.detector_models = [str(m) for m in data["detector_models"]]
|
||||
if "threshold" in data:
|
||||
config.default_threshold = float(data["threshold"])
|
||||
if data.get("restorer"):
|
||||
@@ -49,7 +49,7 @@ def save(config: AppConfig) -> None:
|
||||
data = _read()
|
||||
data.update(
|
||||
detector=config.detector,
|
||||
model_path=config.model_path,
|
||||
detector_models=list(config.detector_models),
|
||||
threshold=config.default_threshold,
|
||||
restorer=config.restorer,
|
||||
dm_dir=config.dm_dir,
|
||||
|
||||
@@ -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]]:
|
||||
|
||||
+103
-48
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user