Enhance HVideoTool's detection and restoration capabilities: introduced per-model overlay threshold settings and cross-model non-maximum suppression (NMS) to improve detection accuracy. Updated configuration management to support these features, and refined the UI for better user experience. Documentation in CLAUDE.md has been updated to reflect these changes.

This commit is contained in:
Leonid Pershin
2026-06-08 04:06:50 +03:00
parent cabb4e3d3d
commit 8a366ed43d
11 changed files with 906 additions and 146 deletions
+94 -23
View File
@@ -138,6 +138,19 @@ hvideotool/
### How it works
- **Toolbar layout (`_build_toolbar`).** To avoid a long flat row spilling into Qt's
"⋯" overflow, related actions are **grouped into `QToolButton` dropdowns** (two helpers:
`_dropdown_button` = InstantPopup menu-only; `_split_button` = MenuButtonPopup, click runs
the primary action, the arrow opens related ones). Top-level groups: **"Проект ▾"**
(создать/открыть/из ролика/импортировать) · **"Детекторы: Модели (N) ▾"** (the model
picker, unchanged) · split **"Рассчитать кадр ▾"** (menu: детектировать все дозапуск/
заново) · split **"Расцензурить кадр ▾"** (menu: расцензурить все/найденное/заново,
движок…, открыть папку результатов, «Сохранить результат…» = `save_restored_action`) ·
**checkable "Показать расцензуренное"** (`toggle_restored_action`, key `R`, kept visible so
the original⇄restored state shows as a pressed button) · **"■ Стоп"** (kept visible,
reachable instantly) · a stretch spacer pushes **"Порог:"** to the right edge. The full
action list also lives in the **menu bar "Файл"** (`_build_menu`). When adding an action,
put it in the matching dropdown — don't add another flat top-level button.
- `MainWindow` holds the config, builds the detector lazily via `build_detector`
(cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps
`_results: dict[path -> list[Detection]]` as the detection cache.
@@ -148,11 +161,28 @@ hvideotool/
`.pt` into `models/yolo/<category>/`. On project open `_ensure_models` prunes vanished
paths and, if nothing is selected, default-ticks every discovered model. `build_detector`
builds one `YoloDetector` per selected model (tagged `label=category`) wrapped in a
`MultiYoloDetector` that concatenates their detections (no cross-model dedup). Each
`Detection` carries `label` (category) **and `model`** (the producing `.pt` stem, tagged in
`YoloDetector` — shown as its own "Модель" column in the detail table since a category folder
may hold several models); overlay colour + table group by `Detection.display`
(label, else the CensorType) via `OverlayConfig.colors` + a stable `palette` fallback.
`MultiYoloDetector`. Each `Detection` carries `label` (category) **and `model`** (the
producing `.pt` stem, tagged in `YoloDetector` — shown as its own "Модель" column in the
detail table since a category folder may hold several models); overlay colour + table
group by `Detection.display` (label, else the CensorType) via `OverlayConfig.colors` + a
stable `palette` fallback.
- **Cross-model NMS (optional).** By default `MultiYoloDetector` just **concatenates** all
models' detections (different categories are meant to coexist). The "Модели" menu has a
checkable **"Объединять пересечения (NMS)"** (`config.cross_model_nms` + `nms_iou`,
`_on_nms_toggled`): when on, `MultiYoloDetector(nms_iou=…)` runs a greedy category-agnostic
IoU NMS (`multi._nms`/`_iou`) that drops the lower-score box of any overlapping pair — kills
the duplicate rects you get when overlapping models fire (e.g. penis + cockAndBall). It
changes the detection result, so it's part of the in-memory detector identity
(`_make_detector` key) **and** the on-disk cache key — but `cache.make_key` adds `nms_iou`
**only when NMS is on**, so the default (off) key is unchanged and an existing cache stays
valid; turning NMS on yields a distinct key (recompute) without clobbering the non-NMS
cache. Toggling also `_invalidate_results` (drops the in-memory cache). Off = concatenate.
- **Per-model display thresholds (optional).** The toolbar "Порог" spin is the global overlay
threshold; the "Модели" menu **"Пороги по моделям…"** (`_edit_model_thresholds`) stores
per-model overrides in `config.model_thresholds` (keyed by `.pt` **stem**). `ImageView`
(`set_model_thresholds`/`_eff_threshold`) draws a detection only if its score clears its
model's override, else the global threshold. **Display-only** — does not change detection or
the "с цензурой" counts (a frame is a hit if it has *any* detection, threshold-independent).
- **Background jobs (`ui/workers.py`).** Detection and restoration are CPU-heavy and
would freeze the GUI, so they run on a `QThreadPool` thread via `Job` (a `QRunnable`
wrapping `fn(job)`); results return to the GUI through queued Qt signals
@@ -164,7 +194,9 @@ hvideotool/
GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables
`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`).
latency can't be hidden with `processEvents`). Progress messages get an **ETA** suffix
("· осталось ~Xм Yс") from `_eta_suffix(done, total)` using `_job_start` (set in
`_start_job`, `time.monotonic`) — average-rate estimate, blank at 0 %/100 %.
- **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
@@ -204,7 +236,10 @@ hvideotool/
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`
**basename** so it survives moving the project). Both `detections.json` and
`project.json` are written **atomically** (sibling `.tmp` + `os.replace`, see
`cache._atomic_write_text` and `Project.save`) — a crash mid-write can't corrupt/truncate
a large cache (29k entries) and lose all detection work. `cache.save_results`/`load_results`
take the cache file and the image `base_dir` (= `project.frames_dir`) separately, since
the cache lives at the project root, not next to the images. It's tagged with
the detector identity (`_results_key` = detector + model + conf/imgsz); on open,
@@ -229,13 +264,30 @@ hvideotool/
**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
a background job** (`_restore_current`; `done` stores `_restored[path]`, **auto-saves it
to the project's `restored/`** (so a single restore persists like the batch, not just in
memory), and shows it). **"Показать расцензуренное/оригинал" (`_toggle_restored`, key `R`)
is a GLOBAL view mode** (`_showing_restored`): when on, `_show` displays each frame's
restored version if one exists — loaded lazily from memory `_restored` **or disk
`restored/`** via `_restored_image_for` (so the whole batch result is browsable, not just
the last frame) — else falls back to the original; overlays are hidden on restored. The
mode persists across navigation (reset to off on project open); a batch/single restore
auto-switches it on. `_has_restored`/`_restored_disk_path` are the cheap (no-decode)
existence checks driving the toggle/save enabled-state. "Сохранить результат" additionally
exports `<stem>_restored.jpg` beside the frame (an explicit one-off export, via
`_restored_image_for`). "Открыть папку результатов" opens `restored/` in Explorer.
**Batch ("Расцензурить все" / "Все заново" /
"Расцензурить найденное", `_restore_all(force, only_detected)`)** mirrors `_detect_all`:
a single background job restores frames 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). **`only_detected`** ("Расцензурить найденное")
uses the YOLO detection cache to skip frames known clean: per-frame restores only the
frames with detections; the temporal engine restricts the run to the contiguous span
`[first hit … last hit]` (recurrence needs continuity). It's a separate, *faster* action
— NOT the default — because LADA misses some mosaic (esp. anime), so it can miss
censorship YOLO didn't flag; "Расцензурить все" stays the thorough option. Returns early
(status hint) if detection isn't computed or no frame has a detection. 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):
@@ -269,11 +321,19 @@ hvideotool/
(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 `[`/`]`,
`_step_hit` scans `_results` for the next non-empty frame). The slider and file list
are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all
navigation ultimately drives `file_list.setCurrentRow`. The scrubber is a custom
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a clickable `pos_label`
(a flat `QPushButton` "row / n"`_jump_to_frame`, a "go to frame N" `QInputDialog`
needed on 29k-frame projects where the scrubber is ~50 frames/px), and jump-to-detection
(◀ детекция / детекция ▶, keys `[`/`]`, `_step_hit` scans `_results` for the next
non-empty frame). The slider and file list are kept in sync via `_update_nav` guarded by
`_nav_sync` (avoids signal loops); all navigation ultimately drives
`file_list.setCurrentRow`. `_step` (◀ ▶) **skips rows hidden by the filter**.
- **File-list filter (`filter_combo`, left pane).** A combobox above the list shows only a
subset of the (possibly huge) frame list: Все · С цензурой · Чистые · Не рассчитано ·
Расцензуренные · Без расцензуривания. `_filter_mode` + `_row_matches_filter(path)` (over
`_results`/`_row_restored`); `_on_filter_changed` re-labels (each `_relabel_row` calls
`item.setHidden(...)`) and jumps off a now-hidden current row. Pure show/hide — doesn't
touch `_files`/cache. The scrubber is a custom
`MarkerSlider` (`ui/marker_slider.py`) that paints **two mark layers**: cyan ticks
(upper half) at frames with detections (`_refresh_marks` projects `_results`) and
**green ticks (lower half) at restored frames** (`_refresh_restored_marks` scans
@@ -281,8 +341,11 @@ hvideotool/
per-frame); per-pixel deduped so big folders stay cheap. Under the scrubber a
**progress summary** `stats_label` reads "Кадров: N · детектировано: D/N (с цензурой:
H) · расцензурено: R/N" (`_update_counts_label`, cheap counts; `_restored_count`
cached by `_refresh_restored_marks`). File-list rows are tinted too (`_tag_file`):
red = censorship found, green = checked & clean. Both reset on `_invalidate_results`.
cached by `_refresh_restored_marks`). File-list rows are labelled too via a single
`_relabel_row` (used by `_tag_file`/`_relabel_all`): tint red = censorship found, green =
checked & clean; a trailing **✓** marks frames with a restored version (`_row_restored`,
populated by `_refresh_restored_marks`). The ✓ is independent of detection — it survives
`_clear_results`. Detection tints reset on `_invalidate_results`.
- **Cancellation (cooperative).** A single "■ Стоп" toolbar action (Esc) cancels the
running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button +
the progress bar (`total=None` → indeterminate). For **background jobs** (detection,
@@ -322,7 +385,12 @@ python -m hvideotool # reopen the last project (or create/open
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
No formal test suite, but there's a **headless smoke test**: `scripts/smoke_test.py`
(run with `QT_QPA_PLATFORM=offscreen` + `PYTHONIOENCODING=utf-8`) covers the pure core
(atomic cache round-trip, cross-model NMS, list-filter predicate, ETA formatting,
extract-dialog options, per-project settings round-trip) and an offscreen `MainWindow`
build on a throwaway project — no torch/weights (detections are injected into `_results`).
Exits non-zero on failure; run it after touching core/UI plumbing. Ad-hoc check: build a
`MainWindow`, `Project.create(tmp)` + copy a few images into `frames/`,
`_open_project(project)`, drive `file_list.setCurrentRow(...)`, and read
`detail_table` / `detail_header`. Or run `build_detector(config).detect(...)` on a
@@ -386,4 +454,7 @@ frame directly.
the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is
~10× faster than every-frame; `-hwaccel` does NOT help (GPU transfer overhead). Use
ffmpeg/cv2, not PyAV, and keep it synchronous. Decoding every frame is the inherent
cost — the speed lever is decoding *fewer* frames (keyframes).
cost — the speed lever is decoding *fewer* frames (keyframes). `ExtractDialog.options()`
returns `(keyframes_only, step, max_dim, jpg_quality)`; **jpg_quality** (1100, default 92,
via `-q:v` `_quality_to_qscale` / cv2 `IMWRITE_JPEG_QUALITY`) trades quality for a bit of
encode speed + smaller files. Default sampling is **every frame** (step=1, not keyframes).