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
@@ -47,12 +47,15 @@ what it found.
|
||||
Keep this scope sharp:
|
||||
|
||||
- Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить")
|
||||
step was added later (user-requested): per-frame, on-demand, behind a `Restorer`
|
||||
interface. Two engines: a cv2 **inpaint baseline** (fills, does NOT reconstruct) and
|
||||
**DeepMosaics** — real generative mosaic removal, its GPL-3.0 network code **vendored**
|
||||
under `core/restore/_deepmosaics/` and run in-process (user supplies only the weights).
|
||||
Because of that vendoring the **whole project is GPL-3.0**. LADA (BasicVSR++) is a
|
||||
possible future engine, not wired.
|
||||
step was added later (user-requested): on-demand (single frame **or** whole-project
|
||||
batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
|
||||
restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
|
||||
composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were
|
||||
**removed** as "works poorly". DeepMosaics has two engines: **image** (per-frame) and
|
||||
**video** (BVDNet, temporal — uses neighbour frames). Its GPL-3.0 network code is
|
||||
**vendored** under `core/restore/_deepmosaics/` and run in-process (user supplies only
|
||||
the weights). Because of that vendoring the **whole project is GPL-3.0**. LADA
|
||||
(BasicVSR++) is a possible future engine, not wired.
|
||||
- Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was
|
||||
rejected early — a generative *conditioning* model, not a censorship restorer. Don't
|
||||
reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/
|
||||
@@ -67,8 +70,8 @@ Keep this scope sharp:
|
||||
|
||||
- **OS:** Windows 11 x64 (primary). Use PowerShell syntax in commands.
|
||||
- **Python:** 3.11+.
|
||||
- **GPU:** NVIDIA + CUDA via PyTorch, only for the YOLO detector. CPU fallback works
|
||||
but is slow. The `classic` detector needs no torch and no GPU.
|
||||
- **GPU:** NVIDIA + CUDA via PyTorch, for the YOLO detector and the DeepMosaics
|
||||
restorer. CPU fallback works but is slow (esp. DeepMosaics / the temporal BVDNet).
|
||||
|
||||
## Tech stack (decided)
|
||||
|
||||
@@ -76,12 +79,12 @@ Keep this scope sharp:
|
||||
|----------------|---------------------------------------------|
|
||||
| GUI | PySide6 (Qt 6) — LGPL |
|
||||
| Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` |
|
||||
| Detector | classic-CV heuristic; Ultralytics YOLO (LADA) behind a pluggable interface |
|
||||
| Detector | Ultralytics YOLO (LADA weights) behind a pluggable interface (YOLO-only) |
|
||||
|
||||
Torch/CUDA + Ultralytics enter only with the YOLO detector. Keep that dependency
|
||||
optional (the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is
|
||||
installed separately per the README). The classic detector must keep running with no
|
||||
torch present.
|
||||
Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional
|
||||
(the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed
|
||||
separately per the README). Both detection (YOLO) and restoration (DeepMosaics) now
|
||||
require torch — there's no longer a torch-free detector.
|
||||
|
||||
## Architecture (as implemented)
|
||||
|
||||
@@ -115,21 +118,17 @@ hvideotool/
|
||||
├── video/
|
||||
│ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale
|
||||
│ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
|
||||
├── restore/ # "un-censor" detected regions (per-frame)
|
||||
│ ├── base.py # Restorer ABC: restore(image, detections, should_cancel=None) -> image; Cancelled exc
|
||||
│ ├── factory.py # build_restorer(name, config) -> inpaint | deepmosaics (lada = TODO)
|
||||
│ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs
|
||||
│ ├── deepmosaics.py # DeepMosaicsRestorer — in-process, loads models once; uses _deepmosaics/
|
||||
│ ├── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
|
||||
│ └── mask.py # detections_to_mask(shape, dets, dilate)
|
||||
└── detection/
|
||||
├── restore/ # "un-censor" (DeepMosaics only; per-frame OR temporal)
|
||||
│ ├── base.py # Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc
|
||||
│ ├── factory.py # build_restorer(name, config) -> deepmosaics | deepmosaics_video (lada = TODO)
|
||||
│ ├── deepmosaics.py # DeepMosaicsRestorer (image, per-frame) + DeepMosaicsVideoRestorer (BVDNet, temporal); in-process, load once; uses _deepmosaics/
|
||||
│ └── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
|
||||
└── detection/ # YOLO only
|
||||
├── base.py # Detector ABC: detect(frame) -> list[Detection]
|
||||
├── factory.py # build_detector(config) -> classic | yolo | combined
|
||||
├── factory.py # build_detector(config) -> yolo (the only kind)
|
||||
├── types.py # Detection (+ to_dict/from_dict), CensorType enum
|
||||
├── cache.py # save/load the project detection cache (detections.json): cache_file + base_dir args
|
||||
├── classic_cv.py # ClassicCVDetector — heuristic; accepts a `types` filter
|
||||
├── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics
|
||||
└── composite.py # CompositeDetector — merge detectors + IoU dedup
|
||||
└── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics
|
||||
```
|
||||
|
||||
### How it works
|
||||
@@ -146,15 +145,21 @@ hvideotool/
|
||||
`engine.detect/restore` all run **inside the worker** (`_compute` is the pure
|
||||
read+detect helper); the `fn` must touch NO Qt widgets — it emits plain data that the
|
||||
GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables
|
||||
`detector_combo`/`model_action` for the duration (they'd race the running detector).
|
||||
This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics
|
||||
per-call latency can't be hidden with `processEvents`).
|
||||
- **Device badge.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green)
|
||||
or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background
|
||||
`Job` at startup (importing torch is slow, so it's off the GUI thread) → `_set_device_badge`.
|
||||
Clicking (`_show_device_info`) opens a diagnostic dialog: `torch_info.reason()` explains
|
||||
why CPU (no torch / CPU-only `+cpu` build / built-with-CUDA-but-no-GPU) plus
|
||||
`install_hint()` (the cu121 pip command). `core/torch_info.py` is pure (no Qt).
|
||||
`model_action` for the duration (it'd race the running detector). This is the
|
||||
deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics per-call
|
||||
latency can't be hidden with `processEvents`).
|
||||
- **Device badge + CUDA diagnostics.** A clickable status-bar chip (`device_badge`) shows
|
||||
"⚡ CUDA" (green) or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in
|
||||
a background `Job` at startup (it imports torch AND shells out to `nvidia-smi`, so it's off
|
||||
the GUI thread) → `_set_device_badge`. `gather()` collects torch facts (version, built_cuda,
|
||||
cuda_available, device_name) **and** NVIDIA facts (gpus, driver_version, max cuda_driver).
|
||||
Clicking (`_show_device_info`) opens a diagnostic **QDialog** (not QMessageBox — its text
|
||||
wasn't copyable): a read-only monospace `QPlainTextEdit` with `torch_info.analyze(info)` →
|
||||
`{summary, details, steps, command}` — a verdict on *why* it's on CPU (CPU-only `+cpu`
|
||||
build / no GPU / driver-too-old-for-built-CUDA) and the exact pip fix. Buttons: "Скопировать
|
||||
команду установки" (`_copy_install_command` → the recommended cu121/cu118 command, picked by
|
||||
`recommend_channel` from the driver's CUDA) and "Проверить заново" (re-runs `_probe_device`).
|
||||
`core/torch_info.py` is pure (no Qt); subprocess uses `CREATE_NO_WINDOW` on Windows.
|
||||
- **Projects (`core/project.py`).** `MainWindow._project` is the open `Project`;
|
||||
`_folder` is kept as a synonym for `project.frames_dir` so the rest of the code
|
||||
(navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…"
|
||||
@@ -179,7 +184,7 @@ hvideotool/
|
||||
"Рассчитать кадр" action (Space → `_recompute_current`, force-recomputes current), or
|
||||
"Детектировать все" (whole folder, progress bar). Do NOT re-add auto-detect-on-select.
|
||||
Results cache in `_results`; the file-list row gets a count suffix when computed.
|
||||
Switching detector/model clears the cache (`_invalidate_results`).
|
||||
Switching the model clears the cache (`_choose_model` → `_invalidate_results`).
|
||||
- **Detection cache (persisted).** `_results` is mirrored to the project's
|
||||
`detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by
|
||||
**basename** so it survives moving the project). `cache.save_results`/`load_results`
|
||||
@@ -203,27 +208,43 @@ hvideotool/
|
||||
created lazily on first move), removing them from list/`_files`/cache. `_unique_dest`
|
||||
avoids clobbering (`foo.jpg` → `foo (1).jpg`). Use case: flag good frames into a
|
||||
training/example set while inspecting detections.
|
||||
- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via
|
||||
`build_restorer`) on the current frame's detections (computing them first if needed),
|
||||
**on a background job** (`_restore_current` builds an `fn` that detects-if-needed +
|
||||
restores in the worker; a `tick` caches freshly-computed detections, `done` stores
|
||||
`_restored[path]` + shows it). "Показать оригинал/
|
||||
результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
||||
`<stem>_restored.jpg` beside the frame. The baseline
|
||||
is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run
|
||||
**in-process** from the vendored `_deepmosaics/` code: it loads the BiSeNet locator +
|
||||
clean generator **once** (lazy, cached on the instance) and per frame reproduces their
|
||||
`cleanmosaic_img_server` (locate mosaic → run generator on the crop → feather back),
|
||||
~0.3 s/frame cached on CPU vs ~7 s when it spawned a subprocess. Use the **image**
|
||||
model `clean_youknow_resnet_9blocks.pth` — the video model (BVDNet) is rejected per
|
||||
frame (needs a neighbour). `should_cancel` (= `lambda: job.cancelled`) is polled so
|
||||
"■ Стоп" stops it; the engine raises `Cancelled`, which `Job.run` reports as a clean cancel.
|
||||
The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…),
|
||||
persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE:
|
||||
DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the
|
||||
clean weights) — our detections aren't passed to it. `_show` resets `_showing_restored`
|
||||
+ `_update_restore_actions`. To add another engine (e.g. LADA), implement
|
||||
`core/restore/base.Restorer` and register it in `restore/factory.build_restorer`.
|
||||
- **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it
|
||||
**locates the mosaic itself** — so restoration is **fully decoupled from detection**: no
|
||||
detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar
|
||||
action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on
|
||||
a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать
|
||||
оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
||||
`<stem>_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново",
|
||||
`_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every
|
||||
frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`,
|
||||
basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the
|
||||
per-frame engine **skips frames already in `restored/`** unless `force` (resume). The
|
||||
engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the
|
||||
vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy,
|
||||
cached on the instance):
|
||||
- `deepmosaics` (image, per-frame): reproduces `cleanmosaic_img_server` (locate mosaic
|
||||
→ run generator on the crop → feather back), ~0.3 s/frame cached on CPU. Image model
|
||||
`clean_youknow_resnet_9blocks.pth`.
|
||||
- `deepmosaics_video` (**temporal, BVDNet**): `DeepMosaicsVideoRestorer`, `.temporal=True`.
|
||||
Reproduces `cleanmosaic_video_fusion` — per target frame it feeds the net a window of
|
||||
`T=5` neighbour frames sampled at step `S=3` around it (`N=2` each side, clamped at the
|
||||
sequence edges) **plus its own previous output (recurrent)**, for temporal coherence.
|
||||
Because of that recurrence it must run a **contiguous, ordered range** — it implements
|
||||
`restore_sequence(count, get_frame, get_dets, emit, should_cancel)` (the batch run uses
|
||||
it; single-frame `restore` degrades to a window of the same frame). Needs the **video**
|
||||
weights `clean_youknow_video.pth` (+ `mosaic_position.pth` beside). `INPUT_SIZE=256`.
|
||||
|
||||
`restore_sequence` is on the `Restorer` ABC (default = independent per-frame loop);
|
||||
`_restore_all` dispatches on `restorer.temporal` (temporal → `restore_sequence` over the
|
||||
whole range; per-frame → resumable loop with skip-existing). `should_cancel`
|
||||
(= `lambda: job.cancelled`) is polled so "■ Стоп" stops it; engines raise `Cancelled`,
|
||||
which `Job.run` reports as a clean cancel. Engine + weights are set in `RestoreDialog`
|
||||
(Файл → Движок восстановления…) — the model dropdown shows image vs video weights per
|
||||
selected engine — persisted, and built lazily/cached in `_make_restorer`. NOTE:
|
||||
DeepMosaics locates mosaics itself (its `mosaic_position.pth`) — our detections aren't
|
||||
passed to it. To add another engine (e.g. LADA), implement `core/restore/base.Restorer`
|
||||
(set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
|
||||
`restore/factory.build_restorer`.
|
||||
- **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys
|
||||
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label`
|
||||
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
|
||||
@@ -255,20 +276,20 @@ hvideotool/
|
||||
- `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
|
||||
`core/detection/factory.build_detector` and talks to `core/` through the `Detector`
|
||||
interface and the `Detection`/`CensorType` types.
|
||||
- New detector kinds: implement `core/detection/base.Detector`, register the string in
|
||||
`core/detection/factory.build_detector`, and add it to `_DETECTORS` in
|
||||
`ui/main_window.py`.
|
||||
- Detection is YOLO-only and restoration is DeepMosaics-only. If you re-add an engine
|
||||
kind, implement `core/detection/base.Detector` / `core/restore/base.Restorer`, register
|
||||
the string in the respective `factory`, and add it to `DETECTORS`/`RESTORERS` in
|
||||
`config.py` (and `normalize_config`). There's no detector dropdown anymore — the toolbar
|
||||
just shows "Детектор: YOLO"; restoration engines are chosen in `RestoreDialog`.
|
||||
|
||||
## Commands
|
||||
|
||||
```powershell
|
||||
python -m venv .venv; .\.venv\Scripts\Activate.ps1
|
||||
pip install -e . # classic detector needs no torch/CUDA
|
||||
pip install -e ".[yolo]" # YOLO needs ultralytics; install torch separately (README)
|
||||
|
||||
python -m hvideotool # reopen the last project (or create/open one in-app)
|
||||
python -m hvideotool "C:\path\to\MyProject" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
|
||||
pip install -e ".[yolo]" # + install torch separately, see README
|
||||
python -m hvideotool "C:\path\to\MyProject" --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
```
|
||||
|
||||
No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a
|
||||
@@ -292,24 +313,22 @@ frame directly.
|
||||
a generic COCO model (e.g. the `yolo11n-seg.pt` in the repo root, which Ultralytics
|
||||
auto-downloads / is the training base), it detects people/objects and maps them to
|
||||
`CensorType.UNKNOWN` → purple boxes that look like noise. This was a real user trap.
|
||||
**Switching to yolo/combined without a model auto-picks one** via
|
||||
`MainWindow._auto_find_model()`: it scans `./models/**.pt` and matches only filenames
|
||||
containing `lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; it
|
||||
falls back to the "Модель…" file dialog only when nothing suitable is found.
|
||||
- **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low
|
||||
block-reconstruction residual + 2D gradient + contrast) fires on textured real
|
||||
footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing
|
||||
real mosaic after the `proc_max_dim=720` downscale softens block edges (measured:
|
||||
contrast/grad fall below `mosaic_contrast_min`/`mosaic_grad_min`). For real-video
|
||||
mosaic use `yolo`/`combined` + LADA. For anime there is no good public model.
|
||||
**A model is auto-picked on project open** via `MainWindow._ensure_model()` →
|
||||
`_auto_find_model()`: it scans `./models/**.pt` and matches only filenames containing
|
||||
`lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; otherwise the
|
||||
user picks via "Модель…" (`_choose_model`).
|
||||
- **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was
|
||||
noisy/approximate on real footage (false positives on skin/hair/fabric/JPEG; missed
|
||||
real mosaic after downscale) and the `combined` mode + cv2 `inpaint` baseline went with
|
||||
it. Detection is YOLO-only, restoration is DeepMosaics-only. `normalize_config` coerces
|
||||
any leftover `classic`/`combined`/`inpaint` in old settings/projects to `yolo`/`deepmosaics`.
|
||||
- **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic
|
||||
but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`),
|
||||
not tuning more classic thresholds.
|
||||
but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`).
|
||||
- **YOLO detector = LADA weights** ([HF `ladaapp/lada`](https://huggingface.co/ladaapp/lada)).
|
||||
YOLO **segmentation** model, classes `{0: mosaic_nsfw, 1: mosaic_sfw_head}` → both map
|
||||
to `CensorType.MOSAIC` (`_name_to_type` matches "mosaic" in the class name). Detects
|
||||
mosaic only; black bars / blur stay with classic. Weights + Ultralytics are AGPL-3.0
|
||||
(accepted). `yolo.py` lazy-imports `torch`/`ultralytics`.
|
||||
mosaic only. Weights + Ultralytics are AGPL-3.0 (accepted). `yolo.py` lazy-imports
|
||||
`torch`/`ultralytics`.
|
||||
- **No model weights in the repo.** Code must fail with a clear, actionable message
|
||||
when the model path is missing — not a raw stack trace (`factory._require_model`,
|
||||
`YoloDetector.__init__`).
|
||||
@@ -318,10 +337,11 @@ frame directly.
|
||||
- **CPU-only torch must not request CUDA.** A `+cpu` torch build raises "Torch not
|
||||
compiled with CUDA enabled" the moment something calls `.cuda()`. Both engines guard
|
||||
for this: `YoloDetector` picks `cuda` only when `torch.cuda.is_available()` (even an
|
||||
explicit `yolo_device="cuda"` is downgraded to cpu); `DeepMosaicsRestorer._ensure_loaded`
|
||||
forces `gpu_id="-1"` when CUDA is absent (its vendored `model_util.todevice` /
|
||||
`data.im2tensor` call `.cuda()` for any `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default).
|
||||
So a wrong/CPU-only torch falls back to CPU instead of crashing.
|
||||
explicit `yolo_device="cuda"` is downgraded to cpu); both `DeepMosaicsRestorer._ensure_loaded`
|
||||
and `DeepMosaicsVideoRestorer._ensure_loaded` force `gpu_id="-1"` when CUDA is absent (the
|
||||
vendored `model_util.todevice` / `data.im2tensor`/`to_tensor` call `.cuda()` for any
|
||||
`gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). So a wrong/CPU-only torch falls back to
|
||||
CPU instead of crashing (the temporal BVDNet engine is heavy on CPU, though).
|
||||
- **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`),
|
||||
otherwise it aliases a buffer that gets freed → garbage/crash.
|
||||
- **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images —
|
||||
|
||||
Reference in New Issue
Block a user