diff --git a/CLAUDE.md b/CLAUDE.md index b7f2afe..4f3931f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -138,6 +138,19 @@ hvideotool/ ### How it works +- **Toolbar layout (`_build_toolbar`).** To avoid a long flat row spilling into Qt's + "⋯" overflow, related actions are **grouped into `QToolButton` dropdowns** (two helpers: + `_dropdown_button` = InstantPopup menu-only; `_split_button` = MenuButtonPopup, click runs + the primary action, the arrow opens related ones). Top-level groups: **"Проект ▾"** + (создать/открыть/из ролика/импортировать) · **"Детекторы: Модели (N) ▾"** (the model + picker, unchanged) · split **"Рассчитать кадр ▾"** (menu: детектировать все дозапуск/ + заново) · split **"Расцензурить кадр ▾"** (menu: расцензурить все/найденное/заново, + движок…, открыть папку результатов, «Сохранить результат…» = `save_restored_action`) · + **checkable "Показать расцензуренное"** (`toggle_restored_action`, key `R`, kept visible so + the original⇄restored state shows as a pressed button) · **"■ Стоп"** (kept visible, + reachable instantly) · a stretch spacer pushes **"Порог:"** to the right edge. The full + action list also lives in the **menu bar "Файл"** (`_build_menu`). When adding an action, + put it in the matching dropdown — don't add another flat top-level button. - `MainWindow` holds the config, builds the detector lazily via `build_detector` (cached by the **selected-model set** + conf/imgsz in `_make_detector`), and keeps `_results: dict[path -> list[Detection]]` as the detection cache. @@ -148,11 +161,28 @@ hvideotool/ `.pt` into `models/yolo//`. On project open `_ensure_models` prunes vanished paths and, if nothing is selected, default-ticks every discovered model. `build_detector` builds one `YoloDetector` per selected model (tagged `label=category`) wrapped in a - `MultiYoloDetector` that concatenates their detections (no cross-model dedup). Each - `Detection` carries `label` (category) **and `model`** (the producing `.pt` stem, tagged in - `YoloDetector` — shown as its own "Модель" column in the detail table since a category folder - may hold several models); overlay colour + table group by `Detection.display` - (label, else the CensorType) via `OverlayConfig.colors` + a stable `palette` fallback. + `MultiYoloDetector`. Each `Detection` carries `label` (category) **and `model`** (the + producing `.pt` stem, tagged in `YoloDetector` — shown as its own "Модель" column in the + detail table since a category folder may hold several models); overlay colour + table + group by `Detection.display` (label, else the CensorType) via `OverlayConfig.colors` + a + stable `palette` fallback. +- **Cross-model NMS (optional).** By default `MultiYoloDetector` just **concatenates** all + models' detections (different categories are meant to coexist). The "Модели" menu has a + checkable **"Объединять пересечения (NMS)"** (`config.cross_model_nms` + `nms_iou`, + `_on_nms_toggled`): when on, `MultiYoloDetector(nms_iou=…)` runs a greedy category-agnostic + IoU NMS (`multi._nms`/`_iou`) that drops the lower-score box of any overlapping pair — kills + the duplicate rects you get when overlapping models fire (e.g. penis + cockAndBall). It + changes the detection result, so it's part of the in-memory detector identity + (`_make_detector` key) **and** the on-disk cache key — but `cache.make_key` adds `nms_iou` + **only when NMS is on**, so the default (off) key is unchanged and an existing cache stays + valid; turning NMS on yields a distinct key (recompute) without clobbering the non-NMS + cache. Toggling also `_invalidate_results` (drops the in-memory cache). Off = concatenate. +- **Per-model display thresholds (optional).** The toolbar "Порог" spin is the global overlay + threshold; the "Модели" menu **"Пороги по моделям…"** (`_edit_model_thresholds`) stores + per-model overrides in `config.model_thresholds` (keyed by `.pt` **stem**). `ImageView` + (`set_model_thresholds`/`_eff_threshold`) draws a detection only if its score clears its + model's override, else the global threshold. **Display-only** — does not change detection or + the "с цензурой" counts (a frame is a hit if it has *any* detection, threshold-independent). - **Background jobs (`ui/workers.py`).** Detection and restoration are CPU-heavy and would freeze the GUI, so they run on a `QThreadPool` thread via `Job` (a `QRunnable` wrapping `fn(job)`); results return to the GUI through queued Qt signals @@ -164,7 +194,9 @@ hvideotool/ GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables `model_action` for the duration (it'd race the running detector). This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics per-call - latency can't be hidden with `processEvents`). + latency can't be hidden with `processEvents`). Progress messages get an **ETA** suffix + ("· осталось ~Xм Yс") from `_eta_suffix(done, total)` using `_job_start` (set in + `_start_job`, `time.monotonic`) — average-rate estimate, blank at 0 %/100 %. - **Device badge + CUDA diagnostics.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green) or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background `Job` at startup (it imports torch AND shells out to `nvidia-smi`, so it's off @@ -204,7 +236,10 @@ hvideotool/ Switching the model clears the cache (`_choose_model` → `_invalidate_results`). - **Detection cache (persisted).** `_results` is mirrored to the project's `detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by - **basename** so it survives moving the project). `cache.save_results`/`load_results` + **basename** so it survives moving the project). Both `detections.json` and + `project.json` are written **atomically** (sibling `.tmp` + `os.replace`, see + `cache._atomic_write_text` and `Project.save`) — a crash mid-write can't corrupt/truncate + a large cache (29k entries) and lose all detection work. `cache.save_results`/`load_results` take the cache file and the image `base_dir` (= `project.frames_dir`) separately, since the cache lives at the project root, not next to the images. It's tagged with the detector identity (`_results_key` = detector + model + conf/imgsz); on open, @@ -229,13 +264,30 @@ hvideotool/ **locates the mosaic itself** — so restoration is **fully decoupled from detection**: no detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on - a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать - оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes - `_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново", - `_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every - frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`, - basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the - per-frame engine **skips frames already in `restored/`** unless `force` (resume). The + a background job** (`_restore_current`; `done` stores `_restored[path]`, **auto-saves it + to the project's `restored/`** (so a single restore persists like the batch, not just in + memory), and shows it). **"Показать расцензуренное/оригинал" (`_toggle_restored`, key `R`) + is a GLOBAL view mode** (`_showing_restored`): when on, `_show` displays each frame's + restored version if one exists — loaded lazily from memory `_restored` **or disk + `restored/`** via `_restored_image_for` (so the whole batch result is browsable, not just + the last frame) — else falls back to the original; overlays are hidden on restored. The + mode persists across navigation (reset to off on project open); a batch/single restore + auto-switches it on. `_has_restored`/`_restored_disk_path` are the cheap (no-decode) + existence checks driving the toggle/save enabled-state. "Сохранить результат" additionally + exports `_restored.jpg` beside the frame (an explicit one-off export, via + `_restored_image_for`). "Открыть папку результатов" opens `restored/` in Explorer. + **Batch ("Расцензурить все" / "Все заново" / + "Расцензурить найденное", `_restore_all(force, only_detected)`)** mirrors `_detect_all`: + a single background job restores frames and writes results to the project's + **`restored/`** dir (`Project.restored_dir`, basename-mirrored, kept OUT of `frames/` so + outputs aren't re-listed/re-restored); the per-frame engine **skips frames already in + `restored/`** unless `force` (resume). **`only_detected`** ("Расцензурить найденное") + uses the YOLO detection cache to skip frames known clean: per-frame restores only the + frames with detections; the temporal engine restricts the run to the contiguous span + `[first hit … last hit]` (recurrence needs continuity). It's a separate, *faster* action + — NOT the default — because LADA misses some mosaic (esp. anime), so it can miss + censorship YOLO didn't flag; "Расцензурить все" stays the thorough option. Returns early + (status hint) if detection isn't computed or no frame has a detection. The engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy, cached on the instance): @@ -269,11 +321,19 @@ hvideotool/ (set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in `restore/factory.build_restorer`. - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys - `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` - ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, - `_step_hit` scans `_results` for the next non-empty frame). The slider and file list - are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all - navigation ultimately drives `file_list.setCurrentRow`. The scrubber is a custom + `,`/`.`), a scrubber `frame_slider` across the whole sequence, a clickable `pos_label` + (a flat `QPushButton` "row / n" → `_jump_to_frame`, a "go to frame N" `QInputDialog` — + needed on 29k-frame projects where the scrubber is ~50 frames/px), and jump-to-detection + (◀ детекция / детекция ▶, keys `[`/`]`, `_step_hit` scans `_results` for the next + non-empty frame). The slider and file list are kept in sync via `_update_nav` guarded by + `_nav_sync` (avoids signal loops); all navigation ultimately drives + `file_list.setCurrentRow`. `_step` (◀ ▶) **skips rows hidden by the filter**. +- **File-list filter (`filter_combo`, left pane).** A combobox above the list shows only a + subset of the (possibly huge) frame list: Все · С цензурой · Чистые · Не рассчитано · + Расцензуренные · Без расцензуривания. `_filter_mode` + `_row_matches_filter(path)` (over + `_results`/`_row_restored`); `_on_filter_changed` re-labels (each `_relabel_row` calls + `item.setHidden(...)`) and jumps off a now-hidden current row. Pure show/hide — doesn't + touch `_files`/cache. The scrubber is a custom `MarkerSlider` (`ui/marker_slider.py`) that paints **two mark layers**: cyan ticks (upper half) at frames with detections (`_refresh_marks` projects `_results`) and **green ticks (lower half) at restored frames** (`_refresh_restored_marks` scans @@ -281,8 +341,11 @@ hvideotool/ per-frame); per-pixel deduped so big folders stay cheap. Under the scrubber a **progress summary** `stats_label` reads "Кадров: N · детектировано: D/N (с цензурой: H) · расцензурено: R/N" (`_update_counts_label`, cheap counts; `_restored_count` - cached by `_refresh_restored_marks`). File-list rows are tinted too (`_tag_file`): - red = censorship found, green = checked & clean. Both reset on `_invalidate_results`. + cached by `_refresh_restored_marks`). File-list rows are labelled too via a single + `_relabel_row` (used by `_tag_file`/`_relabel_all`): tint red = censorship found, green = + checked & clean; a trailing **✓** marks frames with a restored version (`_row_restored`, + populated by `_refresh_restored_marks`). The ✓ is independent of detection — it survives + `_clear_results`. Detection tints reset on `_invalidate_results`. - **Cancellation (cooperative).** A single "■ Стоп" toolbar action (Esc) cancels the running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button + the progress bar (`total=None` → indeterminate). For **background jobs** (detection, @@ -322,7 +385,12 @@ python -m hvideotool # reopen the last project (or create/open python -m hvideotool "C:\path\to\MyProject" --model models\lada_mosaic_detection_model_v4_accurate.pt ``` -No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a +No formal test suite, but there's a **headless smoke test**: `scripts/smoke_test.py` +(run with `QT_QPA_PLATFORM=offscreen` + `PYTHONIOENCODING=utf-8`) covers the pure core +(atomic cache round-trip, cross-model NMS, list-filter predicate, ETA formatting, +extract-dialog options, per-project settings round-trip) and an offscreen `MainWindow` +build on a throwaway project — no torch/weights (detections are injected into `_results`). +Exits non-zero on failure; run it after touching core/UI plumbing. Ad-hoc check: build a `MainWindow`, `Project.create(tmp)` + copy a few images into `frames/`, `_open_project(project)`, drive `file_list.setCurrentRow(...)`, and read `detail_table` / `detail_header`. Or run `build_detector(config).detect(...)` on a @@ -386,4 +454,7 @@ frame directly. the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is ~10× faster than every-frame; `-hwaccel` does NOT help (GPU transfer overhead). Use ffmpeg/cv2, not PyAV, and keep it synchronous. Decoding every frame is the inherent - cost — the speed lever is decoding *fewer* frames (keyframes). + cost — the speed lever is decoding *fewer* frames (keyframes). `ExtractDialog.options()` + returns `(keyframes_only, step, max_dim, jpg_quality)`; **jpg_quality** (1–100, default 92, + via `-q:v` `_quality_to_qscale` / cv2 `IMWRITE_JPEG_QUALITY`) trades quality for a bit of + encode speed + smaller files. Default sampling is **every frame** (step=1, not keyframes). diff --git a/hvideotool/config.py b/hvideotool/config.py index 72f3a12..425847f 100644 --- a/hvideotool/config.py +++ b/hvideotool/config.py @@ -56,6 +56,15 @@ class AppConfig: # selected model and merges results — see core/detection/multi.MultiYoloDetector. detector_models: list[str] = field(default_factory=list) 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 ("расцензурить") --- restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video" diff --git a/hvideotool/core/detection/cache.py b/hvideotool/core/detection/cache.py index b28d656..35c7b2a 100644 --- a/hvideotool/core/detection/cache.py +++ b/hvideotool/core/detection/cache.py @@ -17,6 +17,7 @@ project folder. from __future__ import annotations import json +import os from pathlib import Path from .types import Detection @@ -24,16 +25,43 @@ from .types import Detection _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. 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), "yolo_conf": round(float(yolo_conf), 4), "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: @@ -48,9 +76,7 @@ def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection] } try: cache_file.parent.mkdir(parents=True, exist_ok=True) - cache_file.write_text( - json.dumps(payload, ensure_ascii=False), encoding="utf-8" - ) + _atomic_write_text(cache_file, json.dumps(payload, ensure_ascii=False)) return True except OSError: return False diff --git a/hvideotool/core/detection/factory.py b/hvideotool/core/detection/factory.py index ca7a13c..fcf9b7c 100644 --- a/hvideotool/core/detection/factory.py +++ b/hvideotool/core/detection/factory.py @@ -31,6 +31,8 @@ def build_detector(config: AppConfig) -> Detector: from .multi import MultiYoloDetector from .yolo import YoloDetector - return MultiYoloDetector([ - YoloDetector(m, config.detection, label=category_of(m)) for m in models - ]) + nms_iou = config.nms_iou if config.cross_model_nms else None + return MultiYoloDetector( + [YoloDetector(m, config.detection, label=category_of(m)) for m in models], + nms_iou=nms_iou, + ) diff --git a/hvideotool/core/detection/multi.py b/hvideotool/core/detection/multi.py index 8a6bbf0..b4ff546 100644 --- a/hvideotool/core/detection/multi.py +++ b/hvideotool/core/detection/multi.py @@ -15,10 +15,13 @@ from .types import Detection 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: raise ValueError("MultiYoloDetector requires at least one detector") 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 def name(self) -> str: @@ -28,4 +31,33 @@ class MultiYoloDetector(Detector): out: list[Detection] = [] for d in self._detectors: out.extend(d.detect(frame)) + if self._nms_iou is not None: + out = _nms(out, self._nms_iou) 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 diff --git a/hvideotool/core/project.py b/hvideotool/core/project.py index ff81747..892a3bc 100644 --- a/hvideotool/core/project.py +++ b/hvideotool/core/project.py @@ -19,6 +19,7 @@ threshold, restore engine). Global ``settings.json`` only seeds the defaults for from __future__ import annotations import json +import os from dataclasses import dataclass, field from datetime import datetime from pathlib import Path @@ -38,6 +39,9 @@ _SETTING_KEYS = ( "detector", "detector_models", "default_threshold", + "model_thresholds", + "cross_model_nms", + "nms_iou", "restorer", "dm_dir", "dm_model", @@ -139,9 +143,19 @@ class Project: "settings": self.settings, } self.root.mkdir(parents=True, exist_ok=True) - self.project_file.write_text( - json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8" - ) + # Crash-safe write (sibling .tmp + atomic os.replace): never truncate a good + # 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 def apply_to_config(self, cfg: AppConfig) -> None: diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py index 86fe642..238b3a9 100644 --- a/hvideotool/settings_store.py +++ b/hvideotool/settings_store.py @@ -37,6 +37,14 @@ def apply(config: AppConfig) -> None: config.detector_models = [str(m) for m in data["detector_models"]] if "threshold" in data: 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"): config.restorer = data["restorer"] for key in ("dm_dir", "dm_model", "dm_gpu"): @@ -53,6 +61,9 @@ def save(config: AppConfig) -> None: detector=config.detector, detector_models=list(config.detector_models), 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, dm_dir=config.dm_dir, dm_model=config.dm_model, diff --git a/hvideotool/ui/extract_dialog.py b/hvideotool/ui/extract_dialog.py index 32aa8a8..6d01bca 100644 --- a/hvideotool/ui/extract_dialog.py +++ b/hvideotool/ui/extract_dialog.py @@ -41,13 +41,23 @@ class ExtractDialog(QDialog): self.max_dim.setValue(0) 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.addRow("Режим:", self.mode) form.addRow("Брать каждый N-й кадр:", self.step) form.addRow("Макс. сторона, px:", self.max_dim) + form.addRow("Качество JPEG:", self.quality) hint = QLabel( "Ключевые кадры — в разы быстрее (декодируются только I-кадры),\n" - "но реже по времени. Даунскейл уменьшает файлы и нагрузку на диск." + "но реже по времени. Даунскейл и меньшее качество уменьшают файлы\n" + "и нагрузку на диск (качество чуть ускоряет кодирование)." ) hint.setWordWrap(True) form.addRow(hint) @@ -61,6 +71,11 @@ class ExtractDialog(QDialog): def _sync(self) -> None: self.step.setEnabled(not self.mode.currentData()) - def options(self) -> tuple[bool, int, int]: - """Return (keyframes_only, step, max_dim).""" - return bool(self.mode.currentData()), self.step.value(), self.max_dim.value() + def options(self) -> tuple[bool, int, int, int]: + """Return (keyframes_only, step, max_dim, jpg_quality).""" + return ( + bool(self.mode.currentData()), + self.step.value(), + self.max_dim.value(), + self.quality.value(), + ) diff --git a/hvideotool/ui/image_view.py b/hvideotool/ui/image_view.py index 1409f9f..6a2837b 100644 --- a/hvideotool/ui/image_view.py +++ b/hvideotool/ui/image_view.py @@ -27,6 +27,7 @@ class ImageView(QWidget): self._qimage: QImage | None = None self._dets: list[Detection] = [] self._threshold = 0.0 + self._model_thresholds: dict[str, float] = {} # model stem -> override threshold self._highlight: int | None = None self.setMinimumSize(480, 360) @@ -47,6 +48,15 @@ class ImageView(QWidget): self._threshold = threshold 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: self._highlight = index self.update() @@ -82,8 +92,9 @@ class ImageView(QWidget): painter.setRenderHint(QPainter.Antialiasing, True) for i, d in enumerate(self._dets): highlighted = i == self._highlight - # A highlighted detection is always drawn; others respect the threshold. - if not highlighted and d.score < self._threshold: + # A highlighted detection is always drawn; others respect the threshold + # (per-model override if any, else the global one). + if not highlighted and d.score < self._eff_threshold(d): continue dim = self._highlight is not None and not highlighted self._draw_detection(painter, d, ox, oy, scale, highlighted, dim) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index 764770c..c67c2b5 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -25,6 +25,7 @@ from __future__ import annotations import contextlib import os import shutil +import time from pathlib import Path from PySide6.QtCore import Qt, QThreadPool @@ -32,9 +33,12 @@ from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcu from PySide6.QtWidgets import ( QAbstractItemView, QApplication, + QComboBox, QDialog, + QDialogButtonBox, QDoubleSpinBox, QFileDialog, + QFormLayout, QHBoxLayout, QInputDialog, QLabel, @@ -46,6 +50,7 @@ from PySide6.QtWidgets import ( QPlainTextEdit, QProgressBar, QPushButton, + QSizePolicy, QSplitter, QTableWidget, QTableWidgetItem, @@ -89,7 +94,9 @@ class MainWindow(QMainWindow): self._restorer_key = None self._restored: dict[str, object] = {} # path -> restored image (BGR ndarray) 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._filter_mode = "all" # file-list filter (see filter_combo) self._nav_sync = False # guard against slider<->list signal loops self._busy = False # a long operation is running 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("Расцензурить все (дозапуск)", lambda: self._restore_all(False)) 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.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M") file_menu.addSeparator() @@ -134,14 +143,19 @@ class MainWindow(QMainWindow): def _build_toolbar(self) -> None: tb = self.addToolBar("Главная") tb.setMovable(False) + tb.setToolButtonStyle(Qt.ToolButtonTextOnly) - tb.addAction(QAction("Создать проект…", self, triggered=self._create_project)) - tb.addAction(QAction("Открыть проект…", self, triggered=self._open_project_dialog)) - from_video = QAction("Создать из ролика…", self, triggered=self._create_from_video) - from_video.setToolTip("Разложить видео на кадры в новый проект и открыть его") - tb.addAction(from_video) + # --- Проект: rarely-touched session actions collapsed into one dropdown. + project_btn = self._dropdown_button("Проект ▾", "Действия с проектом") + m = project_btn.menu() + m.addAction("Создать проект…", self._create_project) + 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() + # --- Детекторы: active YOLO models picker (unchanged). tb.addWidget(QLabel(" Детекторы: ")) self._models_menu = QMenu(self) self.models_button = QToolButton() @@ -150,41 +164,62 @@ class MainWindow(QMainWindow): self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)") tb.addWidget(self.models_button) self._rebuild_models_menu() - 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.setToolTip("Отменить текущую операцию (Esc)") self.stop_action.setEnabled(False) tb.addAction(self.stop_action) - tb.addSeparator() - restore = QAction("Расцензурить кадр", self, triggered=self._restore_current) - restore.setToolTip("Восстановить найденные области на текущем кадре") - tb.addAction(restore) - restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False)) - restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)") - tb.addAction(restore_all) - restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True)) - restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)") - tb.addAction(restore_regen) - self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored) - self.toggle_restored_action.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() + # Push the threshold control to the right edge. + spacer = QWidget() + spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred) + tb.addWidget(spacer) tb.addWidget(QLabel(" Порог: ")) self.threshold_spin = QDoubleSpinBox() self.threshold_spin.setRange(0.0, 1.0) @@ -193,12 +228,47 @@ class MainWindow(QMainWindow): self.threshold_spin.valueChanged.connect(self._on_threshold_changed) 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: self.file_list = QListWidget() self.file_list.setSelectionMode(QAbstractItemView.ExtendedSelection) # multi-select for moves self.file_list.currentItemChanged.connect(self._on_file_selected) 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. move_btn = QPushButton("★ В избранное") move_btn.setToolTip("Переместить выбранные кадры в избранное проекта (Ctrl+M)") @@ -208,6 +278,7 @@ class MainWindow(QMainWindow): left_layout = QVBoxLayout(left) left_layout.setContentsMargins(4, 4, 4, 4) left_layout.setSpacing(4) + left_layout.addWidget(self.filter_combo) left_layout.addWidget(self.file_list, 1) left_layout.addWidget(move_btn) @@ -269,9 +340,14 @@ class MainWindow(QMainWindow): self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией") 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.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.setToolTip("Предыдущий кадр с детекцией ([)") @@ -297,8 +373,13 @@ class MainWindow(QMainWindow): n = self.file_list.count() if n == 0: return - row = max(0, min(n - 1, self.file_list.currentRow() + delta)) - self.file_list.setCurrentRow(row) + # Skip rows hidden by the filter so prev/next walk only the visible frames. + 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: """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: if self._nav_sync: return @@ -469,6 +562,7 @@ class MainWindow(QMainWindow): """ self._begin_busy(total) self._tick_count = 0 + self._job_start = time.monotonic() # for ETA in the progress messages job = Job(fn) self._job = job if on_tick is not None: @@ -483,7 +577,26 @@ class MainWindow(QMainWindow): self.progress.setRange(0, total) self.progress.setValue(done) 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: cancelled = self._job.cancelled if self._job is not None else False @@ -500,7 +613,10 @@ class MainWindow(QMainWindow): # --------------------------------------------------------------- detector def _make_detector(self): 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: self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors self._detector_key = key @@ -540,11 +656,74 @@ class MainWindow(QMainWindow): act.setChecked(e.path in selected) act.toggled.connect(lambda on, p=e.path: self._on_model_toggled(p, on)) 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._open_models_dir) self._models_menu.addAction("Обновить список", self._rebuild_models_menu) 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: n = len([m for m in self._cfg.detector_models if Path(m).is_file()]) self.models_button.setText(f"Модели ({n}) ▾") @@ -592,6 +771,16 @@ class MainWindow(QMainWindow): with contextlib.suppress(OSError, AttributeError): 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: """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.blockSignals(False) self.view.set_threshold(self._cfg.default_threshold) + self.view.set_model_thresholds(self._cfg.model_thresholds) def _persist_settings(self) -> None: """Save settings to the global defaults and (if open) into the project.""" @@ -795,7 +985,7 @@ class MainWindow(QMainWindow): dialog = ExtractDialog(self) if dialog.exec() != QDialog.Accepted: return - keyframes_only, step, max_dim = dialog.options() + keyframes_only, step, max_dim, jpg_quality = dialog.options() video = Path(path) root = video.parent / f"{video.stem}_frames" if Project.is_project(root): @@ -821,7 +1011,7 @@ class MainWindow(QMainWindow): try: saved = extract_frames( 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: 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) self._files = files self._results.clear() + self._row_restored.clear() + self._showing_restored = False # start a project in original-view mode self._current = None self.file_list.blockSignals(True) @@ -926,7 +1118,7 @@ class MainWindow(QMainWindow): return key, dets = result self._results[key] = dets - self._tag_file(Path(key), len(dets)) + self._tag_file(Path(key)) self._refresh_marks() self._save_results() if then_show or self._current == Path(key): @@ -937,13 +1129,22 @@ class MainWindow(QMainWindow): self._start_job(fn, None, on_done=done) 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._showing_restored = False - img = imread_unicode(str(path)) dets = self._results.get(str(path)) # None => not yet computed - self.view.set_image(img, dets or []) - self._fill_detail_table(path, img, dets) + 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._fill_detail_table(path, img, dets) self._update_restore_actions() def _recompute_current(self) -> None: @@ -1000,7 +1201,7 @@ class MainWindow(QMainWindow): """GUI-thread handler for one streamed detect-all result.""" key, dets = payload 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 not self._showing_restored and self._current is not None and str(self._current) == key: self._show(self._current) @@ -1036,28 +1237,61 @@ class MainWindow(QMainWindow): return _, k, restored, engine = result 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: self._showing_restored = True - self.view.set_image(restored, []) - self._update_restore_actions() - self._refresh_restored_marks() - self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}") + self._show(self._current) + self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}{saved_to}") self.statusBar().showMessage(f"Восстановление: {path.name}…") self._start_job(fn, None, on_done=done) - def _restore_all(self, force: bool = False) -> None: - """Restore every frame on a background thread, writing results to ``restored/``. + def _restore_all(self, force: bool = False, only_detected: bool = False) -> None: + """Restore frames on a background thread, writing results to ``restored/``. DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame engine skips frames already restored (resume) unless ``force``. The temporal - engine (DeepMosaics-video) runs the whole contiguous sequence in order via + engine (DeepMosaics-video) runs a contiguous sequence in order via ``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: return files = list(self._files) # snapshot — favorites/move mutate self._files - total = len(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) + out_dir = self._project.restored_dir out_dir.mkdir(parents=True, exist_ok=True) @@ -1067,6 +1301,7 @@ class MainWindow(QMainWindow): def fn(job): restorer = self._make_restorer() # built on the worker (may raise) frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads + state = {"done": 0} # frames processed (for the progress bar) def get_frame(i): img = frame_cache.get(i) @@ -1079,42 +1314,48 @@ class MainWindow(QMainWindow): frame_cache[i] = img return img - def emit(i, restored): - imwrite_unicode(str(out_path(files[i])), restored) - job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}") + def emit(i, restored, *, verb="Расцензуривание"): + if restored is not None: + imwrite_unicode(str(out_path(files[i])), restored) + state["done"] += 1 + job.progress(state["done"], total, f"{verb} {state['done']}/{total}: {files[i].name}") if restorer.temporal: + start = span[0] if span else 0 + end = span[1] if span else len(files) - 1 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: - for i, p in enumerate(files): + indices = hits if only_detected else range(len(files)) + for i in indices: if job.cancelled: break + p = files[i] 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 emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled)) frame_cache.pop(i, None) # per-frame: don't accumulate return None def done(_result, cancelled): - if self._current is not None: # live-preview the current frame's result, if any - rp = out_path(self._current) - if rp.is_file(): - img = imread_unicode(str(rp)) - if img is not None: - self._restored[str(self._current)] = img - self._showing_restored = True - self.view.set_image(img, []) - self._update_restore_actions() self._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( "Расцензуривание отменено" 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) def _make_restorer(self): @@ -1133,30 +1374,58 @@ class MainWindow(QMainWindow): self._restorer_key = None # rebuild on next restore 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: - if self._current is None or str(self._current) not in self._restored: - return + """Flip the global view mode between original and restored, then re-show.""" self._showing_restored = not self._showing_restored - key = str(self._current) - if self._showing_restored: - self.view.set_image(self._restored[key], []) + if self._current is not None: + self._show(self._current) 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: - has = self._current is not None and str(self._current) in self._restored - self.toggle_restored_action.setEnabled(has) + any_restored = bool(self._row_restored) or bool(self._restored) + 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( - "Показать оригинал" 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: - 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 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}") else: QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.") @@ -1227,24 +1496,77 @@ class MainWindow(QMainWindow): _TINT_HIT = QColor(200, 80, 80, 70) _TINT_CLEAN = QColor(90, 160, 90, 50) - def _set_row_tag(self, item: QListWidgetItem, count: int) -> None: - base = item.data(Qt.UserRole + 1) - item.setText(f"{base} · {count}" if count else f"{base} · —") - item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN) + def _relabel_row(self, item: QListWidgetItem) -> None: + """Set a row's text/tint/tooltip from its detection + restoration state. - def _tag_file(self, path: Path, count: int) -> None: + Text: ``name · ✓`` — 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()): item = self.file_list.item(i) if item.data(Qt.UserRole) == str(path): - self._set_row_tag(item, count) + self._relabel_row(item) return # ------------------------------------------------------------- result cache def _results_key(self) -> dict: """Detector identity used to tag/validate the on-disk detection cache.""" d = self._cfg.detection + nms_iou = self._cfg.nms_iou if self._cfg.cross_model_nms else None 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: @@ -1265,20 +1587,16 @@ class MainWindow(QMainWindow): if not cached: return 0 self._results = cached - for i in range(self.file_list.count()): - 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])) + self._relabel_all() return len(cached) 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() - for i in range(self.file_list.count()): - item = self.file_list.item(i) - item.setText(item.data(Qt.UserRole + 1)) - item.setBackground(QBrush()) + self._relabel_all() self._refresh_marks() def _refresh_marks(self) -> None: @@ -1313,13 +1631,17 @@ class MainWindow(QMainWindow): 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) rows: set[int] = set() + paths: set[str] = set() if stems or mem: for i in range(self.file_list.count()): - fp = Path(self.file_list.item(i).data(Qt.UserRole)) - if fp.stem in stems or str(fp) in mem: + sp = self.file_list.item(i).data(Qt.UserRole) + if Path(sp).stem in stems or sp in mem: rows.add(i) + paths.add(sp) + self._row_restored = paths self._restored_count = len(rows) self.frame_slider.set_restored_marks(rows) + self._relabel_all() # show/refresh the ✓ markers in the file list self._update_counts_label() def _fill_detail_table(self, path: Path, img, dets: list[Detection] | None) -> None: diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py new file mode 100644 index 0000000..348b330 --- /dev/null +++ b/scripts/smoke_test.py @@ -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())