diff --git a/CLAUDE.md b/CLAUDE.md index bf887e8..ccf65ba 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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//`, 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//*.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//*.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//`). 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//`. 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//` — 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//`.** 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)). diff --git a/README.md b/README.md index 0066fab..a37d3e0 100644 --- a/README.md +++ b/README.md @@ -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 | ## Лицензия diff --git a/hvideotool/__main__.py b/hvideotool/__main__.py index 071c343..bcd5510 100644 --- a/hvideotool/__main__.py +++ b/hvideotool/__main__.py @@ -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) diff --git a/hvideotool/config.py b/hvideotool/config.py index 3f911c0..223822d 100644 --- a/hvideotool/config.py +++ b/hvideotool/config.py @@ -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/ 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//). 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" diff --git a/hvideotool/core/detection/cache.py b/hvideotool/core/detection/cache.py index fb3dfe8..b28d656 100644 --- a/hvideotool/core/detection/cache.py +++ b/hvideotool/core/detection/cache.py @@ -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), } diff --git a/hvideotool/core/detection/factory.py b/hvideotool/core/detection/factory.py index 669299b..ca7a13c 100644 --- a/hvideotool/core/detection/factory.py +++ b/hvideotool/core/detection/factory.py @@ -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//) 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 + ]) diff --git a/hvideotool/core/detection/multi.py b/hvideotool/core/detection/multi.py new file mode 100644 index 0000000..8a6bbf0 --- /dev/null +++ b/hvideotool/core/detection/multi.py @@ -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//*.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 diff --git a/hvideotool/core/detection/registry.py b/hvideotool/core/detection/registry.py new file mode 100644 index 0000000..8a7d999 --- /dev/null +++ b/hvideotool/core/detection/registry.py @@ -0,0 +1,55 @@ +"""Discover the YOLO models available for detection. + +ADetailer-style layout: drop weights under ``models/yolo//*.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 diff --git a/hvideotool/core/detection/types.py b/hvideotool/core/detection/types.py index 15251c3..4b3ca80 100644 --- a/hvideotool/core/detection/types.py +++ b/hvideotool/core/detection/types.py @@ -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/