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
|
### 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`
|
- `MainWindow` holds the config, builds the detector lazily via `build_detector`
|
||||||
(cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps
|
(cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps
|
||||||
`_results: dict[path -> list[Detection]]` as the detection cache.
|
`_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
|
`.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`
|
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
|
builds one `YoloDetector` per selected model (tagged `label=category`) wrapped in a
|
||||||
`MultiYoloDetector` that concatenates their detections (no cross-model dedup). Each
|
`MultiYoloDetector`. Each `Detection` carries `label` (category) **and `model`** (the
|
||||||
`Detection` carries `label` (category) **and `model`** (the producing `.pt` stem, tagged in
|
producing `.pt` stem, tagged in `YoloDetector` — shown as its own "Модель" column in the
|
||||||
`YoloDetector` — shown as its own "Модель" column in the detail table since a category folder
|
detail table since a category folder may hold several models); overlay colour + table
|
||||||
may hold several models); overlay colour + table group by `Detection.display`
|
group by `Detection.display` (label, else the CensorType) via `OverlayConfig.colors` + a
|
||||||
(label, else the CensorType) via `OverlayConfig.colors` + a stable `palette` fallback.
|
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
|
- **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`
|
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
|
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
|
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
|
`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
|
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
|
- **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
|
"⚡ 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
|
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`).
|
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). 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
|
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 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,
|
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
|
**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
|
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
|
action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on
|
||||||
a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать
|
a background job** (`_restore_current`; `done` stores `_restored[path]`, **auto-saves it
|
||||||
оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
to the project's `restored/`** (so a single restore persists like the batch, not just in
|
||||||
`<stem>_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново",
|
memory), and shows it). **"Показать расцензуренное/оригинал" (`_toggle_restored`, key `R`)
|
||||||
`_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every
|
is a GLOBAL view mode** (`_showing_restored`): when on, `_show` displays each frame's
|
||||||
frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`,
|
restored version if one exists — loaded lazily from memory `_restored` **or disk
|
||||||
basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the
|
`restored/`** via `_restored_image_for` (so the whole batch result is browsable, not just
|
||||||
per-frame engine **skips frames already in `restored/`** unless `force` (resume). The
|
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
|
engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the
|
||||||
vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy,
|
vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy,
|
||||||
cached on the instance):
|
cached on the instance):
|
||||||
@@ -269,11 +321,19 @@ hvideotool/
|
|||||||
(set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
|
(set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
|
||||||
`restore/factory.build_restorer`.
|
`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 clickable `pos_label`
|
||||||
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
|
(a flat `QPushButton` "row / n" → `_jump_to_frame`, a "go to frame N" `QInputDialog` —
|
||||||
`_step_hit` scans `_results` for the next non-empty frame). The slider and file list
|
needed on 29k-frame projects where the scrubber is ~50 frames/px), and jump-to-detection
|
||||||
are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all
|
(◀ детекция / детекция ▶, keys `[`/`]`, `_step_hit` scans `_results` for the next
|
||||||
navigation ultimately drives `file_list.setCurrentRow`. The scrubber is a custom
|
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
|
`MarkerSlider` (`ui/marker_slider.py`) that paints **two mark layers**: cyan ticks
|
||||||
(upper half) at frames with detections (`_refresh_marks` projects `_results`) and
|
(upper half) at frames with detections (`_refresh_marks` projects `_results`) and
|
||||||
**green ticks (lower half) at restored frames** (`_refresh_restored_marks` scans
|
**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
|
per-frame); per-pixel deduped so big folders stay cheap. Under the scrubber a
|
||||||
**progress summary** `stats_label` reads "Кадров: N · детектировано: D/N (с цензурой:
|
**progress summary** `stats_label` reads "Кадров: N · детектировано: D/N (с цензурой:
|
||||||
H) · расцензурено: R/N" (`_update_counts_label`, cheap counts; `_restored_count`
|
H) · расцензурено: R/N" (`_update_counts_label`, cheap counts; `_restored_count`
|
||||||
cached by `_refresh_restored_marks`). File-list rows are tinted too (`_tag_file`):
|
cached by `_refresh_restored_marks`). File-list rows are labelled too via a single
|
||||||
red = censorship found, green = checked & clean. Both reset on `_invalidate_results`.
|
`_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
|
- **Cancellation (cooperative).** A single "■ Стоп" toolbar action (Esc) cancels the
|
||||||
running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button +
|
running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button +
|
||||||
the progress bar (`total=None` → indeterminate). For **background jobs** (detection,
|
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
|
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/`,
|
`MainWindow`, `Project.create(tmp)` + copy a few images into `frames/`,
|
||||||
`_open_project(project)`, drive `file_list.setCurrentRow(...)`, and read
|
`_open_project(project)`, drive `file_list.setCurrentRow(...)`, and read
|
||||||
`detail_table` / `detail_header`. Or run `build_detector(config).detect(...)` on a
|
`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
|
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
|
~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
|
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).
|
||||||
|
|||||||
@@ -56,6 +56,15 @@ class AppConfig:
|
|||||||
# selected model and merges results — see core/detection/multi.MultiYoloDetector.
|
# selected model and merges results — see core/detection/multi.MultiYoloDetector.
|
||||||
detector_models: list[str] = field(default_factory=list)
|
detector_models: list[str] = field(default_factory=list)
|
||||||
default_threshold: float = 0.20 # initial overlay confidence threshold
|
default_threshold: float = 0.20 # initial overlay confidence threshold
|
||||||
|
# Per-model overlay threshold overrides, keyed by the model file **stem** (e.g.
|
||||||
|
# "penis" -> 0.4). A detection from a model with no entry uses default_threshold.
|
||||||
|
# Display-only (filters what's drawn/counted as a hit), not the detection conf.
|
||||||
|
model_thresholds: dict[str, float] = field(default_factory=dict)
|
||||||
|
# Merge overlapping detections across models (greedy IoU NMS, keep higher score).
|
||||||
|
# Off by default — different categories are meant to coexist; on, it removes the
|
||||||
|
# duplicate boxes you get when overlapping models (e.g. penis + cockAndBall) fire.
|
||||||
|
cross_model_nms: bool = False
|
||||||
|
nms_iou: float = 0.6 # IoU above which two boxes are deemed duplicates
|
||||||
|
|
||||||
# --- restoration ("расцензурить") ---
|
# --- restoration ("расцензурить") ---
|
||||||
restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
|
restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ project folder.
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from .types import Detection
|
from .types import Detection
|
||||||
@@ -24,16 +25,43 @@ from .types import Detection
|
|||||||
_VERSION = 1
|
_VERSION = 1
|
||||||
|
|
||||||
|
|
||||||
def make_key(models: list[str], yolo_conf: float, yolo_imgsz: int) -> dict:
|
def _atomic_write_text(path: Path, text: str) -> None:
|
||||||
|
"""Write ``text`` to ``path`` crash-safely: write a sibling .tmp, then os.replace.
|
||||||
|
|
||||||
|
``os.replace`` is atomic on the same filesystem (incl. NTFS), so a crash mid-write
|
||||||
|
leaves the previous file intact instead of a truncated/corrupt one — important for
|
||||||
|
a large detections.json that holds tens of thousands of entries.
|
||||||
|
"""
|
||||||
|
tmp = path.with_name(path.name + ".tmp")
|
||||||
|
try:
|
||||||
|
tmp.write_text(text, encoding="utf-8")
|
||||||
|
os.replace(tmp, path)
|
||||||
|
finally:
|
||||||
|
if tmp.exists():
|
||||||
|
try:
|
||||||
|
tmp.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def make_key(
|
||||||
|
models: list[str], yolo_conf: float, yolo_imgsz: int, nms_iou: float | None = None
|
||||||
|
) -> dict:
|
||||||
"""Identity of the detector set that produced a cache; cache is only valid for a match.
|
"""Identity of the detector set that produced a cache; cache is only valid for a match.
|
||||||
|
|
||||||
Keyed by the (sorted) model **basenames** so it's portable across machines/paths.
|
Keyed by the (sorted) model **basenames** so it's portable across machines/paths.
|
||||||
|
``nms_iou`` is only added to the key when cross-model NMS is enabled — so the default
|
||||||
|
(NMS off) key is unchanged and existing caches stay valid; turning NMS on yields a
|
||||||
|
distinct key (its merged results differ) without invalidating the non-NMS cache.
|
||||||
"""
|
"""
|
||||||
return {
|
key = {
|
||||||
"models": sorted(Path(m).name for m in models),
|
"models": sorted(Path(m).name for m in models),
|
||||||
"yolo_conf": round(float(yolo_conf), 4),
|
"yolo_conf": round(float(yolo_conf), 4),
|
||||||
"yolo_imgsz": int(yolo_imgsz),
|
"yolo_imgsz": int(yolo_imgsz),
|
||||||
}
|
}
|
||||||
|
if nms_iou is not None:
|
||||||
|
key["nms_iou"] = round(float(nms_iou), 4)
|
||||||
|
return key
|
||||||
|
|
||||||
|
|
||||||
def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]]) -> bool:
|
def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]]) -> bool:
|
||||||
@@ -48,9 +76,7 @@ def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]
|
|||||||
}
|
}
|
||||||
try:
|
try:
|
||||||
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
cache_file.parent.mkdir(parents=True, exist_ok=True)
|
||||||
cache_file.write_text(
|
_atomic_write_text(cache_file, json.dumps(payload, ensure_ascii=False))
|
||||||
json.dumps(payload, ensure_ascii=False), encoding="utf-8"
|
|
||||||
)
|
|
||||||
return True
|
return True
|
||||||
except OSError:
|
except OSError:
|
||||||
return False
|
return False
|
||||||
|
|||||||
@@ -31,6 +31,8 @@ def build_detector(config: AppConfig) -> Detector:
|
|||||||
from .multi import MultiYoloDetector
|
from .multi import MultiYoloDetector
|
||||||
from .yolo import YoloDetector
|
from .yolo import YoloDetector
|
||||||
|
|
||||||
return MultiYoloDetector([
|
nms_iou = config.nms_iou if config.cross_model_nms else None
|
||||||
YoloDetector(m, config.detection, label=category_of(m)) for m in models
|
return MultiYoloDetector(
|
||||||
])
|
[YoloDetector(m, config.detection, label=category_of(m)) for m in models],
|
||||||
|
nms_iou=nms_iou,
|
||||||
|
)
|
||||||
|
|||||||
@@ -15,10 +15,13 @@ from .types import Detection
|
|||||||
|
|
||||||
|
|
||||||
class MultiYoloDetector(Detector):
|
class MultiYoloDetector(Detector):
|
||||||
def __init__(self, detectors: list[Detector]) -> None:
|
def __init__(self, detectors: list[Detector], nms_iou: float | None = None) -> None:
|
||||||
if not detectors:
|
if not detectors:
|
||||||
raise ValueError("MultiYoloDetector requires at least one detector")
|
raise ValueError("MultiYoloDetector requires at least one detector")
|
||||||
self._detectors = detectors
|
self._detectors = detectors
|
||||||
|
# When set, overlapping detections (across all models, regardless of category)
|
||||||
|
# are merged by greedy IoU NMS — the higher-score box wins. None = keep all.
|
||||||
|
self._nms_iou = nms_iou
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def name(self) -> str:
|
def name(self) -> str:
|
||||||
@@ -28,4 +31,33 @@ class MultiYoloDetector(Detector):
|
|||||||
out: list[Detection] = []
|
out: list[Detection] = []
|
||||||
for d in self._detectors:
|
for d in self._detectors:
|
||||||
out.extend(d.detect(frame))
|
out.extend(d.detect(frame))
|
||||||
|
if self._nms_iou is not None:
|
||||||
|
out = _nms(out, self._nms_iou)
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
|
||||||
|
"""Intersection-over-union of two (x, y, w, h) boxes."""
|
||||||
|
ax, ay, aw, ah = a
|
||||||
|
bx, by, bw, bh = b
|
||||||
|
ix1, iy1 = max(ax, bx), max(ay, by)
|
||||||
|
ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
|
||||||
|
iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
|
||||||
|
inter = iw * ih
|
||||||
|
if inter == 0:
|
||||||
|
return 0.0
|
||||||
|
union = aw * ah + bw * bh - inter
|
||||||
|
return inter / union if union > 0 else 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def _nms(dets: list[Detection], iou_thresh: float) -> list[Detection]:
|
||||||
|
"""Greedy non-maximum suppression across all detections (category-agnostic).
|
||||||
|
|
||||||
|
Highest score first; a box is dropped if it overlaps an already-kept box by more
|
||||||
|
than ``iou_thresh``. Used to remove duplicate boxes from overlapping models.
|
||||||
|
"""
|
||||||
|
kept: list[Detection] = []
|
||||||
|
for d in sorted(dets, key=lambda x: x.score, reverse=True):
|
||||||
|
if all(_iou(d.bbox, k.bbox) <= iou_thresh for k in kept):
|
||||||
|
kept.append(d)
|
||||||
|
return kept
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ threshold, restore engine). Global ``settings.json`` only seeds the defaults for
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -38,6 +39,9 @@ _SETTING_KEYS = (
|
|||||||
"detector",
|
"detector",
|
||||||
"detector_models",
|
"detector_models",
|
||||||
"default_threshold",
|
"default_threshold",
|
||||||
|
"model_thresholds",
|
||||||
|
"cross_model_nms",
|
||||||
|
"nms_iou",
|
||||||
"restorer",
|
"restorer",
|
||||||
"dm_dir",
|
"dm_dir",
|
||||||
"dm_model",
|
"dm_model",
|
||||||
@@ -139,9 +143,19 @@ class Project:
|
|||||||
"settings": self.settings,
|
"settings": self.settings,
|
||||||
}
|
}
|
||||||
self.root.mkdir(parents=True, exist_ok=True)
|
self.root.mkdir(parents=True, exist_ok=True)
|
||||||
self.project_file.write_text(
|
# Crash-safe write (sibling .tmp + atomic os.replace): never truncate a good
|
||||||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
# project.json if the app/PC dies mid-save.
|
||||||
)
|
text = json.dumps(payload, ensure_ascii=False, indent=2)
|
||||||
|
tmp = self.project_file.with_name(PROJECT_FILE + ".tmp")
|
||||||
|
try:
|
||||||
|
tmp.write_text(text, encoding="utf-8")
|
||||||
|
os.replace(tmp, self.project_file)
|
||||||
|
finally:
|
||||||
|
if tmp.exists():
|
||||||
|
try:
|
||||||
|
tmp.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
# ----------------------------------------------------- settings <-> config
|
# ----------------------------------------------------- settings <-> config
|
||||||
def apply_to_config(self, cfg: AppConfig) -> None:
|
def apply_to_config(self, cfg: AppConfig) -> None:
|
||||||
|
|||||||
@@ -37,6 +37,14 @@ def apply(config: AppConfig) -> None:
|
|||||||
config.detector_models = [str(m) for m in data["detector_models"]]
|
config.detector_models = [str(m) for m in data["detector_models"]]
|
||||||
if "threshold" in data:
|
if "threshold" in data:
|
||||||
config.default_threshold = float(data["threshold"])
|
config.default_threshold = float(data["threshold"])
|
||||||
|
if isinstance(data.get("model_thresholds"), dict):
|
||||||
|
config.model_thresholds = {
|
||||||
|
str(k): float(v) for k, v in data["model_thresholds"].items()
|
||||||
|
}
|
||||||
|
if "cross_model_nms" in data:
|
||||||
|
config.cross_model_nms = bool(data["cross_model_nms"])
|
||||||
|
if "nms_iou" in data:
|
||||||
|
config.nms_iou = float(data["nms_iou"])
|
||||||
if data.get("restorer"):
|
if data.get("restorer"):
|
||||||
config.restorer = data["restorer"]
|
config.restorer = data["restorer"]
|
||||||
for key in ("dm_dir", "dm_model", "dm_gpu"):
|
for key in ("dm_dir", "dm_model", "dm_gpu"):
|
||||||
@@ -53,6 +61,9 @@ def save(config: AppConfig) -> None:
|
|||||||
detector=config.detector,
|
detector=config.detector,
|
||||||
detector_models=list(config.detector_models),
|
detector_models=list(config.detector_models),
|
||||||
threshold=config.default_threshold,
|
threshold=config.default_threshold,
|
||||||
|
model_thresholds=dict(config.model_thresholds),
|
||||||
|
cross_model_nms=config.cross_model_nms,
|
||||||
|
nms_iou=config.nms_iou,
|
||||||
restorer=config.restorer,
|
restorer=config.restorer,
|
||||||
dm_dir=config.dm_dir,
|
dm_dir=config.dm_dir,
|
||||||
dm_model=config.dm_model,
|
dm_model=config.dm_model,
|
||||||
|
|||||||
@@ -41,13 +41,23 @@ class ExtractDialog(QDialog):
|
|||||||
self.max_dim.setValue(0)
|
self.max_dim.setValue(0)
|
||||||
self.max_dim.setSpecialValueText("оригинал")
|
self.max_dim.setSpecialValueText("оригинал")
|
||||||
|
|
||||||
|
self.quality = QSpinBox()
|
||||||
|
self.quality.setRange(1, 100)
|
||||||
|
self.quality.setValue(92)
|
||||||
|
self.quality.setToolTip(
|
||||||
|
"Качество JPEG (1–100). Ниже = быстрее кодирование и меньше файлы,\n"
|
||||||
|
"но больше артефактов. 92 — хороший баланс."
|
||||||
|
)
|
||||||
|
|
||||||
form = QFormLayout(self)
|
form = QFormLayout(self)
|
||||||
form.addRow("Режим:", self.mode)
|
form.addRow("Режим:", self.mode)
|
||||||
form.addRow("Брать каждый N-й кадр:", self.step)
|
form.addRow("Брать каждый N-й кадр:", self.step)
|
||||||
form.addRow("Макс. сторона, px:", self.max_dim)
|
form.addRow("Макс. сторона, px:", self.max_dim)
|
||||||
|
form.addRow("Качество JPEG:", self.quality)
|
||||||
hint = QLabel(
|
hint = QLabel(
|
||||||
"Ключевые кадры — в разы быстрее (декодируются только I-кадры),\n"
|
"Ключевые кадры — в разы быстрее (декодируются только I-кадры),\n"
|
||||||
"но реже по времени. Даунскейл уменьшает файлы и нагрузку на диск."
|
"но реже по времени. Даунскейл и меньшее качество уменьшают файлы\n"
|
||||||
|
"и нагрузку на диск (качество чуть ускоряет кодирование)."
|
||||||
)
|
)
|
||||||
hint.setWordWrap(True)
|
hint.setWordWrap(True)
|
||||||
form.addRow(hint)
|
form.addRow(hint)
|
||||||
@@ -61,6 +71,11 @@ class ExtractDialog(QDialog):
|
|||||||
def _sync(self) -> None:
|
def _sync(self) -> None:
|
||||||
self.step.setEnabled(not self.mode.currentData())
|
self.step.setEnabled(not self.mode.currentData())
|
||||||
|
|
||||||
def options(self) -> tuple[bool, int, int]:
|
def options(self) -> tuple[bool, int, int, int]:
|
||||||
"""Return (keyframes_only, step, max_dim)."""
|
"""Return (keyframes_only, step, max_dim, jpg_quality)."""
|
||||||
return bool(self.mode.currentData()), self.step.value(), self.max_dim.value()
|
return (
|
||||||
|
bool(self.mode.currentData()),
|
||||||
|
self.step.value(),
|
||||||
|
self.max_dim.value(),
|
||||||
|
self.quality.value(),
|
||||||
|
)
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ class ImageView(QWidget):
|
|||||||
self._qimage: QImage | None = None
|
self._qimage: QImage | None = None
|
||||||
self._dets: list[Detection] = []
|
self._dets: list[Detection] = []
|
||||||
self._threshold = 0.0
|
self._threshold = 0.0
|
||||||
|
self._model_thresholds: dict[str, float] = {} # model stem -> override threshold
|
||||||
self._highlight: int | None = None
|
self._highlight: int | None = None
|
||||||
self.setMinimumSize(480, 360)
|
self.setMinimumSize(480, 360)
|
||||||
|
|
||||||
@@ -47,6 +48,15 @@ class ImageView(QWidget):
|
|||||||
self._threshold = threshold
|
self._threshold = threshold
|
||||||
self.update()
|
self.update()
|
||||||
|
|
||||||
|
def set_model_thresholds(self, thresholds: dict[str, float]) -> None:
|
||||||
|
"""Per-model overlay threshold overrides (keyed by model stem). Empty = none."""
|
||||||
|
self._model_thresholds = dict(thresholds)
|
||||||
|
self.update()
|
||||||
|
|
||||||
|
def _eff_threshold(self, d: Detection) -> float:
|
||||||
|
"""The threshold this detection must clear — its model's override, else global."""
|
||||||
|
return self._model_thresholds.get(d.model, self._threshold)
|
||||||
|
|
||||||
def set_highlight(self, index: int | None) -> None:
|
def set_highlight(self, index: int | None) -> None:
|
||||||
self._highlight = index
|
self._highlight = index
|
||||||
self.update()
|
self.update()
|
||||||
@@ -82,8 +92,9 @@ class ImageView(QWidget):
|
|||||||
painter.setRenderHint(QPainter.Antialiasing, True)
|
painter.setRenderHint(QPainter.Antialiasing, True)
|
||||||
for i, d in enumerate(self._dets):
|
for i, d in enumerate(self._dets):
|
||||||
highlighted = i == self._highlight
|
highlighted = i == self._highlight
|
||||||
# A highlighted detection is always drawn; others respect the threshold.
|
# A highlighted detection is always drawn; others respect the threshold
|
||||||
if not highlighted and d.score < self._threshold:
|
# (per-model override if any, else the global one).
|
||||||
|
if not highlighted and d.score < self._eff_threshold(d):
|
||||||
continue
|
continue
|
||||||
dim = self._highlight is not None and not highlighted
|
dim = self._highlight is not None and not highlighted
|
||||||
self._draw_detection(painter, d, ox, oy, scale, highlighted, dim)
|
self._draw_detection(painter, d, ox, oy, scale, highlighted, dim)
|
||||||
|
|||||||
+422
-100
@@ -25,6 +25,7 @@ from __future__ import annotations
|
|||||||
import contextlib
|
import contextlib
|
||||||
import os
|
import os
|
||||||
import shutil
|
import shutil
|
||||||
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from PySide6.QtCore import Qt, QThreadPool
|
from PySide6.QtCore import Qt, QThreadPool
|
||||||
@@ -32,9 +33,12 @@ from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcu
|
|||||||
from PySide6.QtWidgets import (
|
from PySide6.QtWidgets import (
|
||||||
QAbstractItemView,
|
QAbstractItemView,
|
||||||
QApplication,
|
QApplication,
|
||||||
|
QComboBox,
|
||||||
QDialog,
|
QDialog,
|
||||||
|
QDialogButtonBox,
|
||||||
QDoubleSpinBox,
|
QDoubleSpinBox,
|
||||||
QFileDialog,
|
QFileDialog,
|
||||||
|
QFormLayout,
|
||||||
QHBoxLayout,
|
QHBoxLayout,
|
||||||
QInputDialog,
|
QInputDialog,
|
||||||
QLabel,
|
QLabel,
|
||||||
@@ -46,6 +50,7 @@ from PySide6.QtWidgets import (
|
|||||||
QPlainTextEdit,
|
QPlainTextEdit,
|
||||||
QProgressBar,
|
QProgressBar,
|
||||||
QPushButton,
|
QPushButton,
|
||||||
|
QSizePolicy,
|
||||||
QSplitter,
|
QSplitter,
|
||||||
QTableWidget,
|
QTableWidget,
|
||||||
QTableWidgetItem,
|
QTableWidgetItem,
|
||||||
@@ -89,7 +94,9 @@ class MainWindow(QMainWindow):
|
|||||||
self._restorer_key = None
|
self._restorer_key = None
|
||||||
self._restored: dict[str, object] = {} # path -> restored image (BGR ndarray)
|
self._restored: dict[str, object] = {} # path -> restored image (BGR ndarray)
|
||||||
self._restored_count = 0 # frames with output in restored/ (for the progress summary)
|
self._restored_count = 0 # frames with output in restored/ (for the progress summary)
|
||||||
|
self._row_restored: set[str] = set() # frame paths that have a restored version (row ✓)
|
||||||
self._showing_restored = False
|
self._showing_restored = False
|
||||||
|
self._filter_mode = "all" # file-list filter (see filter_combo)
|
||||||
self._nav_sync = False # guard against slider<->list signal loops
|
self._nav_sync = False # guard against slider<->list signal loops
|
||||||
self._busy = False # a long operation is running
|
self._busy = False # a long operation is running
|
||||||
self._cancel = False # the user asked to stop it
|
self._cancel = False # the user asked to stop it
|
||||||
@@ -126,6 +133,8 @@ class MainWindow(QMainWindow):
|
|||||||
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(False))
|
||||||
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||||||
|
file_menu.addAction("Расцензурить найденное (по детекции)", lambda: self._restore_all(only_detected=True))
|
||||||
|
file_menu.addAction("Открыть папку результатов", self._open_restored_dir)
|
||||||
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()
|
||||||
@@ -134,14 +143,19 @@ class MainWindow(QMainWindow):
|
|||||||
def _build_toolbar(self) -> None:
|
def _build_toolbar(self) -> None:
|
||||||
tb = self.addToolBar("Главная")
|
tb = self.addToolBar("Главная")
|
||||||
tb.setMovable(False)
|
tb.setMovable(False)
|
||||||
|
tb.setToolButtonStyle(Qt.ToolButtonTextOnly)
|
||||||
|
|
||||||
tb.addAction(QAction("Создать проект…", self, triggered=self._create_project))
|
# --- Проект: rarely-touched session actions collapsed into one dropdown.
|
||||||
tb.addAction(QAction("Открыть проект…", self, triggered=self._open_project_dialog))
|
project_btn = self._dropdown_button("Проект ▾", "Действия с проектом")
|
||||||
from_video = QAction("Создать из ролика…", self, triggered=self._create_from_video)
|
m = project_btn.menu()
|
||||||
from_video.setToolTip("Разложить видео на кадры в новый проект и открыть его")
|
m.addAction("Создать проект…", self._create_project)
|
||||||
tb.addAction(from_video)
|
m.addAction("Открыть проект…", self._open_project_dialog)
|
||||||
|
m.addAction("Создать из ролика…", self._create_from_video)
|
||||||
|
m.addAction("Импортировать папку как проект…", self._import_folder_as_project)
|
||||||
|
tb.addWidget(project_btn)
|
||||||
tb.addSeparator()
|
tb.addSeparator()
|
||||||
|
|
||||||
|
# --- Детекторы: active YOLO models picker (unchanged).
|
||||||
tb.addWidget(QLabel(" Детекторы: "))
|
tb.addWidget(QLabel(" Детекторы: "))
|
||||||
self._models_menu = QMenu(self)
|
self._models_menu = QMenu(self)
|
||||||
self.models_button = QToolButton()
|
self.models_button = QToolButton()
|
||||||
@@ -150,41 +164,62 @@ class MainWindow(QMainWindow):
|
|||||||
self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)")
|
self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)")
|
||||||
tb.addWidget(self.models_button)
|
tb.addWidget(self.models_button)
|
||||||
self._rebuild_models_menu()
|
self._rebuild_models_menu()
|
||||||
|
|
||||||
tb.addSeparator()
|
tb.addSeparator()
|
||||||
calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
|
|
||||||
calc.setToolTip("Запустить детектор на выбранном кадре (Space / двойной клик по файлу)")
|
|
||||||
tb.addAction(calc)
|
|
||||||
detect_all = QAction("Детектировать все", self, triggered=lambda: self._detect_all(False))
|
|
||||||
detect_all.setToolTip("Рассчитать все ещё не посчитанные кадры (дозапуск; кэш сохраняется)")
|
|
||||||
tb.addAction(detect_all)
|
|
||||||
regen = QAction("Все заново", self, triggered=lambda: self._detect_all(True))
|
|
||||||
regen.setToolTip("Очистить кэш детекций и пересчитать всю папку заново")
|
|
||||||
tb.addAction(regen)
|
|
||||||
|
|
||||||
|
# --- Детекция: primary "Рассчитать кадр" + the bulk runs in its dropdown.
|
||||||
|
detect_btn = self._split_button(
|
||||||
|
"Рассчитать кадр", self._recompute_current,
|
||||||
|
"Запустить детектор на выбранном кадре (Space / двойной клик по файлу)",
|
||||||
|
)
|
||||||
|
dm = detect_btn.menu()
|
||||||
|
dm.addAction("Детектировать все (дозапуск)", lambda: self._detect_all(False))
|
||||||
|
dm.addAction("Детектировать все заново", lambda: self._detect_all(True))
|
||||||
|
tb.addWidget(detect_btn)
|
||||||
|
|
||||||
|
# --- Расцензурить: primary "Расцензурить кадр" + the bulk/engine actions.
|
||||||
|
restore_btn = self._split_button(
|
||||||
|
"Расцензурить кадр", self._restore_current,
|
||||||
|
"Восстановить мозаику на текущем кадре (результат сохраняется в restored/)",
|
||||||
|
)
|
||||||
|
rm = restore_btn.menu()
|
||||||
|
rm.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
|
||||||
|
hits = rm.addAction("Расцензурить найденное (по детекции)", lambda: self._restore_all(only_detected=True))
|
||||||
|
hits.setToolTip(
|
||||||
|
"Расцензурить только кадры с детекцией (быстро, пропускает чистые).\n"
|
||||||
|
"ВНИМАНИЕ: YOLO ловит не всю мозаику — может пропустить."
|
||||||
|
)
|
||||||
|
rm.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||||||
|
rm.addSeparator()
|
||||||
|
rm.addAction("Движок восстановления…", self._open_restore_settings)
|
||||||
|
rm.addAction("Открыть папку результатов", self._open_restored_dir)
|
||||||
|
rm.addSeparator()
|
||||||
|
self.save_restored_action = QAction("Сохранить результат…", self, triggered=self._save_restored)
|
||||||
|
self.save_restored_action.setEnabled(False)
|
||||||
|
self.save_restored_action.setToolTip("Экспортировать <имя>_restored.jpg рядом с кадром")
|
||||||
|
rm.addAction(self.save_restored_action)
|
||||||
|
tb.addWidget(restore_btn)
|
||||||
|
|
||||||
|
# --- View toggle: kept visible (checkable) so the result/original state is obvious.
|
||||||
|
self.toggle_restored_action = QAction("Показать расцензуренное", self, triggered=self._toggle_restored)
|
||||||
|
self.toggle_restored_action.setCheckable(True)
|
||||||
|
self.toggle_restored_action.setEnabled(False)
|
||||||
|
self.toggle_restored_action.setShortcut("R")
|
||||||
|
self.toggle_restored_action.setToolTip(
|
||||||
|
"Переключить просмотр оригинал ⇄ расцензуренное для всего проекта (R)"
|
||||||
|
)
|
||||||
|
tb.addAction(self.toggle_restored_action)
|
||||||
|
tb.addSeparator()
|
||||||
|
|
||||||
|
# --- Stop: kept visible — must be reachable instantly during a long run.
|
||||||
self.stop_action = QAction("■ Стоп", self, triggered=self._request_cancel)
|
self.stop_action = QAction("■ Стоп", self, triggered=self._request_cancel)
|
||||||
self.stop_action.setToolTip("Отменить текущую операцию (Esc)")
|
self.stop_action.setToolTip("Отменить текущую операцию (Esc)")
|
||||||
self.stop_action.setEnabled(False)
|
self.stop_action.setEnabled(False)
|
||||||
tb.addAction(self.stop_action)
|
tb.addAction(self.stop_action)
|
||||||
|
|
||||||
tb.addSeparator()
|
# Push the threshold control to the right edge.
|
||||||
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
|
spacer = QWidget()
|
||||||
restore.setToolTip("Восстановить найденные области на текущем кадре")
|
spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)
|
||||||
tb.addAction(restore)
|
tb.addWidget(spacer)
|
||||||
restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False))
|
|
||||||
restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)")
|
|
||||||
tb.addAction(restore_all)
|
|
||||||
restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True))
|
|
||||||
restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)")
|
|
||||||
tb.addAction(restore_regen)
|
|
||||||
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
|
|
||||||
self.toggle_restored_action.setEnabled(False)
|
|
||||||
tb.addAction(self.toggle_restored_action)
|
|
||||||
self.save_restored_action = QAction("Сохранить результат", self, triggered=self._save_restored)
|
|
||||||
self.save_restored_action.setEnabled(False)
|
|
||||||
tb.addAction(self.save_restored_action)
|
|
||||||
|
|
||||||
tb.addSeparator()
|
|
||||||
tb.addWidget(QLabel(" Порог: "))
|
tb.addWidget(QLabel(" Порог: "))
|
||||||
self.threshold_spin = QDoubleSpinBox()
|
self.threshold_spin = QDoubleSpinBox()
|
||||||
self.threshold_spin.setRange(0.0, 1.0)
|
self.threshold_spin.setRange(0.0, 1.0)
|
||||||
@@ -193,12 +228,47 @@ class MainWindow(QMainWindow):
|
|||||||
self.threshold_spin.valueChanged.connect(self._on_threshold_changed)
|
self.threshold_spin.valueChanged.connect(self._on_threshold_changed)
|
||||||
tb.addWidget(self.threshold_spin)
|
tb.addWidget(self.threshold_spin)
|
||||||
|
|
||||||
|
def _dropdown_button(self, text: str, tooltip: str) -> QToolButton:
|
||||||
|
"""A toolbar button that just opens a menu (no default action)."""
|
||||||
|
btn = QToolButton()
|
||||||
|
btn.setText(text)
|
||||||
|
btn.setToolTip(tooltip)
|
||||||
|
btn.setPopupMode(QToolButton.InstantPopup)
|
||||||
|
btn.setMenu(QMenu(btn))
|
||||||
|
return btn
|
||||||
|
|
||||||
|
def _split_button(self, text: str, slot, tooltip: str) -> QToolButton:
|
||||||
|
"""A split button: click runs ``slot``; the arrow opens a menu of related actions."""
|
||||||
|
btn = QToolButton()
|
||||||
|
btn.setText(text)
|
||||||
|
btn.setToolTip(tooltip)
|
||||||
|
btn.setPopupMode(QToolButton.MenuButtonPopup)
|
||||||
|
action = QAction(text, btn, triggered=slot)
|
||||||
|
action.setToolTip(tooltip)
|
||||||
|
btn.setDefaultAction(action)
|
||||||
|
btn.setMenu(QMenu(btn))
|
||||||
|
return btn
|
||||||
|
|
||||||
def _build_central(self) -> None:
|
def _build_central(self) -> None:
|
||||||
self.file_list = QListWidget()
|
self.file_list = QListWidget()
|
||||||
self.file_list.setSelectionMode(QAbstractItemView.ExtendedSelection) # multi-select for moves
|
self.file_list.setSelectionMode(QAbstractItemView.ExtendedSelection) # multi-select for moves
|
||||||
self.file_list.currentItemChanged.connect(self._on_file_selected)
|
self.file_list.currentItemChanged.connect(self._on_file_selected)
|
||||||
self.file_list.itemDoubleClicked.connect(self._on_file_activated)
|
self.file_list.itemDoubleClicked.connect(self._on_file_activated)
|
||||||
|
|
||||||
|
# Filter the (possibly huge) list down to the frames you care about.
|
||||||
|
self.filter_combo = QComboBox()
|
||||||
|
for label, mode in (
|
||||||
|
("Все кадры", "all"),
|
||||||
|
("С цензурой", "hits"),
|
||||||
|
("Чистые", "clean"),
|
||||||
|
("Не рассчитано", "uncomputed"),
|
||||||
|
("Расцензуренные", "restored"),
|
||||||
|
("Без расцензуривания", "unrestored"),
|
||||||
|
):
|
||||||
|
self.filter_combo.addItem(label, mode)
|
||||||
|
self.filter_combo.setToolTip("Показывать только кадры выбранной категории")
|
||||||
|
self.filter_combo.currentIndexChanged.connect(self._on_filter_changed)
|
||||||
|
|
||||||
# One default collection ("Избранное"); the button acts on the list selection.
|
# One default collection ("Избранное"); the button acts on the list selection.
|
||||||
move_btn = QPushButton("★ В избранное")
|
move_btn = QPushButton("★ В избранное")
|
||||||
move_btn.setToolTip("Переместить выбранные кадры в избранное проекта (Ctrl+M)")
|
move_btn.setToolTip("Переместить выбранные кадры в избранное проекта (Ctrl+M)")
|
||||||
@@ -208,6 +278,7 @@ class MainWindow(QMainWindow):
|
|||||||
left_layout = QVBoxLayout(left)
|
left_layout = QVBoxLayout(left)
|
||||||
left_layout.setContentsMargins(4, 4, 4, 4)
|
left_layout.setContentsMargins(4, 4, 4, 4)
|
||||||
left_layout.setSpacing(4)
|
left_layout.setSpacing(4)
|
||||||
|
left_layout.addWidget(self.filter_combo)
|
||||||
left_layout.addWidget(self.file_list, 1)
|
left_layout.addWidget(self.file_list, 1)
|
||||||
left_layout.addWidget(move_btn)
|
left_layout.addWidget(move_btn)
|
||||||
|
|
||||||
@@ -269,9 +340,14 @@ class MainWindow(QMainWindow):
|
|||||||
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
|
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
|
||||||
self.frame_slider.valueChanged.connect(self._on_slider)
|
self.frame_slider.valueChanged.connect(self._on_slider)
|
||||||
|
|
||||||
self.pos_label = QLabel("0 / 0")
|
# Clickable position readout — opens "go to frame N" (handy on huge sequences
|
||||||
|
# where the scrubber is too coarse, ~50 frames/px on 29k).
|
||||||
|
self.pos_label = QPushButton("0 / 0")
|
||||||
|
self.pos_label.setFlat(True)
|
||||||
self.pos_label.setMinimumWidth(90)
|
self.pos_label.setMinimumWidth(90)
|
||||||
self.pos_label.setAlignment(Qt.AlignCenter)
|
self.pos_label.setCursor(Qt.PointingHandCursor)
|
||||||
|
self.pos_label.setToolTip("Перейти к кадру по номеру")
|
||||||
|
self.pos_label.clicked.connect(self._jump_to_frame)
|
||||||
|
|
||||||
self.prev_hit_btn = QPushButton("◀ детекция")
|
self.prev_hit_btn = QPushButton("◀ детекция")
|
||||||
self.prev_hit_btn.setToolTip("Предыдущий кадр с детекцией ([)")
|
self.prev_hit_btn.setToolTip("Предыдущий кадр с детекцией ([)")
|
||||||
@@ -297,8 +373,13 @@ class MainWindow(QMainWindow):
|
|||||||
n = self.file_list.count()
|
n = self.file_list.count()
|
||||||
if n == 0:
|
if n == 0:
|
||||||
return
|
return
|
||||||
row = max(0, min(n - 1, self.file_list.currentRow() + delta))
|
# Skip rows hidden by the filter so prev/next walk only the visible frames.
|
||||||
self.file_list.setCurrentRow(row)
|
i = self.file_list.currentRow() + delta
|
||||||
|
while 0 <= i < n:
|
||||||
|
if not self.file_list.item(i).isHidden():
|
||||||
|
self.file_list.setCurrentRow(i)
|
||||||
|
return
|
||||||
|
i += delta
|
||||||
|
|
||||||
def _step_hit(self, direction: int) -> None:
|
def _step_hit(self, direction: int) -> None:
|
||||||
"""Jump to the nearest frame (in `direction`) that has detections."""
|
"""Jump to the nearest frame (in `direction`) that has detections."""
|
||||||
@@ -318,6 +399,18 @@ class MainWindow(QMainWindow):
|
|||||||
"(сначала «Детектировать все»)"
|
"(сначала «Детектировать все»)"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _jump_to_frame(self) -> None:
|
||||||
|
"""Ask for a 1-based frame number and select it (clamped to range)."""
|
||||||
|
n = self.file_list.count()
|
||||||
|
if n == 0:
|
||||||
|
return
|
||||||
|
cur = self.file_list.currentRow() + 1
|
||||||
|
num, ok = QInputDialog.getInt(
|
||||||
|
self, "Перейти к кадру", f"Номер кадра (1–{n}):", cur, 1, n
|
||||||
|
)
|
||||||
|
if ok:
|
||||||
|
self.file_list.setCurrentRow(num - 1)
|
||||||
|
|
||||||
def _on_slider(self, value: int) -> None:
|
def _on_slider(self, value: int) -> None:
|
||||||
if self._nav_sync:
|
if self._nav_sync:
|
||||||
return
|
return
|
||||||
@@ -469,6 +562,7 @@ class MainWindow(QMainWindow):
|
|||||||
"""
|
"""
|
||||||
self._begin_busy(total)
|
self._begin_busy(total)
|
||||||
self._tick_count = 0
|
self._tick_count = 0
|
||||||
|
self._job_start = time.monotonic() # for ETA in the progress messages
|
||||||
job = Job(fn)
|
job = Job(fn)
|
||||||
self._job = job
|
self._job = job
|
||||||
if on_tick is not None:
|
if on_tick is not None:
|
||||||
@@ -483,7 +577,26 @@ class MainWindow(QMainWindow):
|
|||||||
self.progress.setRange(0, total)
|
self.progress.setRange(0, total)
|
||||||
self.progress.setValue(done)
|
self.progress.setValue(done)
|
||||||
if message:
|
if message:
|
||||||
self.statusBar().showMessage(message)
|
eta = self._eta_suffix(done, total)
|
||||||
|
self.statusBar().showMessage(message + eta)
|
||||||
|
|
||||||
|
def _eta_suffix(self, done: int, total: int) -> str:
|
||||||
|
""" ' · осталось ~Xм Yс' estimated from the average rate so far (or '' if N/A)."""
|
||||||
|
start = getattr(self, "_job_start", None)
|
||||||
|
if not start or done <= 0 or total <= 0 or done >= total:
|
||||||
|
return ""
|
||||||
|
elapsed = time.monotonic() - start
|
||||||
|
if elapsed < 0.5:
|
||||||
|
return ""
|
||||||
|
remaining = elapsed / done * (total - done)
|
||||||
|
secs = int(remaining)
|
||||||
|
if secs >= 3600:
|
||||||
|
text = f"{secs // 3600}ч {secs % 3600 // 60}м"
|
||||||
|
elif secs >= 60:
|
||||||
|
text = f"{secs // 60}м {secs % 60}с"
|
||||||
|
else:
|
||||||
|
text = f"{secs}с"
|
||||||
|
return f" · осталось ~{text}"
|
||||||
|
|
||||||
def _finish_job(self, result, on_done) -> None:
|
def _finish_job(self, result, on_done) -> None:
|
||||||
cancelled = self._job.cancelled if self._job is not None else False
|
cancelled = self._job.cancelled if self._job is not None else False
|
||||||
@@ -500,7 +613,10 @@ class MainWindow(QMainWindow):
|
|||||||
# --------------------------------------------------------------- detector
|
# --------------------------------------------------------------- detector
|
||||||
def _make_detector(self):
|
def _make_detector(self):
|
||||||
d = self._cfg.detection
|
d = self._cfg.detection
|
||||||
key = (tuple(sorted(self._cfg.detector_models)), d.yolo_conf, d.yolo_imgsz)
|
key = (
|
||||||
|
tuple(sorted(self._cfg.detector_models)), d.yolo_conf, d.yolo_imgsz,
|
||||||
|
self._cfg.cross_model_nms, self._cfg.nms_iou,
|
||||||
|
)
|
||||||
if key != self._detector_key:
|
if key != self._detector_key:
|
||||||
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
|
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
|
||||||
self._detector_key = key
|
self._detector_key = key
|
||||||
@@ -540,11 +656,74 @@ class MainWindow(QMainWindow):
|
|||||||
act.setChecked(e.path in selected)
|
act.setChecked(e.path in selected)
|
||||||
act.toggled.connect(lambda on, p=e.path: self._on_model_toggled(p, on))
|
act.toggled.connect(lambda on, p=e.path: self._on_model_toggled(p, on))
|
||||||
self._models_menu.addSeparator()
|
self._models_menu.addSeparator()
|
||||||
|
nms = self._models_menu.addAction("Объединять пересечения (NMS)")
|
||||||
|
nms.setCheckable(True)
|
||||||
|
nms.setChecked(self._cfg.cross_model_nms)
|
||||||
|
nms.setToolTip(
|
||||||
|
"Убирать дублирующие рамки от перекрывающихся моделей (по IoU; остаётся\n"
|
||||||
|
"рамка с большей уверенностью). Меняет результат — кэш пересчитывается."
|
||||||
|
)
|
||||||
|
nms.toggled.connect(self._on_nms_toggled)
|
||||||
|
self._models_menu.addAction("Пороги по моделям…", self._edit_model_thresholds)
|
||||||
|
self._models_menu.addSeparator()
|
||||||
self._models_menu.addAction("Добавить модель…", self._add_model)
|
self._models_menu.addAction("Добавить модель…", self._add_model)
|
||||||
self._models_menu.addAction("Открыть папку моделей", self._open_models_dir)
|
self._models_menu.addAction("Открыть папку моделей", self._open_models_dir)
|
||||||
self._models_menu.addAction("Обновить список", self._rebuild_models_menu)
|
self._models_menu.addAction("Обновить список", self._rebuild_models_menu)
|
||||||
self._update_models_button()
|
self._update_models_button()
|
||||||
|
|
||||||
|
def _on_nms_toggled(self, on: bool) -> None:
|
||||||
|
self._cfg.cross_model_nms = on
|
||||||
|
self._persist_settings()
|
||||||
|
self._invalidate_results() # merging changes detections => recompute
|
||||||
|
self.statusBar().showMessage(
|
||||||
|
"Объединение пересечений (NMS): " + ("вкл" if on else "выкл")
|
||||||
|
)
|
||||||
|
|
||||||
|
def _edit_model_thresholds(self) -> None:
|
||||||
|
"""Dialog: per-model overlay threshold overrides (display-only, not detection)."""
|
||||||
|
models = [m for m in self._cfg.detector_models if Path(m).is_file()]
|
||||||
|
if not models:
|
||||||
|
QMessageBox.information(
|
||||||
|
self, "Нет моделей", "Сначала отметьте хотя бы одну модель."
|
||||||
|
)
|
||||||
|
return
|
||||||
|
dlg = QDialog(self)
|
||||||
|
dlg.setWindowTitle("Пороги отображения по моделям")
|
||||||
|
layout = QVBoxLayout(dlg)
|
||||||
|
layout.addWidget(QLabel(
|
||||||
|
f"Порог отображения для каждой модели (по умолчанию {self._cfg.default_threshold:.2f}).\n"
|
||||||
|
"Влияет только на отрисовку/подсветку, не на саму детекцию."
|
||||||
|
))
|
||||||
|
form = QFormLayout()
|
||||||
|
spins: dict[str, QDoubleSpinBox] = {}
|
||||||
|
for m in models:
|
||||||
|
stem = Path(m).stem
|
||||||
|
spin = QDoubleSpinBox()
|
||||||
|
spin.setRange(0.0, 1.0)
|
||||||
|
spin.setSingleStep(0.05)
|
||||||
|
spin.setValue(self._cfg.model_thresholds.get(stem, self._cfg.default_threshold))
|
||||||
|
spins[stem] = spin
|
||||||
|
form.addRow(stem, spin)
|
||||||
|
layout.addLayout(form)
|
||||||
|
buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
|
||||||
|
buttons.accepted.connect(dlg.accept)
|
||||||
|
buttons.rejected.connect(dlg.reject)
|
||||||
|
layout.addWidget(buttons)
|
||||||
|
if dlg.exec() != QDialog.Accepted:
|
||||||
|
return
|
||||||
|
# Store only the overrides that differ from the global default — keeps it tidy.
|
||||||
|
thresholds = {
|
||||||
|
stem: round(spin.value(), 4)
|
||||||
|
for stem, spin in spins.items()
|
||||||
|
if abs(spin.value() - self._cfg.default_threshold) > 1e-9
|
||||||
|
}
|
||||||
|
self._cfg.model_thresholds = thresholds
|
||||||
|
self.view.set_model_thresholds(thresholds)
|
||||||
|
self._persist_settings()
|
||||||
|
self.statusBar().showMessage(
|
||||||
|
f"Пороги по моделям обновлены ({len(thresholds)} переопределений)"
|
||||||
|
)
|
||||||
|
|
||||||
def _update_models_button(self) -> None:
|
def _update_models_button(self) -> None:
|
||||||
n = len([m for m in self._cfg.detector_models if Path(m).is_file()])
|
n = len([m for m in self._cfg.detector_models if Path(m).is_file()])
|
||||||
self.models_button.setText(f"Модели ({n}) ▾")
|
self.models_button.setText(f"Модели ({n}) ▾")
|
||||||
@@ -592,6 +771,16 @@ class MainWindow(QMainWindow):
|
|||||||
with contextlib.suppress(OSError, AttributeError):
|
with contextlib.suppress(OSError, AttributeError):
|
||||||
os.startfile(str(root)) # noqa: S606 - Windows: open in Explorer
|
os.startfile(str(root)) # noqa: S606 - Windows: open in Explorer
|
||||||
|
|
||||||
|
def _open_restored_dir(self) -> None:
|
||||||
|
"""Open the project's restored/ folder (where restored images are saved)."""
|
||||||
|
if self._project is None:
|
||||||
|
QMessageBox.information(self, "Нет проекта", "Сначала откройте проект.")
|
||||||
|
return
|
||||||
|
d = self._project.restored_dir
|
||||||
|
d.mkdir(parents=True, exist_ok=True)
|
||||||
|
with contextlib.suppress(OSError, AttributeError):
|
||||||
|
os.startfile(str(d)) # noqa: S606 - Windows: open in Explorer
|
||||||
|
|
||||||
def _invalidate_results(self) -> None:
|
def _invalidate_results(self) -> None:
|
||||||
"""Detector changed — drop the in-memory cache and refresh the current image.
|
"""Detector changed — drop the in-memory cache and refresh the current image.
|
||||||
|
|
||||||
@@ -775,6 +964,7 @@ class MainWindow(QMainWindow):
|
|||||||
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)
|
||||||
self.view.set_threshold(self._cfg.default_threshold)
|
self.view.set_threshold(self._cfg.default_threshold)
|
||||||
|
self.view.set_model_thresholds(self._cfg.model_thresholds)
|
||||||
|
|
||||||
def _persist_settings(self) -> None:
|
def _persist_settings(self) -> None:
|
||||||
"""Save settings to the global defaults and (if open) into the project."""
|
"""Save settings to the global defaults and (if open) into the project."""
|
||||||
@@ -795,7 +985,7 @@ class MainWindow(QMainWindow):
|
|||||||
dialog = ExtractDialog(self)
|
dialog = ExtractDialog(self)
|
||||||
if dialog.exec() != QDialog.Accepted:
|
if dialog.exec() != QDialog.Accepted:
|
||||||
return
|
return
|
||||||
keyframes_only, step, max_dim = dialog.options()
|
keyframes_only, step, max_dim, jpg_quality = dialog.options()
|
||||||
video = Path(path)
|
video = Path(path)
|
||||||
root = video.parent / f"{video.stem}_frames"
|
root = video.parent / f"{video.stem}_frames"
|
||||||
if Project.is_project(root):
|
if Project.is_project(root):
|
||||||
@@ -821,7 +1011,7 @@ class MainWindow(QMainWindow):
|
|||||||
try:
|
try:
|
||||||
saved = extract_frames(
|
saved = extract_frames(
|
||||||
str(video), str(out), step=step, keyframes_only=keyframes_only,
|
str(video), str(out), step=step, keyframes_only=keyframes_only,
|
||||||
max_dim=max_dim, progress=cb,
|
max_dim=max_dim, jpg_quality=jpg_quality, progress=cb,
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
|
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
|
||||||
@@ -849,6 +1039,8 @@ class MainWindow(QMainWindow):
|
|||||||
files = sorted(p for p in folder.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
|
files = sorted(p for p in folder.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
|
||||||
self._files = files
|
self._files = files
|
||||||
self._results.clear()
|
self._results.clear()
|
||||||
|
self._row_restored.clear()
|
||||||
|
self._showing_restored = False # start a project in original-view mode
|
||||||
self._current = None
|
self._current = None
|
||||||
|
|
||||||
self.file_list.blockSignals(True)
|
self.file_list.blockSignals(True)
|
||||||
@@ -926,7 +1118,7 @@ class MainWindow(QMainWindow):
|
|||||||
return
|
return
|
||||||
key, dets = result
|
key, dets = result
|
||||||
self._results[key] = dets
|
self._results[key] = dets
|
||||||
self._tag_file(Path(key), len(dets))
|
self._tag_file(Path(key))
|
||||||
self._refresh_marks()
|
self._refresh_marks()
|
||||||
self._save_results()
|
self._save_results()
|
||||||
if then_show or self._current == Path(key):
|
if then_show or self._current == Path(key):
|
||||||
@@ -937,11 +1129,20 @@ class MainWindow(QMainWindow):
|
|||||||
self._start_job(fn, None, on_done=done)
|
self._start_job(fn, None, on_done=done)
|
||||||
|
|
||||||
def _show(self, path: Path) -> None:
|
def _show(self, path: Path) -> None:
|
||||||
"""Display the image with its cached detections (does not run the detector)."""
|
"""Display the image (does not run the detector).
|
||||||
|
|
||||||
|
Honours the global "show restored" view mode (``_showing_restored``): when on and
|
||||||
|
a restored version exists (memory or ``restored/``), the restored image is shown
|
||||||
|
(no overlays); otherwise the original frame with its cached detections.
|
||||||
|
"""
|
||||||
self._current = path
|
self._current = path
|
||||||
self._showing_restored = False
|
|
||||||
img = imread_unicode(str(path))
|
|
||||||
dets = self._results.get(str(path)) # None => not yet computed
|
dets = self._results.get(str(path)) # None => not yet computed
|
||||||
|
restored = self._restored_image_for(path) if self._showing_restored else None
|
||||||
|
if restored is not None:
|
||||||
|
self.view.set_image(restored, []) # restored: no overlays
|
||||||
|
self._fill_detail_table(path, restored, dets)
|
||||||
|
else:
|
||||||
|
img = imread_unicode(str(path))
|
||||||
self.view.set_image(img, dets or [])
|
self.view.set_image(img, dets or [])
|
||||||
self._fill_detail_table(path, img, dets)
|
self._fill_detail_table(path, img, dets)
|
||||||
self._update_restore_actions()
|
self._update_restore_actions()
|
||||||
@@ -1000,7 +1201,7 @@ class MainWindow(QMainWindow):
|
|||||||
"""GUI-thread handler for one streamed detect-all result."""
|
"""GUI-thread handler for one streamed detect-all result."""
|
||||||
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))
|
||||||
# If the frame being viewed was just computed, show its overlay live.
|
# 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:
|
if not self._showing_restored and self._current is not None and str(self._current) == key:
|
||||||
self._show(self._current)
|
self._show(self._current)
|
||||||
@@ -1036,28 +1237,61 @@ class MainWindow(QMainWindow):
|
|||||||
return
|
return
|
||||||
_, k, restored, engine = result
|
_, k, restored, engine = result
|
||||||
self._restored[k] = restored
|
self._restored[k] = restored
|
||||||
|
# Persist to the project's restored/ folder (like the batch run), so a single
|
||||||
|
# restore is saved on disk and survives reopening — not just held in memory.
|
||||||
|
saved_to = ""
|
||||||
|
if self._project is not None:
|
||||||
|
try:
|
||||||
|
self._project.restored_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
dst = self._project.restored_dir / f"{Path(k).stem}.jpg"
|
||||||
|
if imwrite_unicode(str(dst), restored):
|
||||||
|
saved_to = f" → {self._project.restored_dir.name}/"
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
self._refresh_restored_marks()
|
||||||
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._show(self._current)
|
||||||
self._update_restore_actions()
|
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}{saved_to}")
|
||||||
self._refresh_restored_marks()
|
|
||||||
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_done=done)
|
self._start_job(fn, None, on_done=done)
|
||||||
|
|
||||||
def _restore_all(self, force: bool = False) -> None:
|
def _restore_all(self, force: bool = False, only_detected: bool = False) -> None:
|
||||||
"""Restore every frame on a background thread, writing results to ``restored/``.
|
"""Restore frames on a background thread, writing results to ``restored/``.
|
||||||
|
|
||||||
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
|
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
|
||||||
engine skips frames already restored (resume) unless ``force``. The temporal
|
engine skips frames already restored (resume) unless ``force``. The temporal
|
||||||
engine (DeepMosaics-video) runs the whole contiguous sequence in order via
|
engine (DeepMosaics-video) runs a contiguous sequence in order via
|
||||||
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
|
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
|
||||||
|
|
||||||
|
``only_detected`` uses the YOLO detection cache to skip frames known clean:
|
||||||
|
per-frame → restore just the frames with detections; temporal → restrict the run
|
||||||
|
to the contiguous span [first hit … last hit] (clean frames inside it still run,
|
||||||
|
for recurrence). NOTE: LADA misses some mosaic, so this can miss censorship YOLO
|
||||||
|
didn't flag — "Расцензурить все" stays the thorough option.
|
||||||
"""
|
"""
|
||||||
if not self._files or self._project is None or self._busy:
|
if not self._files or self._project is None or self._busy:
|
||||||
return
|
return
|
||||||
files = list(self._files) # snapshot — favorites/move mutate self._files
|
files = list(self._files) # snapshot — favorites/move mutate self._files
|
||||||
|
is_temporal = self._cfg.restorer == "deepmosaics_video"
|
||||||
|
hits = [i for i, p in enumerate(files) if self._results.get(str(p))]
|
||||||
|
|
||||||
|
if only_detected:
|
||||||
|
if not self._results:
|
||||||
|
self.statusBar().showMessage(
|
||||||
|
"Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)"
|
||||||
|
)
|
||||||
|
return
|
||||||
|
if not hits:
|
||||||
|
self.statusBar().showMessage("Цензура не найдена ни на одном кадре — нечего расцензуривать")
|
||||||
|
return
|
||||||
|
span = (hits[0], hits[-1]) if is_temporal else None
|
||||||
|
total = (span[1] - span[0] + 1) if span else len(hits)
|
||||||
|
else:
|
||||||
|
span = None
|
||||||
total = len(files)
|
total = len(files)
|
||||||
|
|
||||||
out_dir = self._project.restored_dir
|
out_dir = self._project.restored_dir
|
||||||
out_dir.mkdir(parents=True, exist_ok=True)
|
out_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
@@ -1067,6 +1301,7 @@ class MainWindow(QMainWindow):
|
|||||||
def fn(job):
|
def fn(job):
|
||||||
restorer = self._make_restorer() # built on the worker (may raise)
|
restorer = self._make_restorer() # built on the worker (may raise)
|
||||||
frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads
|
frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads
|
||||||
|
state = {"done": 0} # frames processed (for the progress bar)
|
||||||
|
|
||||||
def get_frame(i):
|
def get_frame(i):
|
||||||
img = frame_cache.get(i)
|
img = frame_cache.get(i)
|
||||||
@@ -1079,42 +1314,48 @@ class MainWindow(QMainWindow):
|
|||||||
frame_cache[i] = img
|
frame_cache[i] = img
|
||||||
return img
|
return img
|
||||||
|
|
||||||
def emit(i, restored):
|
def emit(i, restored, *, verb="Расцензуривание"):
|
||||||
|
if restored is not None:
|
||||||
imwrite_unicode(str(out_path(files[i])), restored)
|
imwrite_unicode(str(out_path(files[i])), restored)
|
||||||
job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}")
|
state["done"] += 1
|
||||||
|
job.progress(state["done"], total, f"{verb} {state['done']}/{total}: {files[i].name}")
|
||||||
|
|
||||||
if restorer.temporal:
|
if restorer.temporal:
|
||||||
|
start = span[0] if span else 0
|
||||||
|
end = span[1] if span else len(files) - 1
|
||||||
restorer.restore_sequence(
|
restorer.restore_sequence(
|
||||||
total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled
|
end - start + 1,
|
||||||
|
lambda li: get_frame(start + li),
|
||||||
|
lambda _li: [],
|
||||||
|
lambda li, res: emit(start + li, res),
|
||||||
|
should_cancel=lambda: job.cancelled,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
for i, p in enumerate(files):
|
indices = hits if only_detected else range(len(files))
|
||||||
|
for i in indices:
|
||||||
if job.cancelled:
|
if job.cancelled:
|
||||||
break
|
break
|
||||||
|
p = files[i]
|
||||||
if not force and out_path(p).is_file():
|
if not force and out_path(p).is_file():
|
||||||
job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}")
|
emit(i, None, verb="Пропуск") # already restored — count, don't rewrite
|
||||||
continue
|
continue
|
||||||
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
||||||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def done(_result, cancelled):
|
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._refresh_restored_marks()
|
self._refresh_restored_marks()
|
||||||
|
if not cancelled and self._row_restored:
|
||||||
|
self._showing_restored = True # auto-switch to viewing the results
|
||||||
|
if self._current is not None: # re-show current frame in the (new) mode
|
||||||
|
self._show(self._current)
|
||||||
self.statusBar().showMessage(
|
self.statusBar().showMessage(
|
||||||
"Расцензуривание отменено" if cancelled
|
"Расцензуривание отменено" if cancelled
|
||||||
else f"Готово: результаты в {out_dir.name}/ ({total} кадров)"
|
else f"Готово: результаты в {out_dir.name}/ ({total} кадров) — показываю расцензуренное"
|
||||||
)
|
)
|
||||||
|
|
||||||
self.statusBar().showMessage("Пакетное расцензуривание…")
|
scope = " (по детекции)" if only_detected else ""
|
||||||
|
self.statusBar().showMessage(f"Пакетное расцензуривание{scope}…")
|
||||||
self._start_job(fn, total, on_done=done)
|
self._start_job(fn, total, on_done=done)
|
||||||
|
|
||||||
def _make_restorer(self):
|
def _make_restorer(self):
|
||||||
@@ -1133,30 +1374,58 @@ class MainWindow(QMainWindow):
|
|||||||
self._restorer_key = None # rebuild on next restore
|
self._restorer_key = None # rebuild on next restore
|
||||||
self.statusBar().showMessage(f"Движок восстановления: {self._cfg.restorer}")
|
self.statusBar().showMessage(f"Движок восстановления: {self._cfg.restorer}")
|
||||||
|
|
||||||
|
# ---------------------------------------------------------- restored access
|
||||||
|
def _restored_disk_path(self, path: Path) -> Path | None:
|
||||||
|
"""Path of the restored output for ``path`` in ``restored/``, if it exists."""
|
||||||
|
if self._project is None:
|
||||||
|
return None
|
||||||
|
rp = self._project.restored_dir / f"{Path(path).stem}.jpg"
|
||||||
|
return rp if rp.is_file() else None
|
||||||
|
|
||||||
|
def _has_restored(self, path: Path) -> bool:
|
||||||
|
"""Cheap check (no decode): is there a restored version of ``path``?"""
|
||||||
|
return str(path) in self._restored or self._restored_disk_path(path) is not None
|
||||||
|
|
||||||
|
def _restored_image_for(self, path: Path):
|
||||||
|
"""Return the restored image for ``path`` (from memory or ``restored/``), or None."""
|
||||||
|
img = self._restored.get(str(path))
|
||||||
|
if img is not None:
|
||||||
|
return img
|
||||||
|
rp = self._restored_disk_path(path)
|
||||||
|
return imread_unicode(str(rp)) if rp is not None else None
|
||||||
|
|
||||||
def _toggle_restored(self) -> None:
|
def _toggle_restored(self) -> None:
|
||||||
if self._current is None or str(self._current) not in self._restored:
|
"""Flip the global view mode between original and restored, then re-show."""
|
||||||
return
|
|
||||||
self._showing_restored = not self._showing_restored
|
self._showing_restored = not self._showing_restored
|
||||||
key = str(self._current)
|
if self._current is not None:
|
||||||
if self._showing_restored:
|
self._show(self._current)
|
||||||
self.view.set_image(self._restored[key], [])
|
|
||||||
else:
|
else:
|
||||||
self.view.set_image(imread_unicode(key), self._results.get(key) or [])
|
|
||||||
self._update_restore_actions()
|
self._update_restore_actions()
|
||||||
|
self.statusBar().showMessage(
|
||||||
|
"Показ: расцензуренное (где есть)" if self._showing_restored else "Показ: оригинал"
|
||||||
|
)
|
||||||
|
|
||||||
def _update_restore_actions(self) -> None:
|
def _update_restore_actions(self) -> None:
|
||||||
has = self._current is not None and str(self._current) in self._restored
|
any_restored = bool(self._row_restored) or bool(self._restored)
|
||||||
self.toggle_restored_action.setEnabled(has)
|
self.toggle_restored_action.setEnabled(any_restored)
|
||||||
|
# Keep the checkable state in sync with the mode (setChecked emits `toggled`,
|
||||||
|
# not `triggered`, so this never re-enters `_toggle_restored`).
|
||||||
|
self.toggle_restored_action.setChecked(self._showing_restored and any_restored)
|
||||||
self.toggle_restored_action.setText(
|
self.toggle_restored_action.setText(
|
||||||
"Показать оригинал" if self._showing_restored else "Показать результат"
|
"Показать оригинал" if self._showing_restored else "Показать расцензуренное"
|
||||||
|
)
|
||||||
|
self.save_restored_action.setEnabled(
|
||||||
|
self._current is not None and self._has_restored(self._current)
|
||||||
)
|
)
|
||||||
self.save_restored_action.setEnabled(has)
|
|
||||||
|
|
||||||
def _save_restored(self) -> None:
|
def _save_restored(self) -> None:
|
||||||
if self._current is None or str(self._current) not in self._restored:
|
if self._current is None:
|
||||||
|
return
|
||||||
|
restored = self._restored_image_for(self._current)
|
||||||
|
if restored is None:
|
||||||
return
|
return
|
||||||
out = self._unique_dest(self._current.parent, f"{self._current.stem}_restored.jpg")
|
out = self._unique_dest(self._current.parent, f"{self._current.stem}_restored.jpg")
|
||||||
if imwrite_unicode(str(out), self._restored[str(self._current)]):
|
if imwrite_unicode(str(out), restored):
|
||||||
self.statusBar().showMessage(f"Сохранено: {out}")
|
self.statusBar().showMessage(f"Сохранено: {out}")
|
||||||
else:
|
else:
|
||||||
QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.")
|
QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.")
|
||||||
@@ -1227,24 +1496,77 @@ class MainWindow(QMainWindow):
|
|||||||
_TINT_HIT = QColor(200, 80, 80, 70)
|
_TINT_HIT = QColor(200, 80, 80, 70)
|
||||||
_TINT_CLEAN = QColor(90, 160, 90, 50)
|
_TINT_CLEAN = QColor(90, 160, 90, 50)
|
||||||
|
|
||||||
def _set_row_tag(self, item: QListWidgetItem, count: int) -> None:
|
def _relabel_row(self, item: QListWidgetItem) -> None:
|
||||||
base = item.data(Qt.UserRole + 1)
|
"""Set a row's text/tint/tooltip from its detection + restoration state.
|
||||||
item.setText(f"{base} · {count}" if count else f"{base} · —")
|
|
||||||
item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN)
|
|
||||||
|
|
||||||
def _tag_file(self, path: Path, count: int) -> None:
|
Text: ``name · <count|—> ✓`` — the count/— suffix appears once detected (red tint
|
||||||
|
= censorship, green = clean), and a trailing ✓ marks frames that have a restored
|
||||||
|
version in ``restored/``.
|
||||||
|
"""
|
||||||
|
base = item.data(Qt.UserRole + 1)
|
||||||
|
path = item.data(Qt.UserRole)
|
||||||
|
restored = path in self._row_restored
|
||||||
|
if path in self._results: # `in`, not truthy: empty list = clean
|
||||||
|
dets = self._results[path]
|
||||||
|
suffix = f" · {len(dets)}" if dets else " · —"
|
||||||
|
item.setBackground(self._TINT_HIT if dets else self._TINT_CLEAN)
|
||||||
|
else:
|
||||||
|
suffix = ""
|
||||||
|
item.setBackground(QBrush())
|
||||||
|
item.setText(f"{base}{suffix}{' ✓' if restored else ''}")
|
||||||
|
item.setToolTip("Есть расцензуренная версия (restored/)" if restored else "")
|
||||||
|
item.setHidden(not self._row_matches_filter(path))
|
||||||
|
|
||||||
|
def _row_matches_filter(self, path: str) -> bool:
|
||||||
|
"""Whether a row should be visible under the current filter mode."""
|
||||||
|
mode = self._filter_mode
|
||||||
|
if mode == "all":
|
||||||
|
return True
|
||||||
|
if mode == "hits":
|
||||||
|
return bool(self._results.get(path))
|
||||||
|
if mode == "clean":
|
||||||
|
return path in self._results and not self._results[path]
|
||||||
|
if mode == "uncomputed":
|
||||||
|
return path not in self._results
|
||||||
|
if mode == "restored":
|
||||||
|
return path in self._row_restored
|
||||||
|
if mode == "unrestored":
|
||||||
|
return path not in self._row_restored
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _on_filter_changed(self) -> None:
|
||||||
|
self._filter_mode = self.filter_combo.currentData() or "all"
|
||||||
|
self._relabel_all() # re-applies hidden state per row
|
||||||
|
# If the current row got hidden, jump to the first visible one so the view isn't stale.
|
||||||
|
cur = self.file_list.currentItem()
|
||||||
|
if cur is not None and cur.isHidden():
|
||||||
|
for i in range(self.file_list.count()):
|
||||||
|
if not self.file_list.item(i).isHidden():
|
||||||
|
self.file_list.setCurrentRow(i)
|
||||||
|
break
|
||||||
|
n_vis = sum(1 for i in range(self.file_list.count()) if not self.file_list.item(i).isHidden())
|
||||||
|
self.statusBar().showMessage(
|
||||||
|
f"Фильтр: {self.filter_combo.currentText()} — показано {n_vis} из {len(self._files)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _relabel_all(self) -> None:
|
||||||
|
for i in range(self.file_list.count()):
|
||||||
|
self._relabel_row(self.file_list.item(i))
|
||||||
|
|
||||||
|
def _tag_file(self, path: Path) -> None:
|
||||||
for i in range(self.file_list.count()):
|
for i in range(self.file_list.count()):
|
||||||
item = self.file_list.item(i)
|
item = self.file_list.item(i)
|
||||||
if item.data(Qt.UserRole) == str(path):
|
if item.data(Qt.UserRole) == str(path):
|
||||||
self._set_row_tag(item, count)
|
self._relabel_row(item)
|
||||||
return
|
return
|
||||||
|
|
||||||
# ------------------------------------------------------------- result cache
|
# ------------------------------------------------------------- result cache
|
||||||
def _results_key(self) -> dict:
|
def _results_key(self) -> dict:
|
||||||
"""Detector identity used to tag/validate the on-disk detection cache."""
|
"""Detector identity used to tag/validate the on-disk detection cache."""
|
||||||
d = self._cfg.detection
|
d = self._cfg.detection
|
||||||
|
nms_iou = self._cfg.nms_iou if self._cfg.cross_model_nms else None
|
||||||
return detection_cache.make_key(
|
return detection_cache.make_key(
|
||||||
self._cfg.detector_models, d.yolo_conf, d.yolo_imgsz
|
self._cfg.detector_models, d.yolo_conf, d.yolo_imgsz, nms_iou=nms_iou
|
||||||
)
|
)
|
||||||
|
|
||||||
def _save_results(self) -> None:
|
def _save_results(self) -> None:
|
||||||
@@ -1265,20 +1587,16 @@ class MainWindow(QMainWindow):
|
|||||||
if not cached:
|
if not cached:
|
||||||
return 0
|
return 0
|
||||||
self._results = cached
|
self._results = cached
|
||||||
for i in range(self.file_list.count()):
|
self._relabel_all()
|
||||||
item = self.file_list.item(i)
|
|
||||||
path = item.data(Qt.UserRole)
|
|
||||||
if path in self._results: # `in`, not truthy: empty list = checked-clean
|
|
||||||
self._set_row_tag(item, len(self._results[path]))
|
|
||||||
return len(cached)
|
return len(cached)
|
||||||
|
|
||||||
def _clear_results(self) -> None:
|
def _clear_results(self) -> None:
|
||||||
"""Drop all cached detections and reset row labels/tints (keeps the detector)."""
|
"""Drop all cached detections and reset row labels/tints (keeps the detector).
|
||||||
|
|
||||||
|
Restored ✓ markers stay — restoration is independent of detection.
|
||||||
|
"""
|
||||||
self._results.clear()
|
self._results.clear()
|
||||||
for i in range(self.file_list.count()):
|
self._relabel_all()
|
||||||
item = self.file_list.item(i)
|
|
||||||
item.setText(item.data(Qt.UserRole + 1))
|
|
||||||
item.setBackground(QBrush())
|
|
||||||
self._refresh_marks()
|
self._refresh_marks()
|
||||||
|
|
||||||
def _refresh_marks(self) -> None:
|
def _refresh_marks(self) -> None:
|
||||||
@@ -1313,13 +1631,17 @@ class MainWindow(QMainWindow):
|
|||||||
stems = {p.stem for p in self._project.restored_dir.glob("*.jpg")}
|
stems = {p.stem for p in self._project.restored_dir.glob("*.jpg")}
|
||||||
mem = set(self._restored) # single-frame restores held in memory (not on disk yet)
|
mem = set(self._restored) # single-frame restores held in memory (not on disk yet)
|
||||||
rows: set[int] = set()
|
rows: set[int] = set()
|
||||||
|
paths: set[str] = set()
|
||||||
if stems or mem:
|
if stems or mem:
|
||||||
for i in range(self.file_list.count()):
|
for i in range(self.file_list.count()):
|
||||||
fp = Path(self.file_list.item(i).data(Qt.UserRole))
|
sp = self.file_list.item(i).data(Qt.UserRole)
|
||||||
if fp.stem in stems or str(fp) in mem:
|
if Path(sp).stem in stems or sp in mem:
|
||||||
rows.add(i)
|
rows.add(i)
|
||||||
|
paths.add(sp)
|
||||||
|
self._row_restored = paths
|
||||||
self._restored_count = len(rows)
|
self._restored_count = len(rows)
|
||||||
self.frame_slider.set_restored_marks(rows)
|
self.frame_slider.set_restored_marks(rows)
|
||||||
|
self._relabel_all() # show/refresh the ✓ markers in the file list
|
||||||
self._update_counts_label()
|
self._update_counts_label()
|
||||||
|
|
||||||
def _fill_detail_table(self, path: Path, img, dets: list[Detection] | None) -> None:
|
def _fill_detail_table(self, path: Path, img, dets: list[Detection] | None) -> None:
|
||||||
|
|||||||
@@ -0,0 +1,247 @@
|
|||||||
|
"""Headless smoke tests for HVideoTool.
|
||||||
|
|
||||||
|
No formal test suite and no model weights in the repo, so this exercises the pure
|
||||||
|
core logic (detection cache, cross-model NMS, list-filter predicate, ETA formatting,
|
||||||
|
extract-dialog options, per-project settings round-trip) plus a minimal offscreen
|
||||||
|
``MainWindow`` build on a tiny throwaway project. It avoids torch/ultralytics entirely
|
||||||
|
(detection results are injected directly into ``_results``).
|
||||||
|
|
||||||
|
Run::
|
||||||
|
|
||||||
|
set QT_QPA_PLATFORM=offscreen
|
||||||
|
set PYTHONIOENCODING=utf-8
|
||||||
|
.venv\\Scripts\\python.exe scripts\\smoke_test.py
|
||||||
|
|
||||||
|
Exits non-zero on the first failure.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
# Run Qt without a display and keep Unicode console output sane on Windows.
|
||||||
|
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
|
||||||
|
os.environ.setdefault("PYTHONIOENCODING", "utf-8")
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
|
|
||||||
|
from hvideotool.config import AppConfig # noqa: E402
|
||||||
|
from hvideotool.core.detection import cache as detection_cache # noqa: E402
|
||||||
|
from hvideotool.core.detection.multi import MultiYoloDetector, _iou, _nms # noqa: E402
|
||||||
|
from hvideotool.core.detection.types import CensorType, Detection # noqa: E402
|
||||||
|
from hvideotool.core.imageio import imwrite_unicode # noqa: E402
|
||||||
|
from hvideotool.core.project import Project # noqa: E402
|
||||||
|
|
||||||
|
_failures: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
def check(cond: bool, msg: str) -> None:
|
||||||
|
status = "PASS" if cond else "FAIL"
|
||||||
|
print(f" [{status}] {msg}")
|
||||||
|
if not cond:
|
||||||
|
_failures.append(msg)
|
||||||
|
|
||||||
|
|
||||||
|
def _det(score: float, bbox, *, label="", model="", type=CensorType.MOSAIC) -> Detection:
|
||||||
|
return Detection(type=type, score=score, bbox=bbox, label=label, model=model)
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------- core tests
|
||||||
|
def test_cache_atomic_roundtrip() -> None:
|
||||||
|
print("cache: atomic save/load round-trip")
|
||||||
|
with tempfile.TemporaryDirectory() as d:
|
||||||
|
base = Path(d)
|
||||||
|
cache_file = base / "detections.json"
|
||||||
|
key = detection_cache.make_key(["m/a.pt", "m/b.pt"], 0.2, 640)
|
||||||
|
results = {
|
||||||
|
str(base / "001.jpg"): [_det(0.9, (1, 2, 3, 4), label="mosaic", model="a")],
|
||||||
|
str(base / "002.jpg"): [], # checked-clean
|
||||||
|
}
|
||||||
|
ok = detection_cache.save_results(cache_file, key, results)
|
||||||
|
check(ok, "save_results returns True")
|
||||||
|
check(cache_file.is_file(), "cache file created")
|
||||||
|
check(not (base / "detections.json.tmp").exists(), "temp file cleaned up")
|
||||||
|
loaded = detection_cache.load_results(cache_file, key, base)
|
||||||
|
check(loaded is not None, "load returns a dict for matching key")
|
||||||
|
check(set(loaded) == set(results), "all basenames round-trip")
|
||||||
|
d1 = loaded[str(base / "001.jpg")][0]
|
||||||
|
check(d1.score == 0.9 and d1.model == "a", "detection fields preserved")
|
||||||
|
check(loaded[str(base / "002.jpg")] == [], "empty (clean) entry preserved")
|
||||||
|
bad = detection_cache.load_results(
|
||||||
|
cache_file, detection_cache.make_key(["x.pt"], 0.2, 640), base
|
||||||
|
)
|
||||||
|
check(bad is None, "mismatched detector key => None (cache ignored)")
|
||||||
|
# NMS only enters the key when enabled: default key stays valid, NMS key differs.
|
||||||
|
k_off = detection_cache.make_key(["a.pt"], 0.2, 640)
|
||||||
|
k_off2 = detection_cache.make_key(["a.pt"], 0.2, 640, nms_iou=None)
|
||||||
|
k_on = detection_cache.make_key(["a.pt"], 0.2, 640, nms_iou=0.6)
|
||||||
|
check(k_off == k_off2 and "nms_iou" not in k_off, "NMS-off key unchanged (no nms field)")
|
||||||
|
check(k_on != k_off and k_on.get("nms_iou") == 0.6, "NMS-on key is distinct")
|
||||||
|
|
||||||
|
|
||||||
|
def test_nms() -> None:
|
||||||
|
print("detection: cross-model NMS")
|
||||||
|
check(abs(_iou((0, 0, 10, 10), (0, 0, 10, 10)) - 1.0) < 1e-9, "IoU identical = 1.0")
|
||||||
|
check(_iou((0, 0, 10, 10), (100, 100, 5, 5)) == 0.0, "IoU disjoint = 0.0")
|
||||||
|
# Two near-duplicate boxes + one distinct: NMS keeps the higher score + the distinct.
|
||||||
|
dets = [
|
||||||
|
_det(0.6, (0, 0, 10, 10), model="a"),
|
||||||
|
_det(0.9, (1, 1, 10, 10), model="b"), # overlaps the first heavily
|
||||||
|
_det(0.8, (200, 200, 10, 10), model="c"), # separate region
|
||||||
|
]
|
||||||
|
kept = _nms(dets, 0.5)
|
||||||
|
check(len(kept) == 2, "two duplicates merged to one (+ the distinct box)")
|
||||||
|
check(any(k.score == 0.9 for k in kept), "higher-score duplicate survives")
|
||||||
|
check(not any(k.score == 0.6 for k in kept), "lower-score duplicate dropped")
|
||||||
|
|
||||||
|
class _Stub:
|
||||||
|
def __init__(self, ds):
|
||||||
|
self._ds = ds
|
||||||
|
|
||||||
|
@property
|
||||||
|
def name(self):
|
||||||
|
return "stub"
|
||||||
|
|
||||||
|
def detect(self, frame):
|
||||||
|
return list(self._ds)
|
||||||
|
|
||||||
|
merged = MultiYoloDetector(
|
||||||
|
[_Stub([dets[0]]), _Stub([dets[1], dets[2]])], nms_iou=0.5
|
||||||
|
).detect(None)
|
||||||
|
check(len(merged) == 2, "MultiYoloDetector applies NMS across detectors")
|
||||||
|
no_nms = MultiYoloDetector([_Stub([dets[0]]), _Stub([dets[1], dets[2]])]).detect(None)
|
||||||
|
check(len(no_nms) == 3, "without nms_iou, all detections are concatenated")
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_dialog_options() -> None:
|
||||||
|
print("extract dialog: options() includes JPEG quality")
|
||||||
|
from PySide6.QtWidgets import QApplication
|
||||||
|
|
||||||
|
from hvideotool.ui.extract_dialog import ExtractDialog
|
||||||
|
|
||||||
|
_ensure_app(QApplication)
|
||||||
|
dlg = ExtractDialog()
|
||||||
|
opts = dlg.options()
|
||||||
|
check(len(opts) == 4, "options() is a 4-tuple (keyframes, step, max_dim, quality)")
|
||||||
|
keyframes, step, max_dim, quality = opts
|
||||||
|
check(keyframes is False and step == 1, "defaults to every-frame (step=1)")
|
||||||
|
check(1 <= quality <= 100, "quality in 1..100")
|
||||||
|
|
||||||
|
|
||||||
|
def test_settings_roundtrip() -> None:
|
||||||
|
print("project: per-project settings round-trip (new fields)")
|
||||||
|
with tempfile.TemporaryDirectory() as d:
|
||||||
|
cfg = AppConfig()
|
||||||
|
cfg.model_thresholds = {"penis": 0.42}
|
||||||
|
cfg.cross_model_nms = True
|
||||||
|
cfg.nms_iou = 0.55
|
||||||
|
cfg.default_threshold = 0.3
|
||||||
|
proj = Project.create(Path(d) / "P", name="P")
|
||||||
|
proj.update_from_config(cfg)
|
||||||
|
proj.save()
|
||||||
|
reloaded = Project.load(proj.root)
|
||||||
|
cfg2 = AppConfig()
|
||||||
|
reloaded.apply_to_config(cfg2)
|
||||||
|
check(cfg2.model_thresholds == {"penis": 0.42}, "model_thresholds persisted")
|
||||||
|
check(cfg2.cross_model_nms is True, "cross_model_nms persisted")
|
||||||
|
check(abs(cfg2.nms_iou - 0.55) < 1e-9, "nms_iou persisted")
|
||||||
|
check(not (proj.root / "project.json.tmp").exists(), "project.json.tmp cleaned up")
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------- GUI smoke tests
|
||||||
|
_APP = None
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_app(QApplication):
|
||||||
|
global _APP
|
||||||
|
if _APP is None:
|
||||||
|
_APP = QApplication.instance() or QApplication([])
|
||||||
|
return _APP
|
||||||
|
|
||||||
|
|
||||||
|
def test_mainwindow_filter_and_jump() -> None:
|
||||||
|
print("MainWindow: build, filter, jump, ETA (offscreen)")
|
||||||
|
from PySide6.QtWidgets import QApplication
|
||||||
|
|
||||||
|
from hvideotool.ui.main_window import MainWindow
|
||||||
|
|
||||||
|
_ensure_app(QApplication)
|
||||||
|
with tempfile.TemporaryDirectory() as d:
|
||||||
|
proj = Project.create(Path(d) / "Proj", name="Proj")
|
||||||
|
img = np.zeros((16, 16, 3), dtype=np.uint8)
|
||||||
|
names = [f"{i:03d}.jpg" for i in range(1, 6)]
|
||||||
|
for n in names:
|
||||||
|
imwrite_unicode(str(proj.frames_dir / n), img)
|
||||||
|
|
||||||
|
w = MainWindow(AppConfig())
|
||||||
|
w._open_project(proj)
|
||||||
|
check(w.file_list.count() == 5, "all 5 frames listed")
|
||||||
|
|
||||||
|
files = list(w._files)
|
||||||
|
# Inject detection results: frame0 has a hit, frame1 is clean, rest uncomputed.
|
||||||
|
w._results[str(files[0])] = [_det(0.9, (1, 1, 4, 4), label="mosaic", model="a")]
|
||||||
|
w._results[str(files[1])] = []
|
||||||
|
w._relabel_all()
|
||||||
|
w._refresh_marks()
|
||||||
|
|
||||||
|
# Filter: "С цензурой" (hits) shows only frame0.
|
||||||
|
w.filter_combo.setCurrentIndex(w.filter_combo.findData("hits"))
|
||||||
|
vis = [i for i in range(w.file_list.count()) if not w.file_list.item(i).isHidden()]
|
||||||
|
check(vis == [0], "filter 'hits' shows only the censored frame")
|
||||||
|
|
||||||
|
# Filter: "Не рассчитано" shows the 3 uncomputed frames.
|
||||||
|
w.filter_combo.setCurrentIndex(w.filter_combo.findData("uncomputed"))
|
||||||
|
vis = [i for i in range(w.file_list.count()) if not w.file_list.item(i).isHidden()]
|
||||||
|
check(vis == [2, 3, 4], "filter 'uncomputed' shows the not-yet-detected frames")
|
||||||
|
|
||||||
|
# _step skips hidden rows: from row2, next visible is row3.
|
||||||
|
w.file_list.setCurrentRow(2)
|
||||||
|
w._step(1)
|
||||||
|
check(w.file_list.currentRow() == 3, "_step skips filtered-out rows")
|
||||||
|
|
||||||
|
# Back to all.
|
||||||
|
w.filter_combo.setCurrentIndex(w.filter_combo.findData("all"))
|
||||||
|
vis = [i for i in range(w.file_list.count()) if not w.file_list.item(i).isHidden()]
|
||||||
|
check(len(vis) == 5, "filter 'all' shows everything again")
|
||||||
|
|
||||||
|
# Jump-to-frame core (bypassing the modal dialog): selecting row 4.
|
||||||
|
w.file_list.setCurrentRow(4)
|
||||||
|
check(w.pos_label.text() == "5 / 5", "position label reflects the current frame")
|
||||||
|
|
||||||
|
# ETA formatting.
|
||||||
|
w._job_start = __import__("time").monotonic() - 10.0 # 10s elapsed
|
||||||
|
suffix = w._eta_suffix(2, 10) # 2/10 done in 10s => ~40s remaining
|
||||||
|
check("осталось" in suffix, "ETA suffix produced for an in-progress job")
|
||||||
|
check(w._eta_suffix(0, 10) == "" and w._eta_suffix(10, 10) == "",
|
||||||
|
"no ETA at 0% or 100%")
|
||||||
|
|
||||||
|
w.close()
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
tests = [
|
||||||
|
test_cache_atomic_roundtrip,
|
||||||
|
test_nms,
|
||||||
|
test_extract_dialog_options,
|
||||||
|
test_settings_roundtrip,
|
||||||
|
test_mainwindow_filter_and_jump,
|
||||||
|
]
|
||||||
|
for t in tests:
|
||||||
|
t()
|
||||||
|
print()
|
||||||
|
if _failures:
|
||||||
|
print(f"FAILED ({len(_failures)}):")
|
||||||
|
for f in _failures:
|
||||||
|
print(" -", f)
|
||||||
|
return 1
|
||||||
|
print("ALL SMOKE TESTS PASSED")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
Reference in New Issue
Block a user