Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.

This commit is contained in:
Leonid Pershin
2026-06-07 06:16:05 +03:00
parent cc518cc3e6
commit 9c471ca701
17 changed files with 798 additions and 676 deletions
+101 -81
View File
@@ -47,12 +47,15 @@ what it found.
Keep this scope sharp: Keep this scope sharp:
- Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить") - Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить")
step was added later (user-requested): per-frame, on-demand, behind a `Restorer` step was added later (user-requested): on-demand (single frame **or** whole-project
interface. Two engines: a cv2 **inpaint baseline** (fills, does NOT reconstruct) and batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
**DeepMosaics**real generative mosaic removal, its GPL-3.0 network code **vendored** restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
under `core/restore/_deepmosaics/` and run in-process (user supplies only the weights). composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were
Because of that vendoring the **whole project is GPL-3.0**. LADA (BasicVSR++) is a **removed** as "works poorly". DeepMosaics has two engines: **image** (per-frame) and
possible future engine, not wired. **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 - Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was
rejected early — a generative *conditioning* model, not a censorship restorer. Don't rejected early — a generative *conditioning* model, not a censorship restorer. Don't
reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/ 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. - **OS:** Windows 11 x64 (primary). Use PowerShell syntax in commands.
- **Python:** 3.11+. - **Python:** 3.11+.
- **GPU:** NVIDIA + CUDA via PyTorch, only for the YOLO detector. CPU fallback works - **GPU:** NVIDIA + CUDA via PyTorch, for the YOLO detector and the DeepMosaics
but is slow. The `classic` detector needs no torch and no GPU. restorer. CPU fallback works but is slow (esp. DeepMosaics / the temporal BVDNet).
## Tech stack (decided) ## Tech stack (decided)
@@ -76,12 +79,12 @@ Keep this scope sharp:
|----------------|---------------------------------------------| |----------------|---------------------------------------------|
| GUI | PySide6 (Qt 6) — LGPL | | GUI | PySide6 (Qt 6) — LGPL |
| Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` | | Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` |
| Detector | 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 Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional
optional (the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is (the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed
installed separately per the README). The classic detector must keep running with no separately per the README). Both detection (YOLO) and restoration (DeepMosaics) now
torch present. require torch — there's no longer a torch-free detector.
## Architecture (as implemented) ## Architecture (as implemented)
@@ -115,21 +118,17 @@ hvideotool/
├── video/ ├── video/
│ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ ├── 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 │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
├── restore/ # "un-censor" detected regions (per-frame) ├── restore/ # "un-censor" (DeepMosaics only; per-frame OR temporal)
│ ├── base.py # Restorer ABC: restore(image, detections, should_cancel=None) -> image; Cancelled exc │ ├── base.py # Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc
│ ├── factory.py # build_restorer(name, config) -> inpaint | deepmosaics (lada = TODO) │ ├── factory.py # build_restorer(name, config) -> deepmosaics | deepmosaics_video (lada = TODO)
│ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs │ ├── deepmosaics.py # DeepMosaicsRestorer (image, per-frame) + DeepMosaicsVideoRestorer (BVDNet, temporal); in-process, load once; uses _deepmosaics/
── deepmosaics.py # DeepMosaicsRestorer — in-process, loads models once; uses _deepmosaics/ ── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
│ ├── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import ── detection/ # YOLO only
│ └── mask.py # detections_to_mask(shape, dets, dilate)
└── detection/
├── base.py # Detector ABC: detect(frame) -> list[Detection] ├── 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 ├── types.py # Detection (+ to_dict/from_dict), CensorType enum
├── cache.py # save/load the project detection cache (detections.json): cache_file + base_dir args ├── 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
├── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics
└── composite.py # CompositeDetector — merge detectors + IoU dedup
``` ```
### How it works ### How it works
@@ -146,15 +145,21 @@ hvideotool/
`engine.detect/restore` all run **inside the worker** (`_compute` is the pure `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 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 GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables
`detector_combo`/`model_action` for the duration (they'd race the running detector). `model_action` for the duration (it'd race the running detector). This is the
This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics per-call
per-call latency can't be hidden with `processEvents`). latency can't be hidden with `processEvents`).
- **Device badge.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green) - **Device badge + CUDA diagnostics.** A clickable status-bar chip (`device_badge`) shows
or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background "⚡ CUDA" (green) or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in
`Job` at startup (importing torch is slow, so it's off the GUI thread) → `_set_device_badge`. a background `Job` at startup (it imports torch AND shells out to `nvidia-smi`, so it's off
Clicking (`_show_device_info`) opens a diagnostic dialog: `torch_info.reason()` explains the GUI thread) → `_set_device_badge`. `gather()` collects torch facts (version, built_cuda,
why CPU (no torch / CPU-only `+cpu` build / built-with-CUDA-but-no-GPU) plus cuda_available, device_name) **and** NVIDIA facts (gpus, driver_version, max cuda_driver).
`install_hint()` (the cu121 pip command). `core/torch_info.py` is pure (no Qt). 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`; - **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 `_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: "Создать проект…" (navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…"
@@ -179,7 +184,7 @@ hvideotool/
"Рассчитать кадр" action (Space → `_recompute_current`, force-recomputes current), or "Рассчитать кадр" action (Space → `_recompute_current`, force-recomputes current), or
"Детектировать все" (whole folder, progress bar). Do NOT re-add auto-detect-on-select. "Детектировать все" (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. 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 - **Detection cache (persisted).** `_results` is mirrored to the project's
`detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by `detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by
**basename** so it survives moving the project). `cache.save_results`/`load_results` **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` 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 avoids clobbering (`foo.jpg``foo (1).jpg`). Use case: flag good frames into a
training/example set while inspecting detections. training/example set while inspecting detections.
- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via - **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it
`build_restorer`) on the current frame's detections (computing them first if needed), **locates the mosaic itself** — so restoration is **fully decoupled from detection**: no
**on a background job** (`_restore_current` builds an `fn` that detects-if-needed + detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar
restores in the worker; a `tick` caches freshly-computed detections, `done` stores action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on
`_restored[path]` + shows it). "Показать оригинал/ a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать
результат" toggles (`_showing_restored`); "Сохранить результат" writes оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes
`<stem>_restored.jpg` beside the frame. The baseline `<stem>_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново",
is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run `_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every
**in-process** from the vendored `_deepmosaics/` code: it loads the BiSeNet locator + frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`,
clean generator **once** (lazy, cached on the instance) and per frame reproduces their basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the
`cleanmosaic_img_server` (locate mosaic → run generator on the crop → feather back), per-frame engine **skips frames already in `restored/`** unless `force` (resume). The
~0.3 s/frame cached on CPU vs ~7 s when it spawned a subprocess. Use the **image** engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the
model `clean_youknow_resnet_9blocks.pth` — the video model (BVDNet) is rejected per vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy,
frame (needs a neighbour). `should_cancel` (= `lambda: job.cancelled`) is polled so cached on the instance):
"■ Стоп" stops it; the engine raises `Cancelled`, which `Job.run` reports as a clean cancel. - `deepmosaics` (image, per-frame): reproduces `cleanmosaic_img_server` (locate mosaic
The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…), → run generator on the crop → feather back), ~0.3 s/frame cached on CPU. Image model
persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE: `clean_youknow_resnet_9blocks.pth`.
DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the - `deepmosaics_video` (**temporal, BVDNet**): `DeepMosaicsVideoRestorer`, `.temporal=True`.
clean weights) — our detections aren't passed to it. `_show` resets `_showing_restored` Reproduces `cleanmosaic_video_fusion` — per target frame it feeds the net a window of
+ `_update_restore_actions`. To add another engine (e.g. LADA), implement `T=5` neighbour frames sampled at step `S=3` around it (`N=2` each side, clamped at the
`core/restore/base.Restorer` and register it in `restore/factory.build_restorer`. 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 - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label`
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
@@ -255,20 +276,20 @@ hvideotool/
- `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via - `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
`core/detection/factory.build_detector` and talks to `core/` through the `Detector` `core/detection/factory.build_detector` and talks to `core/` through the `Detector`
interface and the `Detection`/`CensorType` types. interface and the `Detection`/`CensorType` types.
- New detector kinds: implement `core/detection/base.Detector`, register the string in - Detection is YOLO-only and restoration is DeepMosaics-only. If you re-add an engine
`core/detection/factory.build_detector`, and add it to `_DETECTORS` in kind, implement `core/detection/base.Detector` / `core/restore/base.Restorer`, register
`ui/main_window.py`. 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 ## Commands
```powershell ```powershell
python -m venv .venv; .\.venv\Scripts\Activate.ps1 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 # 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 python -m hvideotool "C:\path\to\MyProject" --model models\lada_mosaic_detection_model_v4_accurate.pt
pip install -e ".[yolo]" # + install torch separately, see README
``` ```
No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a 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 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 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. `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 **A model is auto-picked on project open** via `MainWindow._ensure_model()`
`MainWindow._auto_find_model()`: it scans `./models/**.pt` and matches only filenames `_auto_find_model()`: it scans `./models/**.pt` and matches only filenames containing
containing `lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; it `lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; otherwise the
falls back to the "Модель…" file dialog only when nothing suitable is found. user picks via "Модель…" (`_choose_model`).
- **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low - **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was
block-reconstruction residual + 2D gradient + contrast) fires on textured real noisy/approximate on real footage (false positives on skin/hair/fabric/JPEG; missed
footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing real mosaic after downscale) and the `combined` mode + cv2 `inpaint` baseline went with
real mosaic after the `proc_max_dim=720` downscale softens block edges (measured: it. Detection is YOLO-only, restoration is DeepMosaics-only. `normalize_config` coerces
contrast/grad fall below `mosaic_contrast_min`/`mosaic_grad_min`). For real-video any leftover `classic`/`combined`/`inpaint` in old settings/projects to `yolo`/`deepmosaics`.
mosaic use `yolo`/`combined` + LADA. For anime there is no good public model.
- **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic - **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic
but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`), but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`).
not tuning more classic thresholds.
- **YOLO detector = LADA weights** ([HF `ladaapp/lada`](https://huggingface.co/ladaapp/lada)). - **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 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 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 mosaic only. Weights + Ultralytics are AGPL-3.0 (accepted). `yolo.py` lazy-imports
(accepted). `yolo.py` lazy-imports `torch`/`ultralytics`. `torch`/`ultralytics`.
- **No model weights in the repo.** Code must fail with a clear, actionable message - **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`, when the model path is missing — not a raw stack trace (`factory._require_model`,
`YoloDetector.__init__`). `YoloDetector.__init__`).
@@ -318,10 +337,11 @@ frame directly.
- **CPU-only torch must not request CUDA.** A `+cpu` torch build raises "Torch not - **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 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 for this: `YoloDetector` picks `cuda` only when `torch.cuda.is_available()` (even an
explicit `yolo_device="cuda"` is downgraded to cpu); `DeepMosaicsRestorer._ensure_loaded` explicit `yolo_device="cuda"` is downgraded to cpu); both `DeepMosaicsRestorer._ensure_loaded`
forces `gpu_id="-1"` when CUDA is absent (its vendored `model_util.todevice` / and `DeepMosaicsVideoRestorer._ensure_loaded` force `gpu_id="-1"` when CUDA is absent (the
`data.im2tensor` call `.cuda()` for any `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). vendored `model_util.todevice` / `data.im2tensor`/`to_tensor` call `.cuda()` for any
So a wrong/CPU-only torch falls back to CPU instead of crashing. `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`), - **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`),
otherwise it aliases a buffer that gets freed → garbage/crash. otherwise it aliases a buffer that gets freed → garbage/crash.
- **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images —
+66 -57
View File
@@ -64,9 +64,10 @@
кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры
с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена, с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена,
🟢 проверено и чисто). 🟢 проверено и чисто).
- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** - Детекция — через **YOLO** (модель LADA для мозаики). Порог уверенности
уверенности прямо в тулбаре — удобно сравнивать. настраивается прямо в тулбаре.
- Выбор файла весов модели кнопкой **«Модель…»**. - Выбор файла весов модели кнопкой **«Модель…»** (нужная модель ищется в `models/`
автоматически при открытии проекта).
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU» - **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU
@@ -96,14 +97,24 @@
`<имя>_restored.jpg` рядом с кадром. Движок выбирается в `<имя>_restored.jpg` рядом с кадром. Движок выбирается в
меню **Файл → Движок восстановления…**. меню **Файл → Движок восстановления…**.
Два движка: **Пакетный прогон по диапазону.** Кнопки **«Расцензурить все»** (дозапуск — пропускает
уже сделанные) и **«Все заново»** обрабатывают **все кадры проекта** в фоне и пишут
результаты в папку **`restored/`** проекта (имена кадров сохраняются; папка держится
отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс,
отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции.
- **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по Расцензуривание — только через **DeepMosaics** (код **встроен** в приложение, vendored,
окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). GPL-3.0; ставить отдельно не нужно — требуются только **веса** и желательно **GPU
- **DeepMosaics** — реальное генеративное удаление мозаики. Код **встроен** в NVIDIA/CUDA**). Два движка:
приложение (vendored, GPL-3.0), ставить его отдельно не нужно — требуются только
**веса** и (желательно) **GPU NVIDIA/CUDA**. На аниме качество ограничено (модели - **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**.
обучены на реальном видео). - **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры**
(окно ±2 кадра с шагом 3) и собственный предыдущий результат для когерентности на
роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е.
запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к
окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`.
На аниме качество ограничено (модели обучены на реальном видео).
### Настройка DeepMosaics ### Настройка DeepMosaics
@@ -112,18 +123,20 @@
([Google Drive](https://drive.google.com/drive/folders/1LTERcN33McoiztYEwBxMuRjjgxh4DEPs), ([Google Drive](https://drive.google.com/drive/folders/1LTERcN33McoiztYEwBxMuRjjgxh4DEPs),
Baidu код `1x0a`): Baidu код `1x0a`):
- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель; - **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель (движок «картинка»);
- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом). - **`clean_youknow_video.pth`** — видеомодель BVDNet (движок «видео», соседние кадры);
- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом, нужен обоим).
Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, выберите **модель Затем в приложении: **Файл → Движок восстановления…**, выберите движок (**DeepMosaics —
из выпадающего списка** (наполняется из `models/deepmosaics`; есть «Обзор…» для файла в картинка** или **видео**) и **модель из выпадающего списка** (наполняется из
другом месте) и GPU id (`-1` = CPU). Если веса в `models/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/` и поэтому **весь > Код DeepMosaics (GPL-3.0) лежит в `core/restore/_deepmosaics/` и поэтому **весь
> проект распространяется под GPL-3.0**. Запускается на современных `torch 2.x`/ > проект распространяется под GPL-3.0**. Запускается на современных `torch 2.x`/
@@ -143,8 +156,8 @@ Baidu код `1x0a`):
- **ОС:** Windows 11 x64 (основная целевая платформа). - **ОС:** Windows 11 x64 (основная целевая платформа).
- **Python:** 3.11+. - **Python:** 3.11+.
- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора. CPU-режим работает, но - **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора и DeepMosaics. CPU-режим
медленный. Для `classic` детектора ни torch, ни GPU не нужны. работает, но медленный (особенно DeepMosaics / временно́й BVDNet).
## Установка ## Установка
@@ -153,12 +166,11 @@ git clone https://github.com/mrleo1nid/HVideoTool.git
cd HVideoTool cd HVideoTool
python -m venv .venv python -m venv .venv
.\.venv\Scripts\Activate.ps1 .\.venv\Scripts\Activate.ps1
pip install -e . pip install -e ".[yolo]"
``` ```
Этого достаточно для `classic` детектора — **PyTorch/CUDA не требуются**. Они Детекция (YOLO) и расцензуривание (DeepMosaics) требуют **PyTorch** — установите его
нужны только для YOLO/комбинированного детектора (см. отдельно под вашу CUDA (см. [Модель детектора](#модель-детектора)).
[Модель детектора](#модель-детектора)).
## Запуск ## Запуск
@@ -166,29 +178,23 @@ pip install -e .
# Без аргументов — открывается последний проект (или создайте/откройте новый в тулбаре) # Без аргументов — открывается последний проект (или создайте/откройте новый в тулбаре)
python -m hvideotool python -m hvideotool
# Необязательно: сразу открыть проект / переопределить детектор и модель (по умолчанию) # Необязательно: сразу открыть проект и указать модель YOLO
python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt 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`), детектор Детекция — **только YOLO** (интерфейс абстрагирован в `core/detection/base.py`):
выбирается в тулбаре (или флагом `--detector`): ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics)
(`core/detection/yolo.py`), **сегментационная** модель — маски превращаются в
- `classic` — эвристический classic-CV детектор (`core/detection/classic_cv.py`). контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada).
Различает `mosaic` / `blur` / `black_bar`, без весов и без GPU. (Старый эвристический classic-CV детектор и комбинированный режим удалены — давали
**Приблизительный**: на реальном видео даёт много ложных срабатываний, заточен много ложных срабатываний.)
скорее под рисованный/аниме контент — но и там ненадёжен.
- `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 нужна модель **детекции цензуры** > **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры**
> (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать > (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать
@@ -207,7 +213,7 @@ python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model mod
```powershell ```powershell
python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic
python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100 python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100
# затем: тулбар → Детектор yolo → «Модель…» → runs\segment\mosaic\weights\best.pt # затем: тулбар → «Модель…» → runs\segment\mosaic\weights\best.pt
``` ```
Установка YOLO-детектора: Установка 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 ` 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" "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-пути) ├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути)
├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки ├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки
├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора ├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора
── detection/ ── detection/ # только YOLO
├── base.py # Detector (ABC): detect(frame) -> list[Detection] ├── base.py # Detector (ABC): detect(frame) -> list[Detection]
├── factory.py # build_detector(config) -> classic/yolo/combined ├── factory.py # build_detector(config) -> yolo
├── types.py # Detection (+ to_dict/from_dict), CensorType ├── types.py # Detection (+ to_dict/from_dict), CensorType
├── cache.py # сохранение/загрузка кэша детекций (detections.json в проекте) ├── cache.py # сохранение/загрузка кэша детекций (detections.json)
── classic_cv.py # эвристический детектор (mosaic/blur/black_bar) ── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны)
├── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны) └── restore/ # только DeepMosaics
── composite.py # CompositeDetector: объединение детекторов ── 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 Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU
@@ -264,9 +273,9 @@ hvideotool/
| Язык | Python ≥ 3.11 | | Язык | Python ≥ 3.11 |
| GUI | PySide6 (Qt 6) | | GUI | PySide6 (Qt 6) |
| Обработка картинок | OpenCV / NumPy | | Обработка картинок | OpenCV / NumPy |
| Детектор (без весов) | classic-CV эвристика (без GPU) | | Детектор | Ultralytics YOLO + PyTorch/CUDA (веса LADA) |
| Детектор (ML) | Ultralytics YOLO + PyTorch/CUDA (LADA) | | Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA |
## Лицензия ## Лицензия
TBD. **GPL-3.0** — из-за встроенного кода DeepMosaics (GPL-3.0); веса/код YOLO LADA — AGPL-3.0.
+2 -4
View File
@@ -12,7 +12,7 @@ import sys
from . import settings_store from . import settings_store
from .app import run from .app import run
from .config import AppConfig from .config import AppConfig, normalize_config
def main() -> int: def main() -> int:
@@ -21,15 +21,13 @@ def main() -> int:
description="Инспектор детекции уже наложенной цензуры на картинках.", description="Инспектор детекции уже наложенной цензуры на картинках.",
) )
parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)") 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)") parser.add_argument("--model", dest="model_path", default=None, help="путь к весам (YOLO)")
args = parser.parse_args() args = parser.parse_args()
config = AppConfig() config = AppConfig()
settings_store.apply(config) # persisted defaults first settings_store.apply(config) # persisted defaults first
normalize_config(config) # drop any legacy classic/inpaint values
if args.detector:
config.detector = args.detector
if args.model_path: if args.model_path:
config.model_path = args.model_path config.model_path = args.model_path
+21 -27
View File
@@ -1,40 +1,21 @@
"""Application configuration and tunable defaults. """Application configuration and tunable defaults.
Plain dataclasses. The detection thresholds matter most here — this tool is now Plain dataclasses. Detection is YOLO-only and restoration is DeepMosaics-only, so the
an image-folder inspector for tuning the detectors, so keep them easy to tweak. knobs here are the YOLO inference params, the overlay style, and the DeepMosaics weights.
""" """
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass, field from dataclasses import dataclass, field
DETECTORS = ("yolo",)
RESTORERS = ("deepmosaics", "deepmosaics_video")
@dataclass(frozen=True) @dataclass(frozen=True)
class DetectionConfig: class DetectionConfig:
"""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_conf: float = 0.2 # confidence threshold (LADA recommends ~0.2)
yolo_imgsz: int = 640 # inference image size yolo_imgsz: int = 640 # inference image size
yolo_device: str | None = None # None => auto ("cuda" if available, else "cpu") yolo_device: str | None = None # None => auto ("cuda" if available, else "cpu")
@@ -62,13 +43,26 @@ class OverlayConfig:
class AppConfig: class AppConfig:
detection: DetectionConfig = field(default_factory=DetectionConfig) detection: DetectionConfig = field(default_factory=DetectionConfig)
overlay: OverlayConfig = field(default_factory=OverlayConfig) overlay: OverlayConfig = field(default_factory=OverlayConfig)
detector: str = "classic" # "classic" | "yolo" | "combined" detector: str = "yolo" # only "yolo"
model_path: str | None = None # weights path, used by the YOLO detector model_path: str | None = None # weights path, used by the YOLO detector
default_threshold: float = 0.20 # initial overlay confidence threshold default_threshold: float = 0.20 # initial overlay confidence threshold
# --- restoration ("расцензурить") --- # --- restoration ("расцензурить") ---
restorer: str = "inpaint" # "inpaint" | "deepmosaics" restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py) dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py)
dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth) dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
dm_python: str | None = None # python exe for DeepMosaics (None = current) dm_python: str | None = None # python exe for DeepMosaics (None = current)
dm_gpu: str = "0" # CUDA device id, "-1" for CPU dm_gpu: str = "0" # CUDA device id, "-1" for CPU
def normalize_config(cfg: AppConfig) -> None:
"""Coerce legacy/removed settings to supported values (mutates ``cfg``).
Old projects / settings.json may carry the removed ``classic``/``combined``
detectors or the ``inpaint`` restorer — map those onto the survivors so loading
them doesn't blow up at build time.
"""
if cfg.detector not in DETECTORS:
cfg.detector = "yolo"
if cfg.restorer not in RESTORERS:
cfg.restorer = "deepmosaics"
-216
View File
@@ -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
-51
View File
@@ -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
+5 -17
View File
@@ -3,14 +3,15 @@
Kept separate from ``app.py`` so both the app bootstrap and the UI can build Kept separate from ``app.py`` so both the app bootstrap and the UI can build
detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on
bad config so the GUI can show the message instead of exiting. bad config so the GUI can show the message instead of exiting.
Only the YOLO detector is supported — the classic-CV heuristic (and the composite
mode that combined them) were removed: they were noisy/approximate on real footage.
""" """
from __future__ import annotations from __future__ import annotations
from ...config import AppConfig from ...config import AppConfig
from .base import Detector from .base import Detector
from .classic_cv import ClassicCVDetector
from .types import CensorType
def _require_model(config: AppConfig) -> str: def _require_model(config: AppConfig) -> str:
@@ -23,19 +24,6 @@ def _require_model(config: AppConfig) -> str:
def build_detector(config: AppConfig) -> Detector: def build_detector(config: AppConfig) -> Detector:
if config.detector == "classic": from .yolo import YoloDetector # lazy: pulls torch/ultralytics
return ClassicCVDetector(config.detection)
if config.detector == "yolo":
from .yolo import YoloDetector # lazy: pulls torch/ultralytics
return YoloDetector(_require_model(config), config.detection) 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}")
+1 -1
View File
@@ -7,7 +7,7 @@ a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose cl
names map onto :class:`CensorType`. names map onto :class:`CensorType`.
Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the 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. Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
""" """
+7
View File
@@ -30,6 +30,7 @@ FRAMES_DIR = "frames"
CACHE_FILE = "detections.json" CACHE_FILE = "detections.json"
COLLECTIONS_DIR = "collections" COLLECTIONS_DIR = "collections"
FAVORITES_DIR = "Избранное" # the single default collection ("в избранное") FAVORITES_DIR = "Избранное" # the single default collection ("в избранное")
RESTORED_DIR = "restored" # batch "расцензурить все" output (kept out of frames/)
_VERSION = 1 _VERSION = 1
# The subset of AppConfig fields a project remembers (mirrored to/from project.json). # The subset of AppConfig fields a project remembers (mirrored to/from project.json).
@@ -78,6 +79,12 @@ class Project:
"""The single default collection — frames moved "to favorites" land here.""" """The single default collection — frames moved "to favorites" land here."""
return self.collections_dir / FAVORITES_DIR 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 # ------------------------------------------------------------- lifecycle
@classmethod @classmethod
def create( def create(
+36 -3
View File
@@ -1,9 +1,10 @@
"""Restorer interface — "un-censor" detected regions of an image. """Restorer interface — "un-censor" detected regions of an image.
A Restorer takes an image plus the detected censored regions and returns a new A Restorer takes an image plus the detected censored regions and returns a new
image with those regions reconstructed/filled. This mirrors the ``Detector`` image with those regions reconstructed. This mirrors the ``Detector`` abstraction
abstraction so different engines (classic inpaint now; a generative model like so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug
DeepMosaics / LADA later) plug in behind the same interface. 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 from __future__ import annotations
@@ -24,7 +25,18 @@ class Cancelled(Exception):
"""Raised by a Restorer when ``should_cancel`` asked it to stop.""" """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): 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 @property
def name(self) -> str: def name(self) -> str:
return type(self).__name__ return type(self).__name__
@@ -42,3 +54,24 @@ class Restorer(ABC):
True the engine should abort and raise :class:`Cancelled`. True the engine should abort and raise :class:`Cancelled`.
""" """
raise NotImplementedError 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))
+171 -13
View File
@@ -26,18 +26,29 @@ from types import SimpleNamespace
import numpy as np import numpy as np
from ..detection.types import Detection 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" _VENDOR = Path(__file__).parent / "_deepmosaics"
# Default place to drop DeepMosaics clean weights (gitignored — see models/). # Default place to drop DeepMosaics clean weights (gitignored — see models/).
DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics" DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"
def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]: def discover_models(
"""Find usable per-frame clean models: (display_name, full_path). 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) 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] dirs = [DEFAULT_WEIGHTS_DIR]
if extra_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(): if not d.is_dir():
continue continue
for p in sorted(d.glob("clean_*.pth")): 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 continue
seen.add(p.name) seen.add(p.name)
out.append((p.stem, str(p))) out.append((p.stem, str(p)))
return out 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: def _netg_kind(model_name: str) -> str:
"""Pick DeepMosaics' netG type from the weights filename (see their options.py).""" """Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
n = model_name.lower() n = model_name.lower()
@@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer):
) )
model = Path(model_path) model = Path(model_path)
self._netg = _netg_kind(model.name) # raises on a video model 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: if pos is None:
raise ValueError( raise ValueError(
"Рядом с clean-моделью не найден mosaic_position.pth.\n" "Рядом с clean-моделью не найден mosaic_position.pth.\n"
@@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer):
self._gpu = gpu_id self._gpu = gpu_id
self._loaded = False # models loaded lazily on first restore 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 @property
def name(self) -> str: def name(self) -> str:
return f"DeepMosaics(gpu={self._gpu})" 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_mosaic = work[y - size:y + size, x - size:x + size]
img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt) img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather) 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))
+18 -11
View File
@@ -1,9 +1,13 @@
"""Restorer factory: build a Restorer from the app config. """Restorer factory: build a Restorer from the app config.
- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct). Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but
- ``deepmosaics``: real generative mosaic removal. The DeepMosaics network code is did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``,
vendored (``_deepmosaics/``, GPL-3.0) and run in-process; the user supplies only the GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth``
clean weights (+ ``mosaic_position.pth`` alongside). A CUDA GPU is recommended. 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 from __future__ import annotations
@@ -11,25 +15,28 @@ from __future__ import annotations
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from .base import Restorer from .base import Restorer
from .inpaint import InpaintRestorer
if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed) if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed)
from ...config import AppConfig from ...config import AppConfig
def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer: def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer:
if name == "inpaint": if config is None:
return InpaintRestorer() raise ValueError("Для DeepMosaics нужны настройки (config).")
if name == "deepmosaics": if name == "deepmosaics":
from .deepmosaics import DeepMosaicsRestorer from .deepmosaics import DeepMosaicsRestorer
if config is None:
raise ValueError("Для DeepMosaics нужны настройки (config).")
return DeepMosaicsRestorer( return DeepMosaicsRestorer(
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu 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": if name == "lada":
raise ValueError( raise ValueError(
"Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README." "Движок LADA пока не подключён. Используйте DeepMosaics. См. README."
) )
raise ValueError(f"Неизвестный режим восстановления: {name!r}") raise ValueError(f"Неизвестный режим восстановления: {name!r}")
-41
View File
@@ -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)
-26
View File
@@ -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
+191 -48
View File
@@ -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 Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to
thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and ``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on
how to enable the GPU. the gathered dict and figures out the most likely cause + concrete fix.
""" """
from __future__ import annotations from __future__ import annotations
# pip index for the CUDA build (matches the README). import re
CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121" 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: 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 = { info: dict = {
"installed": False, "installed": False,
"version": None, # torch.__version__ (e.g. "2.12.0+cpu") "version": None, # torch.__version__ (e.g. "2.12.0+cpu")
@@ -20,27 +68,38 @@ def gather() -> dict:
"cuda_available": False, "cuda_available": False,
"device_name": None, # the active GPU's name, if any "device_name": None, # the active GPU's name, if any
"import_error": None, "import_error": None,
# filled by nvidia-smi:
"nvidia_smi": False,
"gpus": [],
"driver_version": None,
"cuda_driver": None,
} }
try: try:
import torch 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) info["import_error"] = str(exc)
return info else:
info["installed"] = True info["installed"] = True
info["version"] = getattr(torch, "__version__", None) 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"]:
try: try:
info["device_name"] = torch.cuda.get_device_name(0) info["built_cuda"] = torch.version.cuda
except Exception: # noqa: BLE001 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 return info
@@ -48,33 +107,117 @@ def device_label(info: dict) -> str:
return "CUDA" if info.get("cuda_available") else "CPU" return "CUDA" if info.get("cuda_available") else "CPU"
def reason(info: dict) -> str: def _ver_tuple(v: str | None) -> tuple[int, ...]:
"""One-sentence human explanation of the current device choice.""" try:
if not info.get("installed"): return tuple(int(x) for x in str(v).split(".")[:2])
return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics " except (ValueError, AttributeError):
"работают на CPU (классический детектор torch не требует).") return ()
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 install_hint() -> str: def recommend_channel(info: dict) -> str:
"""Steps to enable the GPU (shown when running on CPU).""" """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 ( return (
"Как включить GPU (NVIDIA):\n" f"{q} -m pip uninstall -y torch torchvision\n"
"1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n" f"{q} -m pip install torch torchvision --index-url {index}"
"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 детекция и расцензуривание "
"просто медленнее."
) )
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}
+156 -68
View File
@@ -26,11 +26,10 @@ import shutil
from pathlib import Path from pathlib import Path
from PySide6.QtCore import Qt, QThreadPool 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 ( from PySide6.QtWidgets import (
QAbstractItemView, QAbstractItemView,
QApplication, QApplication,
QComboBox,
QDialog, QDialog,
QDoubleSpinBox, QDoubleSpinBox,
QFileDialog, QFileDialog,
@@ -41,6 +40,7 @@ from PySide6.QtWidgets import (
QListWidgetItem, QListWidgetItem,
QMainWindow, QMainWindow,
QMessageBox, QMessageBox,
QPlainTextEdit,
QProgressBar, QProgressBar,
QPushButton, QPushButton,
QSplitter, QSplitter,
@@ -51,7 +51,7 @@ from PySide6.QtWidgets import (
) )
from .. import settings_store 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 import cache as detection_cache
from ..core.detection.factory import build_detector from ..core.detection.factory import build_detector
from ..core.detection.types import Detection from ..core.detection.types import Detection
@@ -68,7 +68,6 @@ from .workers import Job
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)" _VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
_DETECTORS = ["classic", "yolo", "combined"]
class MainWindow(QMainWindow): class MainWindow(QMainWindow):
@@ -120,6 +119,8 @@ class MainWindow(QMainWindow):
file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True)) file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True))
file_menu.addSeparator() file_menu.addSeparator()
file_menu.addAction("Движок восстановления…", self._open_restore_settings) 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.addSeparator()
file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M") file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M")
file_menu.addSeparator() file_menu.addSeparator()
@@ -136,14 +137,9 @@ class MainWindow(QMainWindow):
tb.addAction(from_video) tb.addAction(from_video)
tb.addSeparator() tb.addSeparator()
tb.addWidget(QLabel(" Детектор: ")) tb.addWidget(QLabel(" Детектор: YOLO "))
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)
self.model_action = QAction("Модель…", self, triggered=self._choose_model) self.model_action = QAction("Модель…", self, triggered=self._choose_model)
self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики")
tb.addAction(self.model_action) tb.addAction(self.model_action)
tb.addSeparator() tb.addSeparator()
@@ -166,6 +162,12 @@ class MainWindow(QMainWindow):
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current) restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
restore.setToolTip("Восстановить найденные области на текущем кадре") restore.setToolTip("Восстановить найденные области на текущем кадре")
tb.addAction(restore) 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 = QAction("Показать оригинал", self, triggered=self._toggle_restored)
self.toggle_restored_action.setEnabled(False) self.toggle_restored_action.setEnabled(False)
tb.addAction(self.toggle_restored_action) 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() info = self._device_info if self._device_info else torch_info.gather()
cuda = bool(info.get("cuda_available")) cuda = bool(info.get("cuda_available"))
a = torch_info.analyze(info)
self._install_command = a["command"] # what the Copy button will copy
lines = [ lines = [
torch_info.reason(info), "ВЕРДИКТ:",
a["summary"],
"", "",
"Диагностика:", "Диагностика:",
f"PyTorch: {info.get('version') or 'не установлен'}", *(f"{d}" for d in a["details"]),
f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}", "",
f" • CUDA доступна: {'да' if cuda else 'нет'}", "Что делать:",
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) # A real dialog (not QMessageBox) so the text — incl. the install command — is
box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning) # selectable, and a Copy button drops the pip command straight onto the clipboard.
box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU")) dlg = QDialog(self)
box.setText("\n".join(lines)) 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: if not cuda:
box.setInformativeText(torch_info.install_hint()) copy_btn = QPushButton("Скопировать команду установки")
box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands copy_btn.clicked.connect(self._copy_install_command)
box.exec() 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 # ------------------------------------------------------------- cancellation
def _begin_busy(self, total: int | None = None) -> None: def _begin_busy(self, total: int | None = None) -> None:
@@ -388,7 +421,6 @@ class MainWindow(QMainWindow):
self._cancel = False self._cancel = False
self.stop_action.setEnabled(True) self.stop_action.setEnabled(True)
# Disable inputs that would race a running job (they clear cache / rebuild engines). # Disable inputs that would race a running job (they clear cache / rebuild engines).
self.detector_combo.setEnabled(False)
self.model_action.setEnabled(False) self.model_action.setEnabled(False)
if total is None: if total is None:
self.progress.setRange(0, 0) # indeterminate self.progress.setRange(0, 0) # indeterminate
@@ -400,7 +432,6 @@ class MainWindow(QMainWindow):
def _end_busy(self) -> None: def _end_busy(self) -> None:
self._busy = False self._busy = False
self.stop_action.setEnabled(False) self.stop_action.setEnabled(False)
self.detector_combo.setEnabled(True)
self.model_action.setEnabled(True) self.model_action.setEnabled(True)
self.progress.setVisible(False) self.progress.setVisible(False)
self.progress.setRange(0, 100) # leave it determinate for the next user self.progress.setRange(0, 100) # leave it determinate for the next user
@@ -459,19 +490,19 @@ class MainWindow(QMainWindow):
self._detector_key = key self._detector_key = key
return self._detector return self._detector
def _on_detector_changed(self, name: str) -> None: def _ensure_model(self) -> None:
self._cfg.detector = name """Make sure the YOLO detector has weights — auto-pick from ./models silently.
# 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). Called on project open. Doesn't prompt (the user can pick via "Модель…"); the
if name in ("yolo", "combined") and not self._cfg.model_path: detector factory raises a clear message if a detect is attempted without one.
found = self._auto_find_model() """
if found: if self._cfg.model_path and Path(self._cfg.model_path).is_file():
self._cfg.model_path = found return
self.statusBar().showMessage(f"Модель найдена автоматически: {found}") found = self._auto_find_model()
else: if found:
self._choose_model() self._cfg.model_path = found
self._persist_settings() self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}")
self._invalidate_results() self._persist_settings()
@staticmethod @staticmethod
def _auto_find_model() -> str | None: def _auto_find_model() -> str | None:
@@ -667,6 +698,8 @@ class MainWindow(QMainWindow):
self._project = project self._project = project
project.frames_dir.mkdir(parents=True, exist_ok=True) project.frames_dir.mkdir(parents=True, exist_ok=True)
project.apply_to_config(self._cfg) # per-project settings -> live config project.apply_to_config(self._cfg) # per-project settings -> live config
normalize_config(self._cfg) # coerce any legacy classic/inpaint values
self._ensure_model() # YOLO needs weights — auto-pick if missing
self._sync_settings_ui() self._sync_settings_ui()
self._detector_key = None self._detector_key = None
self._restorer_key = None self._restorer_key = None
@@ -679,9 +712,6 @@ class MainWindow(QMainWindow):
def _sync_settings_ui(self) -> None: def _sync_settings_ui(self) -> None:
"""Reflect the (project's) config onto the toolbar widgets without signal loops.""" """Reflect the (project's) config onto the toolbar widgets without signal loops."""
self.detector_combo.blockSignals(True)
self.detector_combo.setCurrentText(self._cfg.detector)
self.detector_combo.blockSignals(False)
self.threshold_spin.blockSignals(True) self.threshold_spin.blockSignals(True)
self.threshold_spin.setValue(self._cfg.default_threshold) self.threshold_spin.setValue(self._cfg.default_threshold)
self.threshold_spin.blockSignals(False) self.threshold_spin.blockSignals(False)
@@ -912,41 +942,32 @@ class MainWindow(QMainWindow):
key, dets = payload key, dets = payload
self._results[key] = dets self._results[key] = dets
self._tag_file(Path(key), len(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 self._tick_count += 1
if self._tick_count % 25 == 0: if self._tick_count % 25 == 0:
self._refresh_marks() # let marks appear progressively (throttled) self._refresh_marks() # let marks appear progressively (throttled)
# ------------------------------------------------------------- restoration # ------------------------------------------------------------- restoration
def _restore_current(self) -> None: 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 DeepMosaics locates the mosaic itself, so no detection step is needed — we just
engine polls ``job.cancelled`` so "■ Стоп" stops it promptly.""" 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: if self._current is None or self._busy:
return return
path = self._current path = self._current
key = str(path) key = str(path)
def fn(job): 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) img = imread_unicode(key)
if img is None: if img is None:
raise RuntimeError(f"Не удалось прочитать: {path.name}") raise RuntimeError(f"Не удалось прочитать: {path.name}")
restorer = self._make_restorer() restorer = self._make_restorer()
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled) restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled)
return ("restored", key, restored, len(dets), restorer.name) return ("restored", key, restored, 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()
def done(result, cancelled): def done(result, cancelled):
if cancelled: if cancelled:
@@ -954,19 +975,86 @@ class MainWindow(QMainWindow):
return return
if result is None: if result is None:
return return
if result[0] == "empty": _, k, restored, engine = result
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
return
_, k, restored, n, engine = result
self._restored[k] = restored self._restored[k] = restored
if self._current is not None and str(self._current) == k: if self._current is not None and str(self._current) == k:
self._showing_restored = True self._showing_restored = True
self.view.set_image(restored, []) self.view.set_image(restored, [])
self._update_restore_actions() 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.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): def _make_restorer(self):
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,
+23 -12
View File
@@ -1,9 +1,11 @@
"""Configure the restoration ("расцензурить") engine. """Configure the restoration ("расцензурить") engine.
inpaint — no setup. deepmosaics — the network code is vendored (built-in); the user Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео")
picks a clean model from a dropdown of the bundled ``models/deepmosaics`` weights engine. The network code is vendored (built-in); the user picks a clean model from a
(or browses to another ``clean_*.pth``). ``mosaic_position.pth`` must sit beside the dropdown of the bundled ``models/deepmosaics`` weights (or browses to another
chosen model. A CUDA GPU is strongly recommended (GPU id, -1 = CPU/slow). ``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 from __future__ import annotations
@@ -35,9 +37,9 @@ class RestoreDialog(QDialog):
self.setMinimumWidth(560) self.setMinimumWidth(560)
self.engine = QComboBox() self.engine = QComboBox()
self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint") self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics")
self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics") self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video")
self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0) self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer)))
self.engine.currentIndexChanged.connect(self._sync) self.engine.currentIndexChanged.connect(self._sync)
# Model dropdown — bundled clean models, plus the configured one if external. # 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("Модель:", self._with_browse(self.model_combo, self._browse_model))
form.addRow("GPU id:", self.dm_gpu) form.addRow("GPU id:", self.dm_gpu)
hint = QLabel( hint = QLabel(
"Модели берутся из models/deepmosaics. Нужны clean_youknow_resnet_9blocks.pth " "Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n"
"и mosaic_position.pth (рядом). Видеомодель clean_*_video.pth покадрово не " "• Картинка: clean_youknow_resnet_9blocks.pth покадрово.\n"
"работает и в списке не показывается. На CPU медленно — лучше GPU. См. README." "• Видео: clean_youknow_video.pth — использует соседние кадры (когерентнее на "
"роликах), требует прогона по диапазону («Расцензурить все»).\n"
"На CPU медленно — лучше GPU. См. README."
) )
hint.setWordWrap(True) hint.setWordWrap(True)
form.addRow(hint) form.addRow(hint)
@@ -67,7 +71,12 @@ class RestoreDialog(QDialog):
def _populate_models(self, current: str | None) -> None: def _populate_models(self, current: str | None) -> None:
self.model_combo.clear() 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) self.model_combo.addItem(name, path)
# Keep an externally-configured model selectable even if it's outside the folder. # Keep an externally-configured model selectable even if it's outside the folder.
if current and self.model_combo.findData(current) < 0: if current and self.model_combo.findData(current) < 0:
@@ -88,7 +97,9 @@ class RestoreDialog(QDialog):
return w return w
def _sync(self) -> None: 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.model_combo.setEnabled(is_dm)
self.dm_gpu.setEnabled(is_dm) self.dm_gpu.setEnabled(is_dm)