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:
@@ -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** (1–100, 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).
|
||||
|
||||
Reference in New Issue
Block a user