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
2026-06-07 07:05:50 +03:00
parent ac02ca27a8
commit 0996ca7bb9
14 changed files with 353 additions and 156 deletions
+41 -25
View File
@@ -51,7 +51,10 @@ Keep this scope sharp:
batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were
**removed** as "works poorly". DeepMosaics has two engines: **image** (per-frame) and
**removed** as "works poorly". Detection is **multi-model** (ADetailer-style): drop YOLO
weights under `models/yolo/<category>/`, tick which ones are active in the toolbar
"Модели" menu, and a detect runs **all** ticked models and merges results (each tagged
with its category → its own overlay colour). DeepMosaics has two engines: **image** (per-frame) and
**video** (BVDNet, temporal — uses neighbour frames). Its GPL-3.0 network code is
**vendored** under `core/restore/_deepmosaics/` and run in-process (user supplies only
the weights). Because of that vendoring the **whole project is GPL-3.0**. LADA
@@ -79,7 +82,7 @@ Keep this scope sharp:
|----------------|---------------------------------------------|
| GUI | PySide6 (Qt 6) — LGPL |
| Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` |
| Detector | Ultralytics YOLO (LADA weights) behind a pluggable interface (YOLO-only) |
| Detector | Ultralytics YOLO, multi-model ensemble (models/yolo/<category>/*.pt) |
Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional
(the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed
@@ -123,19 +126,31 @@ hvideotool/
│ ├── factory.py # build_restorer(name, config) -> deepmosaics | deepmosaics_video (lada = TODO)
│ ├── deepmosaics.py # DeepMosaicsRestorer (image, per-frame) + DeepMosaicsVideoRestorer (BVDNet, temporal); in-process, load once; uses _deepmosaics/
│ └── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
└── detection/ # YOLO only
└── detection/ # YOLO only, multi-model
├── base.py # Detector ABC: detect(frame) -> list[Detection]
├── factory.py # build_detector(config) -> yolo (the only kind)
├── types.py # Detection (+ to_dict/from_dict), CensorType enum
├── cache.py # save/load the project detection cache (detections.json): cache_file + base_dir args
── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics
├── factory.py # build_detector(config) -> MultiYoloDetector over config.detector_models
├── registry.py # discover_models()/category_of() — scans models/yolo/<category>/*.pt
├── multi.py # MultiYoloDetector — runs several YoloDetectors, concatenates results
── types.py # Detection (type + score + bbox + polygon + label/category; .display); CensorType enum
├── cache.py # save/load the detection cache (detections.json); key = set of model basenames + conf/imgsz
└── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy torch/ultralytics; tags dets with a category label
```
### How it works
- `MainWindow` holds the config, builds the detector lazily via `build_detector`
(cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path
-> list[Detection]]` as the detection cache.
(cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps
`_results: dict[path -> list[Detection]]` as the detection cache.
- **Multi-model selection.** `config.detector_models` is the list of ticked YOLO weights
(paths under `models/yolo/<category>/`). The toolbar "Модели" `QToolButton`/`QMenu`
(`_rebuild_models_menu`, items grouped by category via `addSection`) toggles them
(`_on_model_toggled` → persist + `_invalidate_results`); "Добавить модель…" copies a
`.pt` into `models/yolo/<category>/`. On project open `_ensure_models` prunes vanished
paths and, if nothing is selected, default-ticks every discovered model. `build_detector`
builds one `YoloDetector` per selected model (tagged `label=category`) wrapped in a
`MultiYoloDetector` that concatenates their detections (no cross-model dedup). Each
`Detection` carries `label` (category); overlay colour + table group by `Detection.display`
(label, else the CensorType) via `OverlayConfig.colors` + a stable `palette` fallback.
- **Background jobs (`ui/workers.py`).** Detection and restoration are CPU-heavy and
would freeze the GUI, so they run on a `QThreadPool` thread via `Job` (a `QRunnable`
wrapping `fn(job)`); results return to the GUI through queued Qt signals
@@ -276,11 +291,13 @@ hvideotool/
- `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
`core/detection/factory.build_detector` and talks to `core/` through the `Detector`
interface and the `Detection`/`CensorType` types.
- Detection is YOLO-only and restoration is DeepMosaics-only. If you re-add an engine
kind, implement `core/detection/base.Detector` / `core/restore/base.Restorer`, register
the string in the respective `factory`, and add it to `DETECTORS`/`RESTORERS` in
`config.py` (and `normalize_config`). There's no detector dropdown anymore — the toolbar
just shows "Детектор: YOLO"; restoration engines are chosen in `RestoreDialog`.
- Detection is YOLO-only (multi-model) and restoration is DeepMosaics-only. New detection
kinds plug in by adding more `.pt` under `models/yolo/<category>/` — no code change. The
toolbar shows a "Модели" menu of checkable models (no detector dropdown); restoration
engines are chosen in `RestoreDialog`. If you re-add a different engine *kind*, implement
`core/detection/base.Detector` / `core/restore/base.Restorer` and register it in the
respective `factory`. (No legacy-settings migration is kept while in active development —
old `settings.json`/`project.json` keys are simply ignored, not coerced.)
## Commands
@@ -308,20 +325,19 @@ frame directly.
## Gotchas
- **Wrong YOLO model = "noise".** The YOLO detector needs a **censorship** model
(LADA `models\lada_mosaic_detection_model_v4_accurate.pt`). If `model_path` points at
a generic COCO model (e.g. the `yolo11n-seg.pt` in the repo root, which Ultralytics
auto-downloads / is the training base), it detects people/objects and maps them to
`CensorType.UNKNOWN` → purple boxes that look like noise. This was a real user trap.
**A model is auto-picked on project open** via `MainWindow._ensure_model()`
`_auto_find_model()`: it scans `./models/**.pt` and matches only filenames containing
`lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; otherwise the
user picks via "Модель…" (`_choose_model`).
- **Models live under `models/yolo/<category>/`.** Discovery (`detection/registry.py`)
scans that tree; the **category folder is the detection label + overlay colour** (e.g.
`models/yolo/mosaic/lada.pt` → "mosaic", `models/yolo/face/…` → "face"). On open
`_ensure_models` default-ticks all discovered models if the project has no selection.
Put a **censorship** model in `mosaic/` (LADA `lada_mosaic_detection_model_v4_accurate.pt`);
a generic COCO model (e.g. `yolo11n-seg.pt`) would detect people/objects → noise.
Selecting many models multiplies per-frame time (each runs in turn).
- **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was
noisy/approximate on real footage (false positives on skin/hair/fabric/JPEG; missed
real mosaic after downscale) and the `combined` mode + cv2 `inpaint` baseline went with
it. Detection is YOLO-only, restoration is DeepMosaics-only. `normalize_config` coerces
any leftover `classic`/`combined`/`inpaint` in old settings/projects to `yolo`/`deepmosaics`.
it. Detection is YOLO-only, restoration is DeepMosaics-only. No backward-compat shims
while in active development — stale keys in old `settings.json`/`project.json` are just
ignored (a project with no valid model selection default-ticks all discovered models).
- **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic
but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`).
- **YOLO detector = LADA weights** ([HF `ladaapp/lada`](https://huggingface.co/ladaapp/lada)).
+41 -26
View File
@@ -64,10 +64,11 @@
кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры
с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена,
🟢 проверено и чисто).
- Детекция — через **YOLO** (модель LADA для мозаики). Порог уверенности
настраивается прямо в тулбаре.
- Выбор файла весов модели кнопкой **«Модель…»** (нужная модель ищется в `models/`
автоматически при открытии проекта).
- Детекция — через **YOLO**, **несколько моделей сразу** (как ADetailer): сложите веса
в `models/yolo/<категория>/` (например `models/yolo/mosaic/`, `models/yolo/face/`) и
отметьте нужные галочками в меню **«Модели»** тулбара. При расчёте кадра прогоняются
все выбранные модели, их детекции попадают в общую таблицу и рисуются оверлеем —
**цвет по категории-папке**. Порог уверенности — в тулбаре.
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU
@@ -82,8 +83,8 @@
- **дозапуск** — «Детектировать все» пропускает уже посчитанные кадры;
- **отказоустойчивость** — при отмене/закрытии посчитанное не теряется.
Кэш помечен «удостоверением» детектора (детектор + модель + порог `conf`/`imgsz`).
Если открыть проект **другим** детектором, чужой кэш не загружается (чтобы не выдавать
Кэш помечен «удостоверением» **набором выбранных моделей** + `conf`/`imgsz`. Если
открыть проект с **другим набором моделей**, чужой кэш не загружается (чтобы не выдавать
старые результаты за текущие). Полностью пересчитать — кнопка **«Все заново»**.
> Ключи внутри файла — **имена файлов**, поэтому кэш переживает перемещение/
@@ -187,19 +188,30 @@ python -m hvideotool "C:\path\to\МойПроект" --model models\lada_mosaic_
---
## Модель детектора
## Модели детекции (мульти-YOLO)
Детекция — **только YOLO** (интерфейс абстрагирован в `core/detection/base.py`):
ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics)
(`core/detection/yolo.py`), **сегментационная** модель — маски превращаются в
контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada).
Детекция — **только YOLO**, но можно держать и включать **несколько моделей сразу**
(как ADetailer). Структура папок:
```
models/yolo/
├── mosaic/ # модель(и) детекции мозаики (LADA) — категория "mosaic", красный цвет
├── face/ # напр. yolov8-face — категория "face", зелёный цвет
└── <своё>/ # любая категория = имя папки = ярлык + цвет
```
Каждая модель — **сегментационная** YOLO ([Ultralytics](https://github.com/ultralytics/ultralytics),
`core/detection/yolo.py`), маски превращаются в контуры. Включайте модели галочками в
меню **«Модели»** тулбара (там же «Добавить модель…» — скопирует `.pt` в нужную
категорию). При расчёте кадра прогоняются **все включённые** модели, детекции
объединяются (`core/detection/multi.py`), **цвет и ярлык — по категории-папке**.
Рекомендуемые веса для мозаики — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada).
(Старый эвристический classic-CV детектор и комбинированный режим удалены — давали
много ложных срабатываний.)
> **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры**
> (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать
> обычную COCO-модель (`yolo11n-seg.pt`), она будет детектить людей/предметы и
> помечать их как `unknown` — это и есть «шум». Файл LADA лежит в `models/`.
> **⚠️ Берите правильную модель для мозаики.** В `models/yolo/mosaic/` нужна модель
> **детекции цензуры** (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Обычная
> COCO-модель (`yolo11n-seg.pt`) детектит людей/предметы — это «шум».
> **Домен важен.** Модель LADA обучена на **реальном видео** (JAV); на части
> рисованного/аниме контента работает, на части — плохо. Хорошего публичного
@@ -213,7 +225,7 @@ ML-детектор на базе [Ultralytics](https://github.com/ultralytics/u
```powershell
python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic
python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100
# затем: тулбар → «Модель…» → runs\segment\mosaic\weights\best.pt
# затем скопируйте runs\segment\mosaic\weights\best.pt в models\yolo\mosaic\ и включите галочкой
```
Установка YOLO-детектора:
@@ -222,10 +234,11 @@ python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs
pip install -e ".[yolo]"
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
curl.exe -L -o models\lada_mosaic_detection_model_v4_accurate.pt `
New-Item -ItemType Directory -Force models\yolo\mosaic | Out-Null
curl.exe -L -o models\yolo\mosaic\lada_mosaic_detection_model_v4_accurate.pt `
"https://huggingface.co/ladaapp/lada/resolve/main/lada_mosaic_detection_model_v4_accurate.pt?download=true"
python -m hvideotool --model models\lada_mosaic_detection_model_v4_accurate.pt
python -m hvideotool # модель из models\yolo\mosaic подхватится и включится автоматически
```
> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**; код DeepMosaics —
@@ -239,7 +252,7 @@ python -m hvideotool --model models\lada_mosaic_detection_model_v4_accurate.pt
hvideotool/
├── __main__.py # точка входа + CLI (всё опционально)
├── app.py # инициализация QApplication
├── config.py # настройки: пороги детекции, оверлей, детектор/модель
├── config.py # настройки: порог/оверлей, detector_models (мульти-YOLO), движок восстановления
├── settings_store.py # дефолты новых проектов + последний/недавние → settings.json
├── ui/
│ ├── main_window.py # окно: список файлов | картинка | таблица детекций
@@ -248,13 +261,15 @@ hvideotool/
└── core/
├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути)
├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки
├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора
├── detection/ # только YOLO
├── video/frame.py # Frame (картинка BGR + индекс) — вход детектора
├── detection/ # только YOLO, мульти-модель
│ ├── base.py # Detector (ABC): detect(frame) -> list[Detection]
│ ├── factory.py # build_detector(config) -> yolo
│ ├── types.py # Detection (+ to_dict/from_dict), CensorType
│ ├── cache.py # сохранение/загрузка кэша детекций (detections.json)
── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны)
│ ├── factory.py # build_detector -> MultiYoloDetector по выбранным моделям
│ ├── registry.py # поиск моделей в models/yolo/<категория>/*.pt
│ ├── multi.py # MultiYoloDetector: прогон нескольких моделей + объединение
── types.py # Detection (+ label/категория, .display), CensorType
│ ├── cache.py # кэш детекций (ключ = набор моделей + conf/imgsz)
│ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны, ярлык категории)
└── restore/ # только DeepMosaics
├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal
├── factory.py # build_restorer -> deepmosaics | deepmosaics_video
@@ -273,7 +288,7 @@ hvideotool/
| Язык | Python ≥ 3.11 |
| GUI | PySide6 (Qt 6) |
| Обработка картинок | OpenCV / NumPy |
| Детектор | Ultralytics YOLO + PyTorch/CUDA (веса LADA) |
| Детектор | Ultralytics YOLO, мульти-модель (models/yolo/<кат>) |
| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA |
## Лицензия
+4 -5
View File
@@ -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)
+15 -18
View File
@@ -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
"black_bar": (26, 188, 156), # teal
"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"
+6 -4
View File
@@ -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),
}
+20 -13
View File
@@ -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
])
+31
View File
@@ -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
+55
View File
@@ -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
+9 -1
View File
@@ -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", ""),
)
+9 -3
View File
@@ -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
+1 -1
View File
@@ -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",
+3 -3
View File
@@ -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,
+12 -6
View File
@@ -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]]:
+102 -47
View File
@@ -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}")
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()
@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._rebuild_models_menu()
self._invalidate_results()
self.statusBar().showMessage(f"Модель добавлена: {dest.name}{category}")
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)