diff --git a/CLAUDE.md b/CLAUDE.md index eb3496b..aa494c5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,12 +47,15 @@ what it found. Keep this scope sharp: - Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить") - step was added later (user-requested): per-frame, on-demand, behind a `Restorer` - interface. Two engines: a cv2 **inpaint baseline** (fills, does NOT reconstruct) and - **DeepMosaics** — real generative mosaic removal, its GPL-3.0 network code **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 (BasicVSR++) is a - possible future engine, not wired. + step was added later (user-requested): on-demand (single frame **or** whole-project + 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 + **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 + (BasicVSR++) is a possible future engine, not wired. - Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was rejected early — a generative *conditioning* model, not a censorship restorer. Don't reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/ @@ -67,8 +70,8 @@ Keep this scope sharp: - **OS:** Windows 11 x64 (primary). Use PowerShell syntax in commands. - **Python:** 3.11+. -- **GPU:** NVIDIA + CUDA via PyTorch, only for the YOLO detector. CPU fallback works - but is slow. The `classic` detector needs no torch and no GPU. +- **GPU:** NVIDIA + CUDA via PyTorch, for the YOLO detector and the DeepMosaics + restorer. CPU fallback works but is slow (esp. DeepMosaics / the temporal BVDNet). ## Tech stack (decided) @@ -76,12 +79,12 @@ Keep this scope sharp: |----------------|---------------------------------------------| | GUI | PySide6 (Qt 6) — LGPL | | Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` | -| Detector | classic-CV heuristic; Ultralytics YOLO (LADA) behind a pluggable interface | +| Detector | Ultralytics YOLO (LADA weights) behind a pluggable interface (YOLO-only) | -Torch/CUDA + Ultralytics enter only with the YOLO detector. Keep that dependency -optional (the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is -installed separately per the README). The classic detector must keep running with no -torch present. +Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional +(the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed +separately per the README). Both detection (YOLO) and restoration (DeepMosaics) now +require torch — there's no longer a torch-free detector. ## Architecture (as implemented) @@ -115,21 +118,17 @@ hvideotool/ ├── video/ │ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type - ├── restore/ # "un-censor" detected regions (per-frame) - │ ├── base.py # Restorer ABC: restore(image, detections, should_cancel=None) -> image; Cancelled exc - │ ├── factory.py # build_restorer(name, config) -> inpaint | deepmosaics (lada = TODO) - │ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs - │ ├── deepmosaics.py # DeepMosaicsRestorer — in-process, loads models once; uses _deepmosaics/ - │ ├── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import - │ └── mask.py # detections_to_mask(shape, dets, dilate) - └── detection/ + ├── restore/ # "un-censor" (DeepMosaics only; per-frame OR temporal) + │ ├── base.py # Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc + │ ├── 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 ├── base.py # Detector ABC: detect(frame) -> list[Detection] - ├── factory.py # build_detector(config) -> classic | yolo | combined + ├── 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 - ├── classic_cv.py # ClassicCVDetector — heuristic; accepts a `types` filter - ├── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics - └── composite.py # CompositeDetector — merge detectors + IoU dedup + └── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics ``` ### How it works @@ -146,15 +145,21 @@ hvideotool/ `engine.detect/restore` all run **inside the worker** (`_compute` is the pure read+detect helper); the `fn` must touch NO Qt widgets — it emits plain data that the GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables - `detector_combo`/`model_action` for the duration (they'd race the running detector). - This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics - per-call latency can't be hidden with `processEvents`). -- **Device badge.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green) - or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background - `Job` at startup (importing torch is slow, so it's off the GUI thread) → `_set_device_badge`. - Clicking (`_show_device_info`) opens a diagnostic dialog: `torch_info.reason()` explains - why CPU (no torch / CPU-only `+cpu` build / built-with-CUDA-but-no-GPU) plus - `install_hint()` (the cu121 pip command). `core/torch_info.py` is pure (no Qt). + `model_action` for the duration (it'd race the running detector). This is the + deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics per-call + latency can't be hidden with `processEvents`). +- **Device badge + CUDA diagnostics.** A clickable status-bar chip (`device_badge`) shows + "⚡ CUDA" (green) or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in + a background `Job` at startup (it imports torch AND shells out to `nvidia-smi`, so it's off + the GUI thread) → `_set_device_badge`. `gather()` collects torch facts (version, built_cuda, + cuda_available, device_name) **and** NVIDIA facts (gpus, driver_version, max cuda_driver). + Clicking (`_show_device_info`) opens a diagnostic **QDialog** (not QMessageBox — its text + wasn't copyable): a read-only monospace `QPlainTextEdit` with `torch_info.analyze(info)` → + `{summary, details, steps, command}` — a verdict on *why* it's on CPU (CPU-only `+cpu` + build / no GPU / driver-too-old-for-built-CUDA) and the exact pip fix. Buttons: "Скопировать + команду установки" (`_copy_install_command` → the recommended cu121/cu118 command, picked by + `recommend_channel` from the driver's CUDA) and "Проверить заново" (re-runs `_probe_device`). + `core/torch_info.py` is pure (no Qt); subprocess uses `CREATE_NO_WINDOW` on Windows. - **Projects (`core/project.py`).** `MainWindow._project` is the open `Project`; `_folder` is kept as a synonym for `project.frames_dir` so the rest of the code (navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…" @@ -179,7 +184,7 @@ hvideotool/ "Рассчитать кадр" action (Space → `_recompute_current`, force-recomputes current), or "Детектировать все" (whole folder, progress bar). Do NOT re-add auto-detect-on-select. Results cache in `_results`; the file-list row gets a count suffix when computed. - Switching detector/model clears the cache (`_invalidate_results`). + Switching the model clears the cache (`_choose_model` → `_invalidate_results`). - **Detection cache (persisted).** `_results` is mirrored to the project's `detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by **basename** so it survives moving the project). `cache.save_results`/`load_results` @@ -203,27 +208,43 @@ hvideotool/ created lazily on first move), removing them from list/`_files`/cache. `_unique_dest` avoids clobbering (`foo.jpg` → `foo (1).jpg`). Use case: flag good frames into a training/example set while inspecting detections. -- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via - `build_restorer`) on the current frame's detections (computing them first if needed), - **on a background job** (`_restore_current` builds an `fn` that detects-if-needed + - restores in the worker; a `tick` caches freshly-computed detections, `done` stores - `_restored[path]` + shows it). "Показать оригинал/ - результат" toggles (`_showing_restored`); "Сохранить результат" writes - `_restored.jpg` beside the frame. The baseline - is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run - **in-process** from the vendored `_deepmosaics/` code: it loads the BiSeNet locator + - clean generator **once** (lazy, cached on the instance) and per frame reproduces their - `cleanmosaic_img_server` (locate mosaic → run generator on the crop → feather back), - ~0.3 s/frame cached on CPU vs ~7 s when it spawned a subprocess. Use the **image** - model `clean_youknow_resnet_9blocks.pth` — the video model (BVDNet) is rejected per - frame (needs a neighbour). `should_cancel` (= `lambda: job.cancelled`) is polled so - "■ Стоп" stops it; the engine raises `Cancelled`, which `Job.run` reports as a clean cancel. - The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…), - persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE: - DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the - clean weights) — our detections aren't passed to it. `_show` resets `_showing_restored` - + `_update_restore_actions`. To add another engine (e.g. LADA), implement - `core/restore/base.Restorer` and register it in `restore/factory.build_restorer`. +- **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it + **locates the mosaic itself** — so restoration is **fully decoupled from detection**: no + detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar + action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on + a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать + оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes + `_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново", + `_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every + frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`, + basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the + per-frame engine **skips frames already in `restored/`** unless `force` (resume). The + engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the + vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy, + cached on the instance): + - `deepmosaics` (image, per-frame): reproduces `cleanmosaic_img_server` (locate mosaic + → run generator on the crop → feather back), ~0.3 s/frame cached on CPU. Image model + `clean_youknow_resnet_9blocks.pth`. + - `deepmosaics_video` (**temporal, BVDNet**): `DeepMosaicsVideoRestorer`, `.temporal=True`. + Reproduces `cleanmosaic_video_fusion` — per target frame it feeds the net a window of + `T=5` neighbour frames sampled at step `S=3` around it (`N=2` each side, clamped at the + sequence edges) **plus its own previous output (recurrent)**, for temporal coherence. + Because of that recurrence it must run a **contiguous, ordered range** — it implements + `restore_sequence(count, get_frame, get_dets, emit, should_cancel)` (the batch run uses + it; single-frame `restore` degrades to a window of the same frame). Needs the **video** + weights `clean_youknow_video.pth` (+ `mosaic_position.pth` beside). `INPUT_SIZE=256`. + + `restore_sequence` is on the `Restorer` ABC (default = independent per-frame loop); + `_restore_all` dispatches on `restorer.temporal` (temporal → `restore_sequence` over the + whole range; per-frame → resumable loop with skip-existing). `should_cancel` + (= `lambda: job.cancelled`) is polled so "■ Стоп" stops it; engines raise `Cancelled`, + which `Job.run` reports as a clean cancel. Engine + weights are set in `RestoreDialog` + (Файл → Движок восстановления…) — the model dropdown shows image vs video weights per + selected engine — persisted, and built lazily/cached in `_make_restorer`. NOTE: + DeepMosaics locates mosaics itself (its `mosaic_position.pth`) — our detections aren't + passed to it. To add another engine (e.g. LADA), implement `core/restore/base.Restorer` + (set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in + `restore/factory.build_restorer`. - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, @@ -255,20 +276,20 @@ 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. -- New detector kinds: implement `core/detection/base.Detector`, register the string in - `core/detection/factory.build_detector`, and add it to `_DETECTORS` in - `ui/main_window.py`. +- 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`. ## Commands ```powershell python -m venv .venv; .\.venv\Scripts\Activate.ps1 -pip install -e . # classic detector needs no torch/CUDA +pip install -e ".[yolo]" # YOLO needs ultralytics; install torch separately (README) python -m hvideotool # reopen the last project (or create/open one in-app) -python -m hvideotool "C:\path\to\MyProject" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt - -pip install -e ".[yolo]" # + install torch separately, see README +python -m hvideotool "C:\path\to\MyProject" --model models\lada_mosaic_detection_model_v4_accurate.pt ``` No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a @@ -292,24 +313,22 @@ frame directly. 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. - **Switching to yolo/combined without a model auto-picks one** via - `MainWindow._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; it - falls back to the "Модель…" file dialog only when nothing suitable is found. -- **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low - block-reconstruction residual + 2D gradient + contrast) fires on textured real - footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing - real mosaic after the `proc_max_dim=720` downscale softens block edges (measured: - contrast/grad fall below `mosaic_contrast_min`/`mosaic_grad_min`). For real-video - mosaic use `yolo`/`combined` + LADA. For anime there is no good public model. + **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`). +- **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`. - **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/`), - not tuning more classic thresholds. + 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 **segmentation** model, classes `{0: mosaic_nsfw, 1: mosaic_sfw_head}` → both map to `CensorType.MOSAIC` (`_name_to_type` matches "mosaic" in the class name). Detects - mosaic only; black bars / blur stay with classic. Weights + Ultralytics are AGPL-3.0 - (accepted). `yolo.py` lazy-imports `torch`/`ultralytics`. + mosaic only. Weights + Ultralytics are AGPL-3.0 (accepted). `yolo.py` lazy-imports + `torch`/`ultralytics`. - **No model weights in the repo.** Code must fail with a clear, actionable message when the model path is missing — not a raw stack trace (`factory._require_model`, `YoloDetector.__init__`). @@ -318,10 +337,11 @@ frame directly. - **CPU-only torch must not request CUDA.** A `+cpu` torch build raises "Torch not compiled with CUDA enabled" the moment something calls `.cuda()`. Both engines guard for this: `YoloDetector` picks `cuda` only when `torch.cuda.is_available()` (even an - explicit `yolo_device="cuda"` is downgraded to cpu); `DeepMosaicsRestorer._ensure_loaded` - forces `gpu_id="-1"` when CUDA is absent (its vendored `model_util.todevice` / - `data.im2tensor` call `.cuda()` for any `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). - So a wrong/CPU-only torch falls back to CPU instead of crashing. + explicit `yolo_device="cuda"` is downgraded to cpu); both `DeepMosaicsRestorer._ensure_loaded` + and `DeepMosaicsVideoRestorer._ensure_loaded` force `gpu_id="-1"` when CUDA is absent (the + vendored `model_util.todevice` / `data.im2tensor`/`to_tensor` call `.cuda()` for any + `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). So a wrong/CPU-only torch falls back to + CPU instead of crashing (the temporal BVDNet engine is heavy on CPU, though). - **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`), otherwise it aliases a buffer that gets freed → garbage/crash. - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — diff --git a/README.md b/README.md index 83913f1..0066fab 100644 --- a/README.md +++ b/README.md @@ -64,9 +64,10 @@ кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена, 🟢 проверено и чисто). -- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** - уверенности прямо в тулбаре — удобно сравнивать. -- Выбор файла весов модели кнопкой **«Модель…»**. +- Детекция — через **YOLO** (модель LADA для мозаики). Порог уверенности + настраивается прямо в тулбаре. +- Выбор файла весов модели кнопкой **«Модель…»** (нужная модель ищется в `models/` + автоматически при открытии проекта). - **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU» показывает диагностику (почему GPU не задействован) и команды установки PyTorch с CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU @@ -96,14 +97,24 @@ `<имя>_restored.jpg` рядом с кадром. Движок выбирается в меню **Файл → Движок восстановления…**. -Два движка: +**Пакетный прогон по диапазону.** Кнопки **«Расцензурить все»** (дозапуск — пропускает +уже сделанные) и **«Все заново»** обрабатывают **все кадры проекта** в фоне и пишут +результаты в папку **`restored/`** проекта (имена кадров сохраняются; папка держится +отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс, +отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции. -- **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по - окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). -- **DeepMosaics** — реальное генеративное удаление мозаики. Код **встроен** в - приложение (vendored, GPL-3.0), ставить его отдельно не нужно — требуются только - **веса** и (желательно) **GPU NVIDIA/CUDA**. На аниме качество ограничено (модели - обучены на реальном видео). +Расцензуривание — только через **DeepMosaics** (код **встроен** в приложение, vendored, +GPL-3.0; ставить отдельно не нужно — требуются только **веса** и желательно **GPU +NVIDIA/CUDA**). Два движка: + +- **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**. +- **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры** + (окно ±2 кадра с шагом 3) и собственный предыдущий результат для когерентности на + роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е. + запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к + окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`. + +На аниме качество ограничено (модели обучены на реальном видео). ### Настройка DeepMosaics @@ -112,18 +123,20 @@ ([Google Drive](https://drive.google.com/drive/folders/1LTERcN33McoiztYEwBxMuRjjgxh4DEPs), Baidu код `1x0a`): -- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель; -- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом). +- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель (движок «картинка»); +- **`clean_youknow_video.pth`** — видеомодель BVDNet (движок «видео», соседние кадры); +- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом, нужен обоим). -Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, выберите **модель -из выпадающего списка** (наполняется из `models/deepmosaics`; есть «Обзор…» для файла в -другом месте) и GPU id (`-1` = CPU). Если веса в `models/deepmosaics` — работает сразу; -модель грузится один раз, дальше кадры считаются быстро. +Затем в приложении: **Файл → Движок восстановления…**, выберите движок (**DeepMosaics — +картинка** или **видео**) и **модель из выпадающего списка** (наполняется из +`models/deepmosaics` — для видеодвижка показываются только `clean_*_video.pth`; есть +«Обзор…» для файла в другом месте) и GPU id (`-1` = CPU). Если веса в `models/deepmosaics` +— работает сразу; модель грузится один раз, дальше кадры считаются быстро. -> **Важно:** берите именно **картиночную** модель `clean_youknow_resnet_9blocks.pth`. -> Видеомодель `clean_youknow_video.pth` (BVDNet) покадрово **не работает** — ей нужен -> соседний кадр (приложение это распознаёт и подскажет). Если на кадре нет мозаики, -> результат = исходный кадр. +> **Какую модель брать:** для покадрового движка — **картиночную** +> `clean_youknow_resnet_9blocks.pth`; для временно́го — **видео** `clean_youknow_video.pth` +> (она запускается только пакетно, «Расцензурить все», т.к. ей нужны соседние кадры). Если +> на кадре нет мозаики, результат = исходный кадр. > > Код DeepMosaics (GPL-3.0) лежит в `core/restore/_deepmosaics/` и поэтому **весь > проект распространяется под GPL-3.0**. Запускается на современных `torch 2.x`/ @@ -143,8 +156,8 @@ Baidu код `1x0a`): - **ОС:** Windows 11 x64 (основная целевая платформа). - **Python:** 3.11+. -- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора. CPU-режим работает, но - медленный. Для `classic` детектора ни torch, ни GPU не нужны. +- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора и DeepMosaics. CPU-режим + работает, но медленный (особенно DeepMosaics / временно́й BVDNet). ## Установка @@ -153,12 +166,11 @@ git clone https://github.com/mrleo1nid/HVideoTool.git cd HVideoTool python -m venv .venv .\.venv\Scripts\Activate.ps1 -pip install -e . +pip install -e ".[yolo]" ``` -Этого достаточно для `classic` детектора — **PyTorch/CUDA не требуются**. Они -нужны только для YOLO/комбинированного детектора (см. -[Модель детектора](#модель-детектора)). +Детекция (YOLO) и расцензуривание (DeepMosaics) требуют **PyTorch** — установите его +отдельно под вашу CUDA (см. [Модель детектора](#модель-детектора)). ## Запуск @@ -166,29 +178,23 @@ pip install -e . # Без аргументов — открывается последний проект (или создайте/откройте новый в тулбаре) python -m hvideotool -# Необязательно: сразу открыть проект / переопределить детектор и модель (по умолчанию) -python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +# Необязательно: сразу открыть проект и указать модель YOLO +python -m hvideotool "C:\path\to\МойПроект" --model models\lada_mosaic_detection_model_v4_accurate.pt ``` -Выбор детектора, путь к модели и порог сохраняются в `~/HVideoTool/settings.json` -и применяются при следующем запуске. +Путь к модели и порог сохраняются в `~/HVideoTool/settings.json` и применяются при +следующем запуске. --- ## Модель детектора -Интерфейс детектора абстрагирован (`core/detection/base.py`), детектор -выбирается в тулбаре (или флагом `--detector`): - -- `classic` — эвристический classic-CV детектор (`core/detection/classic_cv.py`). - Различает `mosaic` / `blur` / `black_bar`, без весов и без GPU. - **Приблизительный**: на реальном видео даёт много ложных срабатываний, заточен - скорее под рисованный/аниме контент — но и там ненадёжен. -- `yolo` — ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics) - (`core/detection/yolo.py`). **Сегментационная** модель — маски превращаются в - контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada). -- `combined` — `CompositeDetector`: YOLO (мозаика) + classic-CV (плашки/размытие), - результаты объединяются с дедупликацией по IoU. +Детекция — **только YOLO** (интерфейс абстрагирован в `core/detection/base.py`): +ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics) +(`core/detection/yolo.py`), **сегментационная** модель — маски превращаются в +контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada). +(Старый эвристический classic-CV детектор и комбинированный режим удалены — давали +много ложных срабатываний.) > **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры** > (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать @@ -207,7 +213,7 @@ python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model mod ```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 -# затем: тулбар → Детектор yolo → «Модель…» → runs\segment\mosaic\weights\best.pt +# затем: тулбар → «Модель…» → runs\segment\mosaic\weights\best.pt ``` Установка YOLO-детектора: @@ -219,11 +225,11 @@ pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 curl.exe -L -o models\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 --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +python -m hvideotool --model models\lada_mosaic_detection_model_v4_accurate.pt ``` -> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**. Classic-CV детектор -> от этого свободен. +> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**; код DeepMosaics — +> **GPL-3.0**. Поэтому весь проект распространяется под **GPL-3.0**. --- @@ -243,14 +249,17 @@ hvideotool/ ├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути) ├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки ├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора - └── detection/ - ├── base.py # Detector (ABC): detect(frame) -> list[Detection] - ├── factory.py # build_detector(config) -> classic/yolo/combined - ├── types.py # Detection (+ to_dict/from_dict), CensorType - ├── cache.py # сохранение/загрузка кэша детекций (detections.json в проекте) - ├── classic_cv.py # эвристический детектор (mosaic/blur/black_bar) - ├── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны) - └── composite.py # CompositeDetector: объединение детекторов + ├── 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, маски→полигоны) + └── restore/ # только DeepMosaics + ├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal + ├── factory.py # build_restorer -> deepmosaics | deepmosaics_video + ├── deepmosaics.py # движки «картинка» (покадрово) и «видео» (BVDNet) + └── _deepmosaics/ # встроенный код DeepMosaics (GPL-3.0) ``` Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU @@ -264,9 +273,9 @@ hvideotool/ | Язык | Python ≥ 3.11 | | GUI | PySide6 (Qt 6) | | Обработка картинок | OpenCV / NumPy | -| Детектор (без весов) | classic-CV эвристика (без GPU) | -| Детектор (ML) | Ultralytics YOLO + PyTorch/CUDA (LADA) | +| Детектор | Ultralytics YOLO + PyTorch/CUDA (веса LADA) | +| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA | ## Лицензия -TBD. +**GPL-3.0** — из-за встроенного кода DeepMosaics (GPL-3.0); веса/код YOLO LADA — AGPL-3.0. diff --git a/hvideotool/__main__.py b/hvideotool/__main__.py index ad8b735..071c343 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 +from .config import AppConfig, normalize_config def main() -> int: @@ -21,15 +21,13 @@ def main() -> int: description="Инспектор детекции уже наложенной цензуры на картинках.", ) parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)") - parser.add_argument("--detector", choices=["classic", "yolo", "combined"], default=None) parser.add_argument("--model", dest="model_path", default=None, help="путь к весам (YOLO)") args = parser.parse_args() config = AppConfig() settings_store.apply(config) # persisted defaults first + normalize_config(config) # drop any legacy classic/inpaint values - if args.detector: - config.detector = args.detector if args.model_path: config.model_path = args.model_path diff --git a/hvideotool/config.py b/hvideotool/config.py index 66c0075..0725d09 100644 --- a/hvideotool/config.py +++ b/hvideotool/config.py @@ -1,40 +1,21 @@ """Application configuration and tunable defaults. -Plain dataclasses. The detection thresholds matter most here — this tool is now -an image-folder inspector for tuning the detectors, so keep them easy to tweak. +Plain dataclasses. Detection is YOLO-only and restoration is DeepMosaics-only, so the +knobs here are the YOLO inference params, the overlay style, and the DeepMosaics weights. """ from __future__ import annotations from dataclasses import dataclass, field +DETECTORS = ("yolo",) +RESTORERS = ("deepmosaics", "deepmosaics_video") + @dataclass(frozen=True) class DetectionConfig: - """Parameters for the detector. Thresholds tuned for the classic-CV detector.""" + """Parameters for the YOLO detector.""" - proc_max_dim: int = 720 # downscale longer side to this before detection (speed) - min_area_frac: float = 0.0008 # ignore regions smaller than this fraction of the image - - # --- solid bars (black or white, achromatic, rectangular) --- - black_intensity: int = 40 # V below this = dark-bar candidate - white_intensity: int = 225 # V above this = light-bar candidate - bar_saturation_max: int = 45 # S below this = achromatic (excludes colored fills) - bar_min_extent: float = 0.80 # contour area / bbox area — how rectangular a bar must be - - # --- mosaic / pixelation --- - mosaic_block_sizes: tuple[int, ...] = (8, 12, 16, 24) # candidate tile sizes (px, proc space) - mosaic_residual_max: float = 6.0 # max reconstruction error to count as "blocky" - mosaic_contrast_min: float = 14.0 # min local contrast (excludes flat gradients) - mosaic_grad_min: float = 8.0 # min edge energy in BOTH x and y (excludes straight edges) - mosaic_min_side: int = 24 # reject thin regions (px) — kills edge false-positives - - # --- blur --- - blur_window: int = 31 # sliding window for local sharpness (odd) - blur_sharpness_ratio: float = 0.35 # below this fraction of median sharpness => blurry - blur_contrast_min: float = 8.0 # min local contrast (excludes flat regions) - - # --- YOLO detector (used only when detector == "yolo"/"combined") --- yolo_conf: float = 0.2 # confidence threshold (LADA recommends ~0.2) yolo_imgsz: int = 640 # inference image size yolo_device: str | None = None # None => auto ("cuda" if available, else "cpu") @@ -62,13 +43,26 @@ class OverlayConfig: class AppConfig: detection: DetectionConfig = field(default_factory=DetectionConfig) overlay: OverlayConfig = field(default_factory=OverlayConfig) - detector: str = "classic" # "classic" | "yolo" | "combined" + detector: str = "yolo" # only "yolo" model_path: str | None = None # weights path, used by the YOLO detector default_threshold: float = 0.20 # initial overlay confidence threshold # --- restoration ("расцензурить") --- - restorer: str = "inpaint" # "inpaint" | "deepmosaics" + restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video" dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py) dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth) dm_python: str | None = None # python exe for DeepMosaics (None = current) 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/classic_cv.py b/hvideotool/core/detection/classic_cv.py deleted file mode 100644 index cd52f2d..0000000 --- a/hvideotool/core/detection/classic_cv.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Weights-free, heuristic censorship detector (classic computer vision). - -APPROXIMATE BY DESIGN. This detector uses hand-tuned CV heuristics, not a -trained model. Its purpose is to make the whole pipeline runnable end-to-end -and to exercise the :class:`Detector` interface. For real-world accuracy, -replace it with a trained model (see ``yolo.py``, to be implemented) — the rest -of the app does not need to change. - -Heuristics: -- black_bar: large, near-uniform very dark regions (classic censor bars). -- mosaic: regions that reconstruct well from a coarse block grid (low - residual) yet have high coarse-scale contrast (i.e. blocky, not flat). -- blur: regions with local high-frequency energy far below the frame median, - while still being textured (excludes genuinely flat areas). -""" - -from __future__ import annotations - -import cv2 -import numpy as np - -from ...config import DetectionConfig -from ..video.frame import Frame -from .base import Detector -from .types import CensorType, Detection - - -class ClassicCVDetector(Detector): - def __init__( - self, - config: DetectionConfig | None = None, - types: "set[CensorType] | None" = None, - ) -> None: - self.cfg = config or DetectionConfig() - # Which censorship kinds to look for. Default: all. The composite detector - # restricts this to black_bar/blur (mosaic comes from the YOLO model). - self.types = ( - types if types is not None - else {CensorType.MOSAIC, CensorType.BLUR, CensorType.BLACK_BAR} - ) - - # ------------------------------------------------------------------ public - def detect(self, frame: Frame) -> list[Detection]: - bgr = frame.image - h0, w0 = bgr.shape[:2] - scale = self._proc_scale(w0, h0) - proc = ( - cv2.resize(bgr, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) - if scale != 1.0 - else bgr - ) - gray = cv2.cvtColor(proc, cv2.COLOR_BGR2GRAY) - ph, pw = gray.shape - min_area = self.cfg.min_area_frac * pw * ph - - dets: list[Detection] = [] - for ctype, fn, factor in ( - (CensorType.BLACK_BAR, self._detect_bars, 1.0), - (CensorType.MOSAIC, self._detect_mosaic, 4.0), - (CensorType.BLUR, self._detect_blur, 6.0), - ): - if ctype not in self.types: - continue - try: - dets += fn(proc, gray, min_area * factor) - except Exception: - # A failing heuristic must not break playback; skip it for this frame. - continue - - # Map proc-space coordinates back to source-frame pixels. - inv = 1.0 / scale - for d in dets: - x, y, w, h = d.bbox - d.bbox = (round(x * inv), round(y * inv), round(w * inv), round(h * inv)) - d.polygon = [(round(px * inv), round(py * inv)) for px, py in d.polygon] - return self._dedup(dets) - - # ----------------------------------------------------------------- helpers - def _proc_scale(self, w: int, h: int) -> float: - longest = max(w, h) - if longest <= self.cfg.proc_max_dim: - return 1.0 - return self.cfg.proc_max_dim / longest - - @staticmethod - def _local_std(g: np.ndarray, win: int) -> np.ndarray: - """Per-pixel standard deviation over a (win x win) box window.""" - mean = cv2.boxFilter(g, -1, (win, win)) - sqmean = cv2.boxFilter(g * g, -1, (win, win)) - var = np.maximum(sqmean - mean * mean, 0.0) - return np.sqrt(var) - - def _mask_to_detections( - self, - mask: np.ndarray, - ctype: CensorType, - min_area: float, - base_score: float, - min_extent: float = 0.0, - min_side: int = 0, - ) -> list[Detection]: - mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8)) - mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8)) - contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) - - out: list[Detection] = [] - for c in contours: - area = cv2.contourArea(c) - if area < min_area: - continue - x, y, w, h = cv2.boundingRect(c) - if min(w, h) < min_side: - continue # reject thin strips (e.g. edge false-positives) - extent = area / float(w * h + 1e-6) # how rectangular the blob is - if extent < min_extent: - continue - approx = cv2.approxPolyDP(c, 0.01 * cv2.arcLength(c, True), True) - poly = [(int(p[0][0]), int(p[0][1])) for p in approx] - score = float(np.clip(base_score + 0.25 * extent, 0.0, 1.0)) - out.append(Detection(type=ctype, score=score, bbox=(x, y, w, h), polygon=poly)) - return out - - # --------------------------------------------------------------- detectors - def _detect_bars(self, bgr, gray, min_area) -> list[Detection]: - # Solid censor bars are achromatic (black OR white) rectangles. Requiring - # low saturation + high rectangularity excludes large flat *colored* fills - # that are common in drawn/anime backgrounds. - hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) - sat, val = hsv[:, :, 1], hsv[:, :, 2] - achromatic = sat < self.cfg.bar_saturation_max - dark = (val < self.cfg.black_intensity) & achromatic - light = (val > self.cfg.white_intensity) & achromatic - mask = (dark | light).astype(np.uint8) * 255 - return self._mask_to_detections( - mask, CensorType.BLACK_BAR, min_area, base_score=0.55, - min_extent=self.cfg.bar_min_extent, - ) - - def _detect_mosaic(self, bgr, gray, min_area) -> list[Detection]: - g = gray.astype(np.float32) - h, w = gray.shape - win = 17 - # Lowest reconstruction residual across candidate tile sizes AND grid phases. - # Real mosaics aren't aligned to the origin, so we try a few offsets per size - # (phase-invariant) and keep the best fit. - best_residual = np.full((h, w), np.inf, np.float32) - for b in self.cfg.mosaic_block_sizes: - half = b // 2 - for oy, ox in ((0, 0), (half, 0), (0, half), (half, half)): - sub = g[oy:, ox:] - sh, sw = sub.shape - if sh < b or sw < b: - continue - small = cv2.resize(sub, (max(1, sw // b), max(1, sh // b)), interpolation=cv2.INTER_AREA) - restored = cv2.resize(small, (sw, sh), interpolation=cv2.INTER_NEAREST) - region = best_residual[oy:oy + sh, ox:ox + sw] - np.minimum(region, np.abs(sub - restored), out=region) - best_residual = cv2.boxFilter(best_residual, -1, (win, win)) - - contrast = self._local_std(g, win) - # Mosaic has edges in BOTH directions; a lone straight boundary (flat-region - # border, bar edge) has edge energy in only one — exclude those. - gx = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)), -1, (win, win)) - gy = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)), -1, (win, win)) - both_dirs = (gx > self.cfg.mosaic_grad_min) & (gy > self.cfg.mosaic_grad_min) - - blocky = best_residual < self.cfg.mosaic_residual_max - textured = contrast > self.cfg.mosaic_contrast_min - mask = (blocky & textured & both_dirs).astype(np.uint8) * 255 - return self._mask_to_detections( - mask, CensorType.MOSAIC, min_area, base_score=0.50, min_side=self.cfg.mosaic_min_side - ) - - def _detect_blur(self, bgr, gray, min_area) -> list[Detection]: - g = gray.astype(np.float32) - win = self.cfg.blur_window | 1 # force odd - lap = cv2.Laplacian(g, cv2.CV_32F, ksize=3) - sharpness = cv2.boxFilter(lap * lap, -1, (win, win)) # local high-freq energy - median = float(np.median(sharpness)) + 1e-6 - contrast = self._local_std(g, win) - - blurry = sharpness < median * self.cfg.blur_sharpness_ratio - textured = contrast > self.cfg.blur_contrast_min - mask = (blurry & textured).astype(np.uint8) * 255 - return self._mask_to_detections( - mask, CensorType.BLUR, min_area, base_score=0.40, min_side=self.cfg.mosaic_min_side - ) - - # ----------------------------------------------------------------- dedup - def _dedup(self, dets: list[Detection]) -> list[Detection]: - """Greedy IoU suppression; prefer black_bar > mosaic > blur, then score.""" - priority = { - CensorType.BLACK_BAR: 3, - CensorType.MOSAIC: 2, - CensorType.BLUR: 1, - CensorType.UNKNOWN: 0, - } - dets = sorted(dets, key=lambda d: (priority[d.type], d.score), reverse=True) - kept: list[Detection] = [] - for d in dets: - if all(self._iou(d.bbox, k.bbox) < 0.5 for k in kept): - kept.append(d) - return kept - - @staticmethod - def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float: - ax, ay, aw, ah = a - bx, by, bw, bh = b - ix = max(ax, bx) - iy = max(ay, by) - ix2 = min(ax + aw, bx + bw) - iy2 = min(ay + ah, by + bh) - iw, ih = max(0, ix2 - ix), max(0, iy2 - iy) - inter = iw * ih - union = aw * ah + bw * bh - inter - return inter / union if union > 0 else 0.0 diff --git a/hvideotool/core/detection/composite.py b/hvideotool/core/detection/composite.py deleted file mode 100644 index 25b4924..0000000 --- a/hvideotool/core/detection/composite.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Composite detector: runs several detectors and merges their results. - -Used for the "combined" mode = YOLO (mosaic) + classic-CV (black bars / blur). -Detections from all sub-detectors are concatenated, then de-duplicated by IoU -(higher score wins) so overlapping hits from different detectors don't stack. -""" - -from __future__ import annotations - -from ..video.frame import Frame -from .base import Detector -from .types import Detection - - -class CompositeDetector(Detector): - def __init__(self, detectors: list[Detector], iou_threshold: float = 0.6) -> None: - if not detectors: - raise ValueError("CompositeDetector requires at least one detector") - self._detectors = detectors - self._iou = iou_threshold - - @property - def name(self) -> str: - return "Composite(" + " + ".join(d.name for d in self._detectors) + ")" - - def detect(self, frame: Frame) -> list[Detection]: - merged: list[Detection] = [] - for detector in self._detectors: - try: - merged += detector.detect(frame) - except Exception: # noqa: BLE001 - one detector failing must not kill the frame - continue - return self._dedup(merged) - - def _dedup(self, dets: list[Detection]) -> list[Detection]: - dets = sorted(dets, key=lambda d: d.score, reverse=True) - kept: list[Detection] = [] - for d in dets: - if all(self._iou_of(d.bbox, k.bbox) < self._iou for k in kept): - kept.append(d) - return kept - - @staticmethod - def _iou_of(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float: - ax, ay, aw, ah = a - bx, by, bw, bh = b - ix, iy = max(ax, bx), max(ay, by) - ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh) - inter = max(0, ix2 - ix) * max(0, iy2 - iy) - union = aw * ah + bw * bh - inter - return inter / union if union > 0 else 0.0 diff --git a/hvideotool/core/detection/factory.py b/hvideotool/core/detection/factory.py index fc6aaf2..669299b 100644 --- a/hvideotool/core/detection/factory.py +++ b/hvideotool/core/detection/factory.py @@ -3,14 +3,15 @@ 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. """ from __future__ import annotations from ...config import AppConfig from .base import Detector -from .classic_cv import ClassicCVDetector -from .types import CensorType def _require_model(config: AppConfig) -> str: @@ -23,19 +24,6 @@ def _require_model(config: AppConfig) -> str: def build_detector(config: AppConfig) -> Detector: - if config.detector == "classic": - return ClassicCVDetector(config.detection) - if config.detector == "yolo": - from .yolo import YoloDetector # lazy: pulls torch/ultralytics + from .yolo import YoloDetector # lazy: pulls torch/ultralytics - return YoloDetector(_require_model(config), config.detection) - if config.detector == "combined": - # YOLO handles mosaic; classic-CV handles black bars / blur. - from .composite import CompositeDetector - from .yolo import YoloDetector - - return CompositeDetector([ - YoloDetector(_require_model(config), config.detection), - ClassicCVDetector(config.detection, types={CensorType.BLACK_BAR, CensorType.BLUR}), - ]) - raise ValueError(f"Неизвестный детектор: {config.detector!r}") + return YoloDetector(_require_model(config), config.detection) diff --git a/hvideotool/core/detection/yolo.py b/hvideotool/core/detection/yolo.py index ad6e895..7b7f4c2 100644 --- a/hvideotool/core/detection/yolo.py +++ b/hvideotool/core/detection/yolo.py @@ -7,7 +7,7 @@ a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose cl names map onto :class:`CensorType`. Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the -rest of the app — and the classic-CV detector — never pull them in. +rest of the app never pulls them in until detection actually runs. Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README. """ diff --git a/hvideotool/core/project.py b/hvideotool/core/project.py index e7ee9f5..0fd6b34 100644 --- a/hvideotool/core/project.py +++ b/hvideotool/core/project.py @@ -30,6 +30,7 @@ FRAMES_DIR = "frames" CACHE_FILE = "detections.json" COLLECTIONS_DIR = "collections" FAVORITES_DIR = "Избранное" # the single default collection ("в избранное") +RESTORED_DIR = "restored" # batch "расцензурить все" output (kept out of frames/) _VERSION = 1 # The subset of AppConfig fields a project remembers (mirrored to/from project.json). @@ -78,6 +79,12 @@ class Project: """The single default collection — frames moved "to favorites" land here.""" return self.collections_dir / FAVORITES_DIR + @property + def restored_dir(self) -> Path: + """Batch restoration output ("расцензурить все") — mirrors frame basenames. + Kept out of ``frames/`` so results aren't listed/re-detected/re-restored.""" + return self.root / RESTORED_DIR + # ------------------------------------------------------------- lifecycle @classmethod def create( diff --git a/hvideotool/core/restore/base.py b/hvideotool/core/restore/base.py index e1ed427..ee4078e 100644 --- a/hvideotool/core/restore/base.py +++ b/hvideotool/core/restore/base.py @@ -1,9 +1,10 @@ """Restorer interface — "un-censor" detected regions of an image. A Restorer takes an image plus the detected censored regions and returns a new -image with those regions reconstructed/filled. This mirrors the ``Detector`` -abstraction so different engines (classic inpaint now; a generative model like -DeepMosaics / LADA later) plug in behind the same interface. +image with those regions reconstructed. This mirrors the ``Detector`` abstraction +so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug +in behind the same interface. Note DeepMosaics locates the mosaic itself, so the +``detections`` argument is currently advisory (unused by the DeepMosaics engines). """ from __future__ import annotations @@ -24,7 +25,18 @@ class Cancelled(Exception): """Raised by a Restorer when ``should_cancel`` asked it to stop.""" +# Sequence (batch) restore callbacks — see ``Restorer.restore_sequence``. +FrameGetter = Callable[[int], np.ndarray] # index -> BGR image +DetGetter = Callable[[int], list[Detection]] # index -> that frame's detections +ResultSink = Callable[[int, np.ndarray], None] # (index, restored image) -> None + + class Restorer(ABC): + #: Whether this engine uses *neighbouring* frames (so a batch run must feed it a + #: contiguous, ordered sequence — see :meth:`restore_sequence`). Per-frame engines + #: leave this False; the temporal DeepMosaics (BVDNet) sets it True. + temporal: bool = False + @property def name(self) -> str: return type(self).__name__ @@ -42,3 +54,24 @@ class Restorer(ABC): True the engine should abort and raise :class:`Cancelled`. """ raise NotImplementedError + + def restore_sequence( + self, + count: int, + get_frame: FrameGetter, + get_dets: DetGetter, + emit: ResultSink, + should_cancel: CancelCheck | None = None, + ) -> None: + """Restore ``count`` frames *in order*, calling ``emit(i, restored)`` for each. + + The default treats every frame independently (just loops :meth:`restore`). + Temporal engines override this to pull neighbouring frames via ``get_frame`` + and carry recurrent state across the sequence. ``get_frame``/``get_dets`` are + lazy so the engine only reads the frames it needs; ``emit`` lets the caller + stream results to disk instead of holding them all in memory. + """ + for i in range(count): + if should_cancel is not None and should_cancel(): + raise Cancelled("Восстановление отменено") + emit(i, self.restore(get_frame(i), get_dets(i), should_cancel)) diff --git a/hvideotool/core/restore/deepmosaics.py b/hvideotool/core/restore/deepmosaics.py index 3f7b2ec..991d13d 100644 --- a/hvideotool/core/restore/deepmosaics.py +++ b/hvideotool/core/restore/deepmosaics.py @@ -26,18 +26,29 @@ from types import SimpleNamespace import numpy as np from ..detection.types import Detection -from .base import CancelCheck, Cancelled, Restorer +from .base import ( + CancelCheck, + Cancelled, + DetGetter, + FrameGetter, + ResultSink, + Restorer, +) _VENDOR = Path(__file__).parent / "_deepmosaics" # Default place to drop DeepMosaics clean weights (gitignored — see models/). DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics" -def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]: - """Find usable per-frame clean models: (display_name, full_path). +def discover_models( + extra_dir: str | None = None, include_video: bool = False +) -> list[tuple[str, str]]: + """Find usable clean models: (display_name, full_path). Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given) - for ``clean_*.pth``. The video model is skipped — it can't run per-frame. + for ``clean_*.pth``. By default the video model is skipped — it can't run + per-frame; pass ``include_video=True`` for the temporal (BVDNet) engine, which + *needs* ``clean_*_video.pth``. """ dirs = [DEFAULT_WEIGHTS_DIR] if extra_dir: @@ -48,13 +59,23 @@ def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]: if not d.is_dir(): continue for p in sorted(d.glob("clean_*.pth")): - if "video" in p.name.lower() or p.name in seen: + if p.name in seen: + continue + if "video" in p.name.lower() and not include_video: continue seen.add(p.name) out.append((p.stem, str(p))) return out +def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None: + """Locate ``mosaic_position.pth`` (the BiSeNet mosaic locator) for ``model``.""" + candidates = [model.parent / "mosaic_position.pth"] + if dm_dir: + candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth") + return next((p for p in candidates if p.is_file()), None) + + def _netg_kind(model_name: str) -> str: """Pick DeepMosaics' netG type from the weights filename (see their options.py).""" n = model_name.lower() @@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer): ) model = Path(model_path) self._netg = _netg_kind(model.name) # raises on a video model - pos = self._find_mosaic_position(model, deepmosaics_dir) + pos = _find_mosaic_position(model, deepmosaics_dir) if pos is None: raise ValueError( "Рядом с clean-моделью не найден mosaic_position.pth.\n" @@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer): self._gpu = gpu_id self._loaded = False # models loaded lazily on first restore - @staticmethod - def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None: - candidates = [model.parent / "mosaic_position.pth"] - if dm_dir: - candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth") - return next((p for p in candidates if p.is_file()), None) - @property def name(self) -> str: return f"DeepMosaics(gpu={self._gpu})" @@ -167,3 +181,147 @@ class DeepMosaicsRestorer(Restorer): img_mosaic = work[y - size:y + size, x - size:x + size] img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt) return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather) + + +class DeepMosaicsVideoRestorer(Restorer): + """Temporal DeepMosaics (BVDNet) — un-censors using *neighbouring* frames. + + Reproduces DeepMosaics' ``cleanmosaic_video_fusion`` per target frame: for frame + ``i`` it feeds the network a temporal window of ``T`` frames (sampled at step ``S`` + around ``i``) plus its own previous output (recurrent), so the reconstruction is + temporally coherent. Because of that recurrence the frames MUST be processed in + order over a contiguous range — see :meth:`restore_sequence` (the batch run). + + Needs the **video** weights ``clean_youknow_video.pth`` + ``mosaic_position.pth`` + (beside it). Single-frame :meth:`restore` degrades to a window of the same frame. + """ + + temporal = True + + # DeepMosaics fusion window: N before/after at step S → T = 2N+1 frames, INPUT_SIZE px. + _N, _T, _S = 2, 5, 3 + _INPUT_SIZE = 256 + + def __init__( + self, + deepmosaics_dir: str | None, + model_path: str | None, + python_exe: str | None = None, # unused (in-process); kept for factory parity + gpu_id: str = "0", + ) -> None: + chosen: Path | None = None + if model_path and Path(model_path).is_file() and "video" in Path(model_path).name.lower(): + chosen = Path(model_path) + else: # configured model missing or not a video model → auto-pick a video model + vids = [p for _n, p in discover_models(include_video=True) if "video" in Path(p).name.lower()] + if vids: + chosen = Path(vids[0]) + if chosen is None: + raise ValueError( + "Не найдены веса видеомодели DeepMosaics (clean_*_video.pth).\n" + "Положите clean_youknow_video.pth + mosaic_position.pth в models/deepmosaics " + "(или выберите в «Восстановление…»). См. README." + ) + pos = _find_mosaic_position(chosen, deepmosaics_dir) + if pos is None: + raise ValueError( + "Рядом с видеомоделью не найден mosaic_position.pth.\n" + "Положите mosaic_position.pth в ту же папку, что и clean_*_video.pth. См. README." + ) + self._model = str(chosen) + self._pos = str(pos) + self._gpu = gpu_id + self._loaded = False + + @property + def name(self) -> str: + return f"DeepMosaicsVideo(gpu={self._gpu})" + + # ------------------------------------------------------------------ engine + def _ensure_loaded(self) -> None: + if self._loaded: + return + if str(_VENDOR) not in sys.path: + sys.path.insert(0, str(_VENDOR)) + + import torch # noqa: E402 + if self._gpu != "-1" and not torch.cuda.is_available(): + self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why) + + from models import loadmodel, runmodel # type: ignore # noqa: E402 + import util.data as data # type: ignore # noqa: E402 + import util.image_processing as impro # type: ignore # noqa: E402 + + self._torch = torch + self._runmodel = runmodel + self._data = data + self._impro = impro + self._opt = SimpleNamespace( + gpu_id=self._gpu, + model_path=self._model, + mosaic_position_model_path=self._pos, + mask_threshold=64, + all_mosaic_area=False, + ex_mult=1.5, + no_feather=False, + ) + self._netM = loadmodel.bisenet(self._opt, "mosaic") + self._netG = loadmodel.video(self._opt) # BVDNet + self._loaded = True + + def restore( + self, + image: np.ndarray, + detections: list[Detection], + should_cancel: CancelCheck | None = None, + ) -> np.ndarray: + """Single-frame restore — no neighbours, so the window is the same frame.""" + out: dict[int, np.ndarray] = {} + self.restore_sequence( + 1, + lambda _i: image, + lambda _i: detections, + lambda i, r: out.__setitem__(i, r), + should_cancel, + ) + return out.get(0, image.copy()) + + def restore_sequence( + self, + count: int, + get_frame: FrameGetter, + get_dets: DetGetter, + emit: ResultSink, + should_cancel: CancelCheck | None = None, + ) -> None: + self._ensure_loaded() + torch, data, impro, opt = self._torch, self._data, self._impro, self._opt + N, T, S, SZ = self._N, self._T, self._S, self._INPUT_SIZE + + previous = None # recurrent state: the network's previous output (a tensor) + for i in range(count): + if should_cancel is not None and should_cancel(): + raise Cancelled("Восстановление отменено") + img_origin = get_frame(i) + x, y, size, mask = self._runmodel.get_mosaic_position(img_origin, self._netM, opt) + if size <= 50: + emit(i, img_origin.copy()) # no mosaic here; recurrence carries over + continue + + stream = [] + for k in range(T): + j = min(max(i + (k - N) * S, 0), count - 1) # clamp window to range edges + frame = img_origin if j == i else get_frame(j) + crop = frame[y - size:y + size, x - size:x + size] + stream.append(impro.resize(crop, SZ)[:, :, ::-1]) # BGR→RGB, SZ×SZ + + if previous is None: # seed recurrence with the (centre) input crop + previous = data.im2tensor(stream[N], bgr2rgb=False, gpu_id=opt.gpu_id) + + arr = np.array(stream).reshape(1, T, SZ, SZ, 3).transpose((0, 4, 1, 2, 3)) + tensor = data.to_tensor(data.normalize(arr), gpu_id=opt.gpu_id) + with torch.no_grad(): + pred = self._netG(tensor, previous) + previous = pred + img_fake = data.tensor2im(pred, rgb2bgr=True) + emit(i, impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather)) diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py index 07eba75..00e3e75 100644 --- a/hvideotool/core/restore/factory.py +++ b/hvideotool/core/restore/factory.py @@ -1,9 +1,13 @@ """Restorer factory: build a Restorer from the app config. -- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct). -- ``deepmosaics``: real generative mosaic removal. The DeepMosaics network code is - vendored (``_deepmosaics/``, GPL-3.0) and run in-process; the user supplies only the - clean weights (+ ``mosaic_position.pth`` alongside). A CUDA GPU is recommended. +Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but +did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``, +GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth`` +alongside). A CUDA GPU is recommended. + +- ``deepmosaics``: per-frame generative mosaic removal (image model). +- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for + coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence. """ from __future__ import annotations @@ -11,25 +15,28 @@ from __future__ import annotations from typing import TYPE_CHECKING from .base import Restorer -from .inpaint import InpaintRestorer if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed) from ...config import AppConfig -def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer: - if name == "inpaint": - return InpaintRestorer() +def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer: + if config is None: + raise ValueError("Для DeepMosaics нужны настройки (config).") if name == "deepmosaics": from .deepmosaics import DeepMosaicsRestorer - if config is None: - raise ValueError("Для DeepMosaics нужны настройки (config).") return DeepMosaicsRestorer( config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu ) + if name == "deepmosaics_video": + from .deepmosaics import DeepMosaicsVideoRestorer + + return DeepMosaicsVideoRestorer( + config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu + ) if name == "lada": raise ValueError( - "Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README." + "Движок LADA пока не подключён. Используйте DeepMosaics. См. README." ) raise ValueError(f"Неизвестный режим восстановления: {name!r}") diff --git a/hvideotool/core/restore/inpaint.py b/hvideotool/core/restore/inpaint.py deleted file mode 100644 index f5e2b16..0000000 --- a/hvideotool/core/restore/inpaint.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Classic inpainting restorer (cv2) — the always-available baseline. - -HONEST LIMITATION: cv2 inpainting fills the masked region by propagating -surrounding pixels. It removes the mosaic/bar but does NOT reconstruct the hidden -detail — it smooths/guesses. For real reconstruction a generative model -(DeepMosaics / LADA) is needed; this is the no-weights, no-GPU fallback so the -"Расцензурить кадр" flow works end-to-end today. -""" - -from __future__ import annotations - -import cv2 -import numpy as np - -from ..detection.types import Detection -from .base import CancelCheck, Restorer -from .mask import detections_to_mask - - -class InpaintRestorer(Restorer): - def __init__(self, radius: int = 3, dilate: int = 2, method: str = "telea") -> None: - self.radius = radius - self.dilate = dilate - self.method = method - - @property - def name(self) -> str: - return f"InpaintRestorer({self.method})" - - def restore( - self, - image: np.ndarray, - detections: list[Detection], - should_cancel: CancelCheck | None = None, - ) -> np.ndarray: - # Single cv2.inpaint call — effectively instant, so cancellation is moot. - if not detections: - return image.copy() - mask = detections_to_mask(image.shape, detections, dilate=self.dilate) - flags = cv2.INPAINT_TELEA if self.method == "telea" else cv2.INPAINT_NS - return cv2.inpaint(image, mask, self.radius, flags) diff --git a/hvideotool/core/restore/mask.py b/hvideotool/core/restore/mask.py deleted file mode 100644 index 237549b..0000000 --- a/hvideotool/core/restore/mask.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Build a binary mask of the censored regions from detections.""" - -from __future__ import annotations - -import cv2 -import numpy as np - -from ..detection.types import Detection - - -def detections_to_mask( - shape: tuple[int, int], detections: list[Detection], dilate: int = 0 -) -> np.ndarray: - """White (255) over every detected region (polygon if present, else bbox).""" - h, w = shape[:2] - mask = np.zeros((h, w), np.uint8) - for d in detections: - if len(d.polygon) >= 3: - cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255) - else: - x, y, bw, bh = d.bbox - cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1) - if dilate > 0: - k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8) - mask = cv2.dilate(mask, k) - return mask diff --git a/hvideotool/core/torch_info.py b/hvideotool/core/torch_info.py index a45c33a..eb16221 100644 --- a/hvideotool/core/torch_info.py +++ b/hvideotool/core/torch_info.py @@ -1,18 +1,66 @@ -"""Probe the PyTorch / CUDA situation, so the UI can show a device badge. +"""Diagnose the PyTorch / CUDA situation so the UI can explain *why* it's on CPU. -Pure (no Qt). ``gather()`` imports torch (slow / heavy) — call it off the GUI -thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and -how to enable the GPU. +Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to +``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on +the gathered dict and figures out the most likely cause + concrete fix. """ from __future__ import annotations -# pip index for the CUDA build (matches the README). -CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121" +import re +import shutil +import subprocess +import sys + +# pip indexes for the CUDA builds (see README). cu121 needs a driver with CUDA >= 12.1. +CUDA_WHEELS = { + "cu121": "https://download.pytorch.org/whl/cu121", + "cu118": "https://download.pytorch.org/whl/cu118", +} + +_NO_WINDOW = 0x08000000 if sys.platform == "win32" else 0 # CREATE_NO_WINDOW + + +def _run_nvidia_smi() -> dict: + """Probe the NVIDIA driver/GPU via nvidia-smi. Never raises.""" + out: dict = {"found": False, "gpus": [], "driver_version": None, "cuda_driver": None} + exe = shutil.which("nvidia-smi") + if not exe and sys.platform == "win32": + candidate = r"C:\Windows\System32\nvidia-smi.exe" + exe = candidate if shutil.os.path.isfile(candidate) else None + if not exe: + return out + # GPU names + driver version (robust CSV form). + try: + r = subprocess.run( + [exe, "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW, + ) + if r.returncode == 0: + out["found"] = True + for line in r.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if parts and parts[0]: + out["gpus"].append(parts[0]) + if len(parts) > 1 and parts[1]: + out["driver_version"] = parts[1] + except Exception: # noqa: BLE001 - any failure => "not found" + return out + # Max CUDA version the driver supports (only in the plain header). + try: + r2 = subprocess.run( + [exe], capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW, + ) + m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout) + if m: + out["cuda_driver"] = m.group(1) + except Exception: # noqa: BLE001 + pass + return out def gather() -> dict: - """Collect torch/CUDA facts. Never raises — missing torch is a valid result.""" + """Collect torch + NVIDIA facts. Never raises — missing torch/GPU are valid.""" info: dict = { "installed": False, "version": None, # torch.__version__ (e.g. "2.12.0+cpu") @@ -20,27 +68,38 @@ def gather() -> dict: "cuda_available": False, "device_name": None, # the active GPU's name, if any "import_error": None, + # filled by nvidia-smi: + "nvidia_smi": False, + "gpus": [], + "driver_version": None, + "cuda_driver": None, } try: import torch - except Exception as exc: # noqa: BLE001 - report any import failure, not just ImportError + except Exception as exc: # noqa: BLE001 - report any import failure info["import_error"] = str(exc) - return info - info["installed"] = True - info["version"] = getattr(torch, "__version__", None) - try: - info["built_cuda"] = torch.version.cuda - except Exception: # noqa: BLE001 - info["built_cuda"] = None - try: - info["cuda_available"] = bool(torch.cuda.is_available()) - except Exception: # noqa: BLE001 - info["cuda_available"] = False - if info["cuda_available"]: + else: + info["installed"] = True + info["version"] = getattr(torch, "__version__", None) try: - info["device_name"] = torch.cuda.get_device_name(0) + info["built_cuda"] = torch.version.cuda except Exception: # noqa: BLE001 - info["device_name"] = None + info["built_cuda"] = None + try: + info["cuda_available"] = bool(torch.cuda.is_available()) + except Exception: # noqa: BLE001 + info["cuda_available"] = False + if info["cuda_available"]: + try: + info["device_name"] = torch.cuda.get_device_name(0) + except Exception: # noqa: BLE001 + info["device_name"] = None + + smi = _run_nvidia_smi() + info["nvidia_smi"] = smi["found"] + info["gpus"] = smi["gpus"] + info["driver_version"] = smi["driver_version"] + info["cuda_driver"] = smi["cuda_driver"] return info @@ -48,33 +107,117 @@ def device_label(info: dict) -> str: return "CUDA" if info.get("cuda_available") else "CPU" -def reason(info: dict) -> str: - """One-sentence human explanation of the current device choice.""" - if not info.get("installed"): - return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics " - "работают на CPU (классический детектор torch не требует).") - if info.get("cuda_available"): - name = info.get("device_name") or "GPU" - return f"PyTorch использует CUDA: {name}. Вычисления идут на видеокарте." - version = info.get("version") or "?" - built = info.get("built_cuda") - if not built: - return (f"Установлена CPU-сборка PyTorch ({version}) — без поддержки CUDA, " - "поэтому вычисления идут на процессоре (медленно).") - return (f"PyTorch собран с CUDA {built} ({version}), но GPU недоступен: нет " - "NVIDIA-видеокарты, не установлен/устарел драйвер, либо версия CUDA " - "несовместима с драйвером.") +def _ver_tuple(v: str | None) -> tuple[int, ...]: + try: + return tuple(int(x) for x in str(v).split(".")[:2]) + except (ValueError, AttributeError): + return () -def install_hint() -> str: - """Steps to enable the GPU (shown when running on CPU).""" +def recommend_channel(info: dict) -> str: + """Pick the pip CUDA wheel index that matches the driver (cu118 for older).""" + cd = _ver_tuple(info.get("cuda_driver")) + if cd and cd < (12, 1): + return "cu118" + return "cu121" + + +def install_command(channel: str = "cu121") -> str: + """The pip commands to (re)install the chosen CUDA build. + + Targets the **running interpreter** (``sys.executable -m pip``) so the command + hits the same venv that runs the app — not whatever ``pip`` is on PATH. (A common + trap: running bare ``pip`` in a global shell while torch lives in the project venv.) + """ + index = CUDA_WHEELS.get(channel, CUDA_WHEELS["cu121"]) + py = sys.executable or "python" + q = f'"{py}"' if " " in py else py return ( - "Как включить GPU (NVIDIA):\n" - "1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n" - "2. Переустановите PyTorch со сборкой CUDA:\n\n" - " pip uninstall -y torch torchvision\n" - f" pip install torch torchvision --index-url {CUDA_WHEEL_INDEX}\n\n" - "3. Перезапустите приложение.\n\n" - "Классический детектор работает и без CUDA. На CPU детекция и расцензуривание " - "просто медленнее." + f"{q} -m pip uninstall -y torch torchvision\n" + f"{q} -m pip install torch torchvision --index-url {index}" ) + + +def analyze(info: dict) -> dict: + """Turn the raw facts into {summary, details[list], steps, command}.""" + installed = info.get("installed") + cuda = info.get("cuda_available") + built = info.get("built_cuda") + gpus = info.get("gpus") or [] + smi = info.get("nvidia_smi") + driver = info.get("driver_version") + cuda_driver = info.get("cuda_driver") + channel = recommend_channel(info) + command = install_command(channel) + + if not installed: + build_str = "не установлен" + elif built: + build_str = f"CUDA {built}" + else: + build_str = "CPU-only (+cpu)" + gpu_str = ( + ", ".join(gpus) if gpus + else ("не обнаружена" if smi or info.get("nvidia_smi") is False else "nvidia-smi не найден") + ) + if not gpus and not smi: + gpu_str = "nvidia-smi не найден (нет драйвера NVIDIA?)" + details = [ + f"PyTorch: {info.get('version') or 'не установлен'}", + f"Сборка PyTorch: {build_str}", + f"CUDA доступна в PyTorch: {'да' if cuda else 'нет'}", + f"Видеокарта (nvidia-smi): {gpu_str}", + f"Драйвер NVIDIA: {driver or '—'}", + f"Макс. CUDA драйвера: {cuda_driver or '—'}", + f"Интерпретатор (venv): {sys.executable}", + ] + if info.get("import_error"): + details.append(f"Ошибка импорта torch: {info['import_error']}") + + if not installed: + summary = "PyTorch не установлен — детекция YOLO и DeepMosaics идут на CPU." + steps = ("Установите PyTorch (CUDA-сборку, если есть NVIDIA-видеокарта):\n\n" + + command + "\n\nДля YOLO также: pip install -e \".[yolo]\"") + elif cuda: + gpu = info.get("device_name") or (gpus[0] if gpus else "GPU") + summary = f"Всё в порядке: PyTorch использует CUDA. Активный GPU: {gpu}." + steps = "GPU уже задействован — ничего делать не нужно." + elif not built: # CPU-only torch build — the usual case + if gpus: + summary = ( + "Главная причина: установлена CPU-сборка PyTorch (+cpu) — она физически " + f"не умеет в CUDA. Видеокарта ({gpus[0]}) и драйвер {driver or '?'} на месте, " + "поэтому достаточно переустановить PyTorch со сборкой CUDA." + ) + steps = ("Переустановите PyTorch под CUDA, затем перезапустите приложение:\n\n" + + command) + else: + summary = ( + "Установлена CPU-сборка PyTorch (+cpu), и видеокарта NVIDIA не обнаружена " + "(nvidia-smi не отвечает). Либо нет NVIDIA GPU, либо не установлен драйвер." + ) + steps = ("1. Проверьте видеокарту и драйвер: в консоли выполните nvidia-smi\n" + "2. Если NVIDIA GPU есть — переустановите PyTorch под CUDA:\n\n" + + command + + "\n\nБез NVIDIA GPU всё работает на CPU — просто медленнее.") + else: # CUDA-enabled torch build, but CUDA still not available + if not gpus: + summary = ( + f"PyTorch собран под CUDA {built}, но видеокарта/драйвер NVIDIA не найдены. " + "Скорее всего не установлен драйвер NVIDIA или нет GPU." + ) + steps = "Установите свежий драйвер NVIDIA и перезапустите. Проверка: nvidia-smi." + elif cuda_driver and _ver_tuple(cuda_driver) < _ver_tuple(built): + summary = ( + f"Драйвер поддерживает CUDA {cuda_driver}, а PyTorch собран под CUDA {built} — " + "версия драйвера слишком старая." + ) + steps = ("Вариант A — обновите драйвер NVIDIA (рекомендуется).\n" + "Вариант B — поставьте PyTorch под CUDA вашего драйвера:\n\n" + command) + else: + summary = ( + f"PyTorch собран под CUDA {built}, GPU ({gpus[0]}) есть, но CUDA недоступна — " + "вероятен конфликт версий или повреждённая установка." + ) + steps = "Переустановите PyTorch под CUDA:\n\n" + command + return {"summary": summary, "details": details, "steps": steps, "command": command} diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index e0d43f4..4fca337 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -26,11 +26,10 @@ import shutil from pathlib import Path from PySide6.QtCore import Qt, QThreadPool -from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut +from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcut from PySide6.QtWidgets import ( QAbstractItemView, QApplication, - QComboBox, QDialog, QDoubleSpinBox, QFileDialog, @@ -41,6 +40,7 @@ from PySide6.QtWidgets import ( QListWidgetItem, QMainWindow, QMessageBox, + QPlainTextEdit, QProgressBar, QPushButton, QSplitter, @@ -51,7 +51,7 @@ from PySide6.QtWidgets import ( ) from .. import settings_store -from ..config import AppConfig +from ..config import AppConfig, normalize_config from ..core.detection import cache as detection_cache from ..core.detection.factory import build_detector from ..core.detection.types import Detection @@ -68,7 +68,6 @@ from .workers import Job _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} _VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)" -_DETECTORS = ["classic", "yolo", "combined"] class MainWindow(QMainWindow): @@ -120,6 +119,8 @@ class MainWindow(QMainWindow): file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True)) file_menu.addSeparator() file_menu.addAction("Движок восстановления…", self._open_restore_settings) + file_menu.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False)) + file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True)) file_menu.addSeparator() file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M") file_menu.addSeparator() @@ -136,14 +137,9 @@ class MainWindow(QMainWindow): tb.addAction(from_video) tb.addSeparator() - tb.addWidget(QLabel(" Детектор: ")) - self.detector_combo = QComboBox() - self.detector_combo.addItems(_DETECTORS) - self.detector_combo.setCurrentText(self._cfg.detector) - self.detector_combo.currentTextChanged.connect(self._on_detector_changed) - tb.addWidget(self.detector_combo) - + 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.addSeparator() @@ -166,6 +162,12 @@ class MainWindow(QMainWindow): restore = QAction("Расцензурить кадр", self, triggered=self._restore_current) restore.setToolTip("Восстановить найденные области на текущем кадре") tb.addAction(restore) + restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False)) + restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)") + tb.addAction(restore_all) + restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True)) + restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)") + tb.addAction(restore_regen) self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored) self.toggle_restored_action.setEnabled(False) tb.addAction(self.toggle_restored_action) @@ -359,27 +361,58 @@ class MainWindow(QMainWindow): info = self._device_info if self._device_info else torch_info.gather() cuda = bool(info.get("cuda_available")) + a = torch_info.analyze(info) + self._install_command = a["command"] # what the Copy button will copy + lines = [ - torch_info.reason(info), + "ВЕРДИКТ:", + a["summary"], "", "Диагностика:", - f" • PyTorch: {info.get('version') or 'не установлен'}", - f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}", - f" • CUDA доступна: {'да' if cuda else 'нет'}", + *(f" • {d}" for d in a["details"]), + "", + "Что делать:", + a["steps"], ] - if info.get("device_name"): - lines.append(f" • GPU: {info['device_name']}") - if info.get("import_error"): - lines.append(f" • Ошибка импорта torch: {info['import_error']}") - box = QMessageBox(self) - box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning) - box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU")) - box.setText("\n".join(lines)) + # A real dialog (not QMessageBox) so the text — incl. the install command — is + # selectable, and a Copy button drops the pip command straight onto the clipboard. + dlg = QDialog(self) + dlg.setWindowTitle("Почему " + ("GPU" if cuda else "CPU") + " — диагностика PyTorch/CUDA") + dlg.resize(620, 480) + layout = QVBoxLayout(dlg) + + text = QPlainTextEdit() + text.setReadOnly(True) + text.setPlainText("\n".join(lines)) + mono = QFont("Consolas") + mono.setStyleHint(QFont.Monospace) + text.setFont(mono) + layout.addWidget(text, 1) + + buttons = QHBoxLayout() if not cuda: - box.setInformativeText(torch_info.install_hint()) - box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands - box.exec() + copy_btn = QPushButton("Скопировать команду установки") + copy_btn.clicked.connect(self._copy_install_command) + buttons.addWidget(copy_btn) + recheck = QPushButton("Проверить заново") + recheck.setToolTip("Перепроверить torch/CUDA (например, после переустановки)") + recheck.clicked.connect(lambda: (self._probe_device(), dlg.accept())) + buttons.addWidget(recheck) + buttons.addStretch(1) + close_btn = QPushButton("Закрыть") + close_btn.clicked.connect(dlg.accept) + buttons.addWidget(close_btn) + layout.addLayout(buttons) + dlg.exec() + + def _copy_install_command(self) -> None: + command = getattr(self, "_install_command", None) + if not command: + from ..core import torch_info + command = torch_info.install_command() + QApplication.clipboard().setText(command) + self.statusBar().showMessage("Команда установки скопирована в буфер обмена") # ------------------------------------------------------------- cancellation def _begin_busy(self, total: int | None = None) -> None: @@ -388,7 +421,6 @@ 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.detector_combo.setEnabled(False) self.model_action.setEnabled(False) if total is None: self.progress.setRange(0, 0) # indeterminate @@ -400,7 +432,6 @@ class MainWindow(QMainWindow): def _end_busy(self) -> None: self._busy = False self.stop_action.setEnabled(False) - self.detector_combo.setEnabled(True) self.model_action.setEnabled(True) self.progress.setVisible(False) self.progress.setRange(0, 100) # leave it determinate for the next user @@ -459,19 +490,19 @@ class MainWindow(QMainWindow): self._detector_key = key return self._detector - def _on_detector_changed(self, name: str) -> None: - self._cfg.detector = name - # YOLO/combined need a model. Auto-pick a known one from models/ if we have it; - # only prompt when nothing suitable is found (don't nag when the path is obvious). - if name in ("yolo", "combined") and not self._cfg.model_path: - found = self._auto_find_model() - if found: - self._cfg.model_path = found - self.statusBar().showMessage(f"Модель найдена автоматически: {found}") - else: - self._choose_model() - self._persist_settings() - self._invalidate_results() + def _ensure_model(self) -> None: + """Make sure the YOLO detector has weights — auto-pick from ./models silently. + + 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. + """ + if self._cfg.model_path and Path(self._cfg.model_path).is_file(): + return + found = self._auto_find_model() + if found: + self._cfg.model_path = found + self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}") + self._persist_settings() @staticmethod def _auto_find_model() -> str | None: @@ -667,6 +698,8 @@ 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._sync_settings_ui() self._detector_key = None self._restorer_key = None @@ -679,9 +712,6 @@ class MainWindow(QMainWindow): def _sync_settings_ui(self) -> None: """Reflect the (project's) config onto the toolbar widgets without signal loops.""" - self.detector_combo.blockSignals(True) - self.detector_combo.setCurrentText(self._cfg.detector) - self.detector_combo.blockSignals(False) self.threshold_spin.blockSignals(True) self.threshold_spin.setValue(self._cfg.default_threshold) self.threshold_spin.blockSignals(False) @@ -912,41 +942,32 @@ class MainWindow(QMainWindow): key, dets = payload self._results[key] = dets self._tag_file(Path(key), len(dets)) + # If the frame being viewed was just computed, show its overlay live. + if not self._showing_restored and self._current is not None and str(self._current) == key: + self._show(self._current) self._tick_count += 1 if self._tick_count % 25 == 0: self._refresh_marks() # let marks appear progressively (throttled) # ------------------------------------------------------------- restoration def _restore_current(self) -> None: - """Restore the current frame's regions on a background thread, then show it. + """Restore the current frame on a background thread, then show it. - Detections are computed first (in the same job) if not cached. The DeepMosaics - engine polls ``job.cancelled`` so "■ Стоп" stops it promptly.""" + DeepMosaics locates the mosaic itself, so no detection step is needed — we just + run the engine on the frame (if there's no mosaic the frame comes back unchanged). + The engine polls ``job.cancelled`` so "■ Стоп" stops it promptly.""" if self._current is None or self._busy: return path = self._current key = str(path) def fn(job): - dets = self._results.get(key) - if dets is None: - dets = self._compute(self._make_detector(), path) - job.tick(("dets", key, dets)) # cache them on the GUI thread - if not dets: - return ("empty", key) img = imread_unicode(key) if img is None: raise RuntimeError(f"Не удалось прочитать: {path.name}") restorer = self._make_restorer() - restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled) - return ("restored", key, restored, len(dets), restorer.name) - - def tick(payload): - if payload[0] == "dets": - _, k, dets = payload - self._results[k] = dets - self._tag_file(Path(k), len(dets)) - self._refresh_marks() + restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled) + return ("restored", key, restored, restorer.name) def done(result, cancelled): if cancelled: @@ -954,19 +975,86 @@ class MainWindow(QMainWindow): return if result is None: return - if result[0] == "empty": - self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать") - return - _, k, restored, n, engine = result + _, k, restored, engine = result self._restored[k] = restored if self._current is not None and str(self._current) == k: self._showing_restored = True self.view.set_image(restored, []) self._update_restore_actions() - self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name} — {n} обл.") + self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}") self.statusBar().showMessage(f"Восстановление: {path.name}…") - self._start_job(fn, None, on_tick=tick, on_done=done) + self._start_job(fn, None, on_done=done) + + def _restore_all(self, force: bool = False) -> None: + """Restore every frame on a background thread, writing results to ``restored/``. + + DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame + engine skips frames already restored (resume) unless ``force``. The temporal + engine (DeepMosaics-video) runs the whole contiguous sequence in order via + ``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied. + """ + if not self._files or self._project is None or self._busy: + return + files = list(self._files) # snapshot — favorites/move mutate self._files + total = len(files) + out_dir = self._project.restored_dir + out_dir.mkdir(parents=True, exist_ok=True) + + def out_path(p: Path) -> Path: + return out_dir / f"{p.stem}.jpg" + + def fn(job): + restorer = self._make_restorer() # built on the worker (may raise) + frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads + + def get_frame(i): + img = frame_cache.get(i) + if img is None: + img = imread_unicode(str(files[i])) + if img is None: + raise RuntimeError(f"Не удалось прочитать: {files[i].name}") + if len(frame_cache) > 24: + frame_cache.clear() + frame_cache[i] = img + return img + + def emit(i, restored): + imwrite_unicode(str(out_path(files[i])), restored) + job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}") + + if restorer.temporal: + restorer.restore_sequence( + total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled + ) + else: + for i, p in enumerate(files): + if job.cancelled: + break + if not force and out_path(p).is_file(): + job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}") + continue + emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled)) + frame_cache.pop(i, None) # per-frame: don't accumulate + return None + + def done(_result, cancelled): + if self._current is not None: # live-preview the current frame's result, if any + rp = out_path(self._current) + if rp.is_file(): + img = imread_unicode(str(rp)) + if img is not None: + self._restored[str(self._current)] = img + self._showing_restored = True + self.view.set_image(img, []) + self._update_restore_actions() + self.statusBar().showMessage( + "Расцензуривание отменено" if cancelled + else f"Готово: результаты в {out_dir.name}/ ({total} кадров)" + ) + + self.statusBar().showMessage("Пакетное расцензуривание…") + self._start_job(fn, total, on_done=done) def _make_restorer(self): key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, diff --git a/hvideotool/ui/restore_dialog.py b/hvideotool/ui/restore_dialog.py index c046ab8..c6f83af 100644 --- a/hvideotool/ui/restore_dialog.py +++ b/hvideotool/ui/restore_dialog.py @@ -1,9 +1,11 @@ """Configure the restoration ("расцензурить") engine. -inpaint — no setup. deepmosaics — the network code is vendored (built-in); the user -picks a clean model from a dropdown of the bundled ``models/deepmosaics`` weights -(or browses to another ``clean_*.pth``). ``mosaic_position.pth`` must sit beside the -chosen model. A CUDA GPU is strongly recommended (GPU id, -1 = CPU/slow). +Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео") +engine. The network code is vendored (built-in); the user picks a clean model from a +dropdown of the bundled ``models/deepmosaics`` weights (or browses to another +``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered. +``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly +recommended (GPU id, -1 = CPU/slow). """ from __future__ import annotations @@ -35,9 +37,9 @@ class RestoreDialog(QDialog): self.setMinimumWidth(560) self.engine = QComboBox() - self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint") - self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics") - self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0) + self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics") + self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video") + self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer))) self.engine.currentIndexChanged.connect(self._sync) # Model dropdown — bundled clean models, plus the configured one if external. @@ -52,9 +54,11 @@ class RestoreDialog(QDialog): form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model)) form.addRow("GPU id:", self.dm_gpu) hint = QLabel( - "Модели берутся из models/deepmosaics. Нужны clean_youknow_resnet_9blocks.pth " - "и mosaic_position.pth (рядом). Видеомодель clean_*_video.pth покадрово не " - "работает и в списке не показывается. На CPU медленно — лучше GPU. См. README." + "Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n" + "• Картинка: clean_youknow_resnet_9blocks.pth — покадрово.\n" + "• Видео: clean_youknow_video.pth — использует соседние кадры (когерентнее на " + "роликах), требует прогона по диапазону («Расцензурить все»).\n" + "На CPU медленно — лучше GPU. См. README." ) hint.setWordWrap(True) form.addRow(hint) @@ -67,7 +71,12 @@ class RestoreDialog(QDialog): def _populate_models(self, current: str | None) -> None: self.model_combo.clear() - for name, path in discover_models(): + is_video = self.engine.currentData() == "deepmosaics_video" + if is_video: # temporal engine: only the video weights (clean_*_video.pth) + models = [(n, p) for n, p in discover_models(include_video=True) if "video" in n.lower()] + else: + models = discover_models() # per-frame engine: image clean models only + for name, path in models: self.model_combo.addItem(name, path) # Keep an externally-configured model selectable even if it's outside the folder. if current and self.model_combo.findData(current) < 0: @@ -88,7 +97,9 @@ class RestoreDialog(QDialog): return w def _sync(self) -> None: - is_dm = self.engine.currentData() == "deepmosaics" + is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video") + # The model list differs per engine (image vs video weights) — repopulate. + self._populate_models(self._cfg.dm_model) self.model_combo.setEnabled(is_dm) self.dm_gpu.setEnabled(is_dm)