mrleo1nid/HVideoTool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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.

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+41 -25
View File
@@ -51,7 +51,10 @@ Keep this scope sharp:
batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined` restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were 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 **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 **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 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 | | GUI | PySide6 (Qt 6) — LGPL |
| Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` | | 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 Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional
(the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed (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) │ ├── 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.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 │ └── _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] ├── base.py # Detector ABC: detect(frame) -> list[Detection]
├── factory.py # build_detector(config) -> yolo (the only kind) ├── factory.py # build_detector(config) -> MultiYoloDetector over config.detector_models
├── types.py # Detection (+ to_dict/from_dict), CensorType enum ├── registry.py # discover_models()/category_of() — scans models/yolo/<category>/*.pt
├── cache.py # save/load the project detection cache (detections.json): cache_file + base_dir args ├── multi.py # MultiYoloDetector — runs several YoloDetectors, concatenates results
└── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics ├── 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 ### How it works
- `MainWindow` holds the config, builds the detector lazily via `build_detector` - `MainWindow` holds the config, builds the detector lazily via `build_detector`
(cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path (cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps
-> list[Detection]]` as the detection cache. `_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 - **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` 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 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 - `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
`core/detection/factory.build_detector` and talks to `core/` through the `Detector` `core/detection/factory.build_detector` and talks to `core/` through the `Detector`
interface and the `Detection`/`CensorType` types. interface and the `Detection`/`CensorType` types.
- Detection is YOLO-only and restoration is DeepMosaics-only. If you re-add an engine - Detection is YOLO-only (multi-model) and restoration is DeepMosaics-only. New detection
kind, implement `core/detection/base.Detector` / `core/restore/base.Restorer`, register kinds plug in by adding more `.pt` under `models/yolo/<category>/` — no code change. The
the string in the respective `factory`, and add it to `DETECTORS`/`RESTORERS` in toolbar shows a "Модели" menu of checkable models (no detector dropdown); restoration
`config.py` (and `normalize_config`). There's no detector dropdown anymore — the toolbar engines are chosen in `RestoreDialog`. If you re-add a different engine *kind*, implement
just shows "Детектор: YOLO"; restoration engines are chosen in `RestoreDialog`. `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 ## Commands
@@ -308,20 +325,19 @@ frame directly.
## Gotchas ## Gotchas
- **Wrong YOLO model = "noise".** The YOLO detector needs a **censorship** model - **Models live under `models/yolo/<category>/`.** Discovery (`detection/registry.py`)
(LADA `models\lada_mosaic_detection_model_v4_accurate.pt`). If `model_path` points at scans that tree; the **category folder is the detection label + overlay colour** (e.g.
a generic COCO model (e.g. the `yolo11n-seg.pt` in the repo root, which Ultralytics `models/yolo/mosaic/lada.pt` → "mosaic", `models/yolo/face/…` → "face"). On open
auto-downloads / is the training base), it detects people/objects and maps them to `_ensure_models` default-ticks all discovered models if the project has no selection.
`CensorType.UNKNOWN` → purple boxes that look like noise. This was a real user trap. Put a **censorship** model in `mosaic/` (LADA `lada_mosaic_detection_model_v4_accurate.pt`);
**A model is auto-picked on project open** via `MainWindow._ensure_model()` → a generic COCO model (e.g. `yolo11n-seg.pt`) would detect people/objects → noise.
`_auto_find_model()`: it scans `./models/**.pt` and matches only filenames containing Selecting many models multiplies per-frame time (each runs in turn).
`lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; otherwise the
user picks via "Модель…" (`_choose_model`).
- **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was - **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 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 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 it. Detection is YOLO-only, restoration is DeepMosaics-only. No backward-compat shims
any leftover `classic`/`combined`/`inpaint` in old settings/projects to `yolo`/`deepmosaics`. 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 - **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/`). 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)). - **YOLO detector = LADA weights** ([HF `ladaapp/lada`](https://huggingface.co/ladaapp/lada)).
README.md
+41 -26
View File
@@ -64,10 +64,11 @@
кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры
с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена, с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена,
🟢 проверено и чисто). 🟢 проверено и чисто).
- Детекция — через **YOLO** (модель LADA для мозаики). Порог уверенности - Детекция — через **YOLO**, **несколько моделей сразу** (как ADetailer): сложите веса
настраивается прямо в тулбаре. в `models/yolo/<категория>/` (например `models/yolo/mosaic/`, `models/yolo/face/`) и
- Выбор файла весов модели кнопкой **«Модель…»** (нужная модель ищется в `models/` отметьте нужные галочками в меню **«Модели»** тулбара. При расчёте кадра прогоняются
автоматически при открытии проекта). все выбранные модели, их детекции попадают в общую таблицу и рисуются оверлеем —
**цвет по категории-папке**. Порог уверенности — в тулбаре.
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU» - **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU 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`): Детекция — **только YOLO**, но можно держать и включать **несколько моделей сразу**
ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics) (как ADetailer). Структура папок:
(`core/detection/yolo.py`), **сегментационная** модель — маски превращаются в
контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada). ```
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 детектор и комбинированный режим удалены — давали (Старый эвристический classic-CV детектор и комбинированный режим удалены — давали
много ложных срабатываний.) много ложных срабатываний.)
> **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры** > **⚠️ Берите правильную модель для мозаики.** В `models/yolo/mosaic/` нужна модель
> (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать > **детекции цензуры** (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Обычная
> обычную COCO-модель (`yolo11n-seg.pt`), она будет детектить людей/предметы и > COCO-модель (`yolo11n-seg.pt`) детектит людей/предметы — это «шум».
> помечать их как `unknown` — это и есть «шум». Файл LADA лежит в `models/`.
> **Домен важен.** Модель LADA обучена на **реальном видео** (JAV); на части > **Домен важен.** Модель LADA обучена на **реальном видео** (JAV); на части
> рисованного/аниме контента работает, на части — плохо. Хорошего публичного > рисованного/аниме контента работает, на части — плохо. Хорошего публичного
@@ -213,7 +225,7 @@ ML-детектор на базе [Ultralytics](https://github.com/ultralytics/u
```powershell ```powershell
python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic 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 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-детектора: Установка YOLO-детектора:
@@ -222,10 +234,11 @@ python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs
pip install -e ".[yolo]" pip install -e ".[yolo]"
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 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" "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 — > **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**; код DeepMosaics —
@@ -239,7 +252,7 @@ python -m hvideotool --model models\lada_mosaic_detection_model_v4_accurate.pt
hvideotool/ hvideotool/
├── __main__.py # точка входа + CLI (всё опционально) ├── __main__.py # точка входа + CLI (всё опционально)
├── app.py # инициализация QApplication ├── app.py # инициализация QApplication
├── config.py # настройки: пороги детекции, оверлей, детектор/модель ├── config.py # настройки: порог/оверлей, detector_models (мульти-YOLO), движок восстановления
├── settings_store.py # дефолты новых проектов + последний/недавние → settings.json ├── settings_store.py # дефолты новых проектов + последний/недавние → settings.json
├── ui/ ├── ui/
│ ├── main_window.py # окно: список файлов | картинка | таблица детекций │ ├── main_window.py # окно: список файлов | картинка | таблица детекций
@@ -248,13 +261,15 @@ hvideotool/
└── core/ └── core/
├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути) ├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути)
├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки ├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки
├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора ├── video/frame.py # Frame (картинка BGR + индекс) — вход детектора
├── detection/ # только YOLO ├── detection/ # только YOLO, мульти-модель
│ ├── base.py # Detector (ABC): detect(frame) -> list[Detection] │ ├── base.py # Detector (ABC): detect(frame) -> list[Detection]
│ ├── factory.py # build_detector(config) -> yolo │ ├── factory.py # build_detector -> MultiYoloDetector по выбранным моделям
│ ├── types.py # Detection (+ to_dict/from_dict), CensorType │ ├── registry.py # поиск моделей в models/yolo/<категория>/*.pt
│ ├── cache.py # сохранение/загрузка кэша детекций (detections.json) │ ├── multi.py # MultiYoloDetector: прогон нескольких моделей + объединение
│ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны) │ ├── types.py # Detection (+ label/категория, .display), CensorType
│ ├── cache.py # кэш детекций (ключ = набор моделей + conf/imgsz)
│ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны, ярлык категории)
└── restore/ # только DeepMosaics └── restore/ # только DeepMosaics
├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal ├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal
├── factory.py # build_restorer -> deepmosaics | deepmosaics_video ├── factory.py # build_restorer -> deepmosaics | deepmosaics_video
@@ -273,7 +288,7 @@ hvideotool/
| Язык | Python ≥ 3.11 | | Язык | Python ≥ 3.11 |
| GUI | PySide6 (Qt 6) | | GUI | PySide6 (Qt 6) |
| Обработка картинок | OpenCV / NumPy | | Обработка картинок | OpenCV / NumPy |
| Детектор | Ultralytics YOLO + PyTorch/CUDA (веса LADA) | | Детектор | Ultralytics YOLO, мульти-модель (models/yolo/<кат>) |
| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA | | Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA |
## Лицензия ## Лицензия
hvideotool/__main__.py
+4 -5
View File
@@ -12,7 +12,7 @@ import sys
from . import settings_store from . import settings_store
from .app import run from .app import run
from .config import AppConfig, normalize_config from .config import AppConfig
def main() -> int: def main() -> int:
@@ -21,15 +21,14 @@ def main() -> int:
description="Инспектор детекции уже наложенной цензуры на картинках.", description="Инспектор детекции уже наложенной цензуры на картинках.",
) )
parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)") 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() args = parser.parse_args()
config = AppConfig() config = AppConfig()
settings_store.apply(config) # persisted defaults first settings_store.apply(config) # persisted defaults first
normalize_config(config) # drop any legacy classic/inpaint values
if args.model_path: if args.model:
config.model_path = args.model_path config.detector_models = [args.model]
return run(config, target=args.target) return run(config, target=args.target)
hvideotool/config.py
+18 -21
View File
@@ -8,9 +8,6 @@ from __future__ import annotations
from dataclasses import dataclass, field from dataclasses import dataclass, field
DETECTORS = ("yolo",)
RESTORERS = ("deepmosaics", "deepmosaics_video")
@dataclass(frozen=True) @dataclass(frozen=True)
class DetectionConfig: class DetectionConfig:
@@ -25,15 +22,26 @@ class DetectionConfig:
class OverlayConfig: class OverlayConfig:
"""How detections are drawn over the image.""" """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( colors: dict[str, tuple[int, int, int]] = field(
default_factory=lambda: { default_factory=lambda: {
"mosaic": (231, 76, 60), # red "mosaic": (231, 76, 60), # red
"blur": (241, 196, 15), # yellow "blur": (241, 196, 15), # yellow
"black_bar": (26, 188, 156), # teal "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 line_width: int = 2
fill_alpha: int = 48 # 0..255 translucency of the region fill fill_alpha: int = 48 # 0..255 translucency of the region fill
show_labels: bool = True show_labels: bool = True
@@ -44,7 +52,9 @@ class AppConfig:
detection: DetectionConfig = field(default_factory=DetectionConfig) detection: DetectionConfig = field(default_factory=DetectionConfig)
overlay: OverlayConfig = field(default_factory=OverlayConfig) overlay: OverlayConfig = field(default_factory=OverlayConfig)
detector: str = "yolo" # only "yolo" 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 default_threshold: float = 0.20 # initial overlay confidence threshold
# --- restoration ("расцензурить") --- # --- restoration ("расцензурить") ---
@@ -52,16 +62,3 @@ class AppConfig:
dm_dir: str | None = None # optional extra dir to search for mosaic_position.pth 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_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
dm_gpu: str = "0" # CUDA device id, "-1" for CPU 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"
hvideotool/core/detection/cache.py
+6 -4
View File
@@ -24,11 +24,13 @@ from .types import Detection
_VERSION = 1 _VERSION = 1
def make_key(detector: str, model_path: str | None, yolo_conf: float, yolo_imgsz: int) -> dict: def make_key(models: list[str], yolo_conf: float, yolo_imgsz: int) -> dict:
"""Identity of the detector that produced a cache; cache is only valid for a match.""" """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 { return {
"detector": detector, "models": sorted(Path(m).name for m in models),
"model_path": model_path or "",
"yolo_conf": round(float(yolo_conf), 4), "yolo_conf": round(float(yolo_conf), 4),
"yolo_imgsz": int(yolo_imgsz), "yolo_imgsz": int(yolo_imgsz),
} }
hvideotool/core/detection/factory.py
+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 detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on
bad config so the GUI can show the message instead of exiting. 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 YOLO-only, but multi-model: every path in ``config.detector_models`` (ticked under
mode that combined them) were removed: they were noisy/approximate on real footage. 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 __future__ import annotations
from pathlib import Path
from ...config import AppConfig from ...config import AppConfig
from .base import Detector from .base import Detector
from .registry import category_of
def _require_model(config: AppConfig) -> str:
if not config.model_path:
raise ValueError(
"Для детектора YOLO укажите путь к весам (.pt) в Параметрах "
"или скачайте модель LADA — см. README."
)
return config.model_path
def build_detector(config: AppConfig) -> Detector: 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
])
hvideotool/core/detection/multi.py
+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
hvideotool/core/detection/registry.py
+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
hvideotool/core/detection/types.py
+9 -1
View File
@@ -17,12 +17,18 @@ class CensorType(StrEnum):
@dataclass @dataclass
class Detection: class Detection:
"""A single detected censored region, in source-frame pixel coordinates.""" """A single detected region, in source-frame pixel coordinates."""
type: CensorType type: CensorType
score: float # confidence, 0..1 score: float # confidence, 0..1
bbox: tuple[int, int, int, int] # x, y, w, h bbox: tuple[int, int, int, int] # x, y, w, h
polygon: list[tuple[int, int]] = field(default_factory=list) # contour points 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: def to_dict(self) -> dict:
return { return {
@@ -30,6 +36,7 @@ class Detection:
"score": self.score, "score": self.score,
"bbox": list(self.bbox), "bbox": list(self.bbox),
"polygon": [list(p) for p in self.polygon], "polygon": [list(p) for p in self.polygon],
"label": self.label,
} }
@classmethod @classmethod
@@ -39,4 +46,5 @@ class Detection:
score=float(data["score"]), score=float(data["score"]),
bbox=tuple(data["bbox"]), # type: ignore[arg-type] bbox=tuple(data["bbox"]), # type: ignore[arg-type]
polygon=[tuple(p) for p in data.get("polygon", [])], polygon=[tuple(p) for p in data.get("polygon", [])],
label=data.get("label", ""),
) )
hvideotool/core/detection/yolo.py
+9 -3
View File
@@ -34,8 +34,11 @@ def _name_to_type(name: str) -> CensorType:
class YoloDetector(Detector): 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.cfg = config or DetectionConfig()
self._label = label # category (models/yolo/<label>) tagged onto every detection
if not os.path.isfile(model_path): if not os.path.isfile(model_path):
raise FileNotFoundError( raise FileNotFoundError(
f"Файл весов не найден: {model_path}\n" f"Файл весов не найден: {model_path}\n"
@@ -71,7 +74,8 @@ class YoloDetector(Detector):
@property @property
def name(self) -> str: 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]: def detect(self, frame: Frame) -> list[Detection]:
results = self._model.predict( results = self._model.predict(
@@ -103,5 +107,7 @@ class YoloDetector(Detector):
if polygons is not None and i < len(polygons): if polygons is not None and i < len(polygons):
poly = [(int(px), int(py)) for px, py in polygons[i]] poly = [(int(px), int(py)) for px, py in polygons[i]]
ctype = _name_to_type(names.get(int(classes[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 return out
hvideotool/core/project.py
+1 -1
View File
@@ -36,7 +36,7 @@ _VERSION = 1
# The subset of AppConfig fields a project remembers (mirrored to/from project.json). # The subset of AppConfig fields a project remembers (mirrored to/from project.json).
_SETTING_KEYS = ( _SETTING_KEYS = (
"detector", "detector",
"model_path", "detector_models",
"default_threshold", "default_threshold",
"restorer", "restorer",
"dm_dir", "dm_dir",
hvideotool/settings_store.py
+3 -3
View File
@@ -33,8 +33,8 @@ def apply(config: AppConfig) -> None:
data = _read() data = _read()
if data.get("detector"): if data.get("detector"):
config.detector = data["detector"] config.detector = data["detector"]
if "model_path" in data: if isinstance(data.get("detector_models"), list):
config.model_path = data["model_path"] config.detector_models = [str(m) for m in data["detector_models"]]
if "threshold" in data: if "threshold" in data:
config.default_threshold = float(data["threshold"]) config.default_threshold = float(data["threshold"])
if data.get("restorer"): if data.get("restorer"):
@@ -49,7 +49,7 @@ def save(config: AppConfig) -> None:
data = _read() data = _read()
data.update( data.update(
detector=config.detector, detector=config.detector,
model_path=config.model_path, detector_models=list(config.detector_models),
threshold=config.default_threshold, threshold=config.default_threshold,
restorer=config.restorer, restorer=config.restorer,
dm_dir=config.dm_dir, dm_dir=config.dm_dir,
hvideotool/ui/image_view.py
+12 -6
View File
@@ -8,6 +8,8 @@ what the detector found.
from __future__ import annotations from __future__ import annotations
import zlib
import cv2 import cv2
import numpy as np import numpy as np
from PySide6.QtCore import QPointF, QRectF, Qt 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 PySide6.QtWidgets import QWidget
from ..config import OverlayConfig from ..config import OverlayConfig
from ..core.detection.types import CensorType, Detection from ..core.detection.types import Detection
class ImageView(QWidget): class ImageView(QWidget):
@@ -50,9 +52,13 @@ class ImageView(QWidget):
self.update() self.update()
# ------------------------------------------------------------------ paint # ------------------------------------------------------------------ paint
def _color(self, ctype: CensorType) -> QColor: def _color(self, key: str) -> QColor:
r, g, b = self._cfg.colors.get(ctype.value, (255, 0, 0)) """Colour for a detection category — fixed if configured, else a stable palette pick."""
return QColor(r, g, b) 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: def paintEvent(self, event) -> None:
painter = QPainter(self) painter = QPainter(self)
@@ -87,7 +93,7 @@ class ImageView(QWidget):
self, painter: QPainter, d: Detection, ox: float, oy: float, self, painter: QPainter, d: Detection, ox: float, oy: float,
scale: float, highlighted: bool, dim: bool, scale: float, highlighted: bool, dim: bool,
) -> None: ) -> None:
color = self._color(d.type) color = self._color(d.display)
width = self._cfg.line_width * (2 if highlighted else 1) width = self._cfg.line_width * (2 if highlighted else 1)
pen_color = QColor(color) pen_color = QColor(color)
if dim: if dim:
@@ -103,7 +109,7 @@ class ImageView(QWidget):
if self._cfg.show_labels and not dim: if self._cfg.show_labels and not dim:
x, y, _w, _h = d.bbox 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 @staticmethod
def _bbox_points(bbox: tuple[int, int, int, int]) -> list[tuple[int, int]]: def _bbox_points(bbox: tuple[int, int, int, int]) -> list[tuple[int, int]]:
hvideotool/ui/main_window.py
+103 -48
View File
@@ -23,6 +23,7 @@ project; switching the model clears the cache.
from __future__ import annotations from __future__ import annotations
import contextlib import contextlib
import os
import shutil import shutil
from pathlib import Path from pathlib import Path
@@ -40,6 +41,7 @@ from PySide6.QtWidgets import (
QListWidget, QListWidget,
QListWidgetItem, QListWidgetItem,
QMainWindow, QMainWindow,
QMenu,
QMessageBox, QMessageBox,
QPlainTextEdit, QPlainTextEdit,
QProgressBar, QProgressBar,
@@ -47,13 +49,15 @@ from PySide6.QtWidgets import (
QSplitter, QSplitter,
QTableWidget, QTableWidget,
QTableWidgetItem, QTableWidgetItem,
QToolButton,
QVBoxLayout, QVBoxLayout,
QWidget, QWidget,
) )
from .. import settings_store 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 cache as detection_cache
from ..core.detection import registry as model_registry
from ..core.detection.factory import build_detector from ..core.detection.factory import build_detector
from ..core.detection.types import Detection from ..core.detection.types import Detection
from ..core.imageio import imread_unicode, imwrite_unicode from ..core.imageio import imread_unicode, imwrite_unicode
@@ -137,10 +141,14 @@ class MainWindow(QMainWindow):
tb.addAction(from_video) tb.addAction(from_video)
tb.addSeparator() tb.addSeparator()
tb.addWidget(QLabel(" Детектор: YOLO ")) tb.addWidget(QLabel(" Детекторы: "))
self.model_action = QAction("Модель…", self, triggered=self._choose_model) self._models_menu = QMenu(self)
self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики") self.models_button = QToolButton()
tb.addAction(self.model_action) 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() tb.addSeparator()
calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current) calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
@@ -218,7 +226,7 @@ class MainWindow(QMainWindow):
self.detail_header.setWordWrap(True) self.detail_header.setWordWrap(True)
rlayout.addWidget(self.detail_header) rlayout.addWidget(self.detail_header)
self.detail_table = QTableWidget(0, 4) 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.verticalHeader().setVisible(False)
self.detail_table.setSelectionBehavior(QTableWidget.SelectRows) self.detail_table.setSelectionBehavior(QTableWidget.SelectRows)
self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers) self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers)
@@ -421,7 +429,7 @@ class MainWindow(QMainWindow):
self._cancel = False self._cancel = False
self.stop_action.setEnabled(True) self.stop_action.setEnabled(True)
# Disable inputs that would race a running job (they clear cache / rebuild engines). # 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: if total is None:
self.progress.setRange(0, 0) # indeterminate self.progress.setRange(0, 0) # indeterminate
else: else:
@@ -432,7 +440,7 @@ class MainWindow(QMainWindow):
def _end_busy(self) -> None: def _end_busy(self) -> None:
self._busy = False self._busy = False
self.stop_action.setEnabled(False) self.stop_action.setEnabled(False)
self.model_action.setEnabled(True) self.models_button.setEnabled(True)
self.progress.setVisible(False) self.progress.setVisible(False)
self.progress.setRange(0, 100) # leave it determinate for the next user self.progress.setRange(0, 100) # leave it determinate for the next user
@@ -484,50 +492,97 @@ class MainWindow(QMainWindow):
# --------------------------------------------------------------- detector # --------------------------------------------------------------- detector
def _make_detector(self): def _make_detector(self):
d = self._cfg.detection 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: if key != self._detector_key:
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
self._detector_key = key self._detector_key = key
return self._detector return self._detector
def _ensure_model(self) -> None: # --------------------------------------------------------- model selection
"""Make sure the YOLO detector has weights — auto-pick from ./models silently. 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 Called on project open. The factory raises a clear message if a detect runs with
detector factory raises a clear message if a detect is attempted without one. 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 return
found = self._auto_find_model() category, ok = QInputDialog.getText(
if found: self, "Категория модели",
self._cfg.model_path = found "Категория (папка под models/yolo, напр. mosaic, face):", text="misc"
self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}") )
self._persist_settings() 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 _open_models_dir(self) -> None:
def _auto_find_model() -> str | None: root = model_registry.yolo_root()
"""Find a censorship YOLO model under ./models without prompting. root.mkdir(parents=True, exist_ok=True)
with contextlib.suppress(OSError, AttributeError):
Matches LADA/mosaic weights by filename; deliberately ignores generic COCO os.startfile(str(root)) # noqa: S606 - Windows: open in Explorer
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 _invalidate_results(self) -> None: def _invalidate_results(self) -> None:
"""Detector changed — drop the in-memory cache and refresh the current image. """Detector changed — drop the in-memory cache and refresh the current image.
@@ -694,8 +749,7 @@ class MainWindow(QMainWindow):
self._project = project self._project = project
project.frames_dir.mkdir(parents=True, exist_ok=True) project.frames_dir.mkdir(parents=True, exist_ok=True)
project.apply_to_config(self._cfg) # per-project settings -> live config project.apply_to_config(self._cfg) # per-project settings -> live config
normalize_config(self._cfg) # coerce any legacy classic/inpaint values self._ensure_models() # prune missing / default-select discovered models
self._ensure_model() # YOLO needs weights — auto-pick if missing
self._sync_settings_ui() self._sync_settings_ui()
self._detector_key = None self._detector_key = None
self._restorer_key = None self._restorer_key = None
@@ -708,6 +762,7 @@ class MainWindow(QMainWindow):
def _sync_settings_ui(self) -> None: def _sync_settings_ui(self) -> None:
"""Reflect the (project's) config onto the toolbar widgets without signal loops.""" """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.blockSignals(True)
self.threshold_spin.setValue(self._cfg.default_threshold) self.threshold_spin.setValue(self._cfg.default_threshold)
self.threshold_spin.blockSignals(False) self.threshold_spin.blockSignals(False)
@@ -1177,7 +1232,7 @@ class MainWindow(QMainWindow):
"""Detector identity used to tag/validate the on-disk detection cache.""" """Detector identity used to tag/validate the on-disk detection cache."""
d = self._cfg.detection d = self._cfg.detection
return detection_cache.make_key( 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: def _save_results(self) -> None:
@@ -1235,7 +1290,7 @@ class MainWindow(QMainWindow):
by_type: dict[str, int] = {} by_type: dict[str, int] = {}
for d in dets: 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 "ничего не найдено" 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})") 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)) self.detail_table.setRowCount(len(dets))
for row, d in enumerate(dets): for row, d in enumerate(dets):
x, y, bw, bh = d.bbox 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): for col, text in enumerate(cells):
self.detail_table.setItem(row, col, QTableWidgetItem(text)) self.detail_table.setItem(row, col, QTableWidgetItem(text))
self.detail_table.blockSignals(False) self.detail_table.blockSignals(False)