diff --git a/CLAUDE.md b/CLAUDE.md index 9ab9ba8..eb3496b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -85,15 +85,18 @@ torch present. ## Architecture (as implemented) -> This reflects the actual code on disk. It is a synchronous, single-threaded GUI app -> — no worker threads. Work is organized into **projects** (`core/project.py`): a -> project folder holds `project.json` (metadata + per-project settings), `frames/` (the -> images), `detections.json` (the detection cache, at the project root — no longer a -> sidecar next to the images), and `collections/Избранное` (the favorites collection). The only -> video touch is a one-shot "Создать из ролика…" that creates a new project and decodes -> a clip into its `frames/` via the **ffmpeg CLI** (cv2.VideoCapture fallback; NOT -> PyAV). Detection runs on the GUI thread (lazily per image, or via "Детектировать -> все"). When code and this file disagree, trust the code. +> This reflects the actual code on disk. The GUI is mostly synchronous, but the +> **heavy compute (detection + restoration) runs on a background thread** so the UI +> stays responsive — see `ui/workers.py` and the "Background jobs" bullet (this reverses +> the earlier "no worker threads" rule; `processEvents` can't unfreeze a single multi- +> second `detector.detect()`/DeepMosaics call). Work is organized into **projects** +> (`core/project.py`): a project folder holds `project.json` (metadata + per-project +> settings), `frames/` (the images), `detections.json` (the detection cache, at the +> project root — no longer a sidecar next to the images), and `collections/Избранное` +> (the favorites collection). The only video touch is a one-shot "Создать из ролика…" +> that creates a new project and decodes a clip into its `frames/` via the **ffmpeg +> CLI** (cv2.VideoCapture fallback; NOT PyAV). When code and this file disagree, trust +> the code. ``` hvideotool/ @@ -103,9 +106,11 @@ hvideotool/ ├── settings_store.py # new-project DEFAULTS + last/recent projects to ~/HVideoTool/settings.json ├── ui/ │ ├── main_window.py # the whole UI: toolbar + [file list | image view | detail table] +│ ├── workers.py # Job (QRunnable): runs detect/restore off-thread, results via Qt signals │ └── image_view.py # renders an image + draws polygon/bbox overlays (QPainter); can highlight one └── core/ ├── imageio.py # unicode-safe imread/imwrite (np.fromfile + imdecode) + ├── torch_info.py # probe torch/CUDA (gather/reason/install_hint) for the device badge; no Qt ├── project.py # Project: layout (project.json/frames/detections.json/collections) + per-project settings ├── video/ │ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale @@ -132,6 +137,24 @@ hvideotool/ - `MainWindow` holds the config, builds the detector lazily via `build_detector` (cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path -> list[Detection]]` as the detection cache. +- **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 + (`tick`/`progress`/`done`/`failed`). `MainWindow._start_job(fn, total, on_tick, on_done)` + starts one (only one at a time — `_busy` guards entry points), `_finish_job`/ + `_on_job_failed` end it. `_make_detector`/`_make_restorer`, image reads, and + `engine.detect/restore` all run **inside the worker** (`_compute` is the pure + read+detect helper); the `fn` must touch NO Qt widgets — it emits plain data that the + GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables + `detector_combo`/`model_action` for the duration (they'd race the running detector). + This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics + per-call latency can't be hidden with `processEvents`). +- **Device badge.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green) + or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background + `Job` at startup (importing torch is slow, so it's off the GUI thread) → `_set_device_badge`. + Clicking (`_show_device_info`) opens a diagnostic dialog: `torch_info.reason()` explains + why CPU (no torch / CPU-only `+cpu` build / built-with-CUDA-but-no-GPU) plus + `install_hint()` (the cu121 pip command). `core/torch_info.py` is pure (no Qt). - **Projects (`core/project.py`).** `MainWindow._project` is the open `Project`; `_folder` is kept as a synonym for `project.frames_dir` so the rest of the code (navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…" @@ -182,7 +205,9 @@ hvideotool/ training/example set while inspecting detections. - **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via `build_restorer`) on the current frame's detections (computing them first if needed), - caches the result in `_restored[path]`, and shows it overlay-free. "Показать оригинал/ + **on a background job** (`_restore_current` builds an `fn` that detects-if-needed + + restores in the worker; a `tick` caches freshly-computed detections, `done` stores + `_restored[path]` + shows it). "Показать оригинал/ результат" toggles (`_showing_restored`); "Сохранить результат" writes `_restored.jpg` beside the frame. The baseline is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run @@ -191,7 +216,8 @@ hvideotool/ `cleanmosaic_img_server` (locate mosaic → run generator on the crop → feather back), ~0.3 s/frame cached on CPU vs ~7 s when it spawned a subprocess. Use the **image** model `clean_youknow_resnet_9blocks.pth` — the video model (BVDNet) is rejected per - frame (needs a neighbour). `should_cancel` is polled at entry (raises `Cancelled`). + frame (needs a neighbour). `should_cancel` (= `lambda: job.cancelled`) is polled so + "■ Стоп" stops it; the engine raises `Cancelled`, which `Job.run` reports as a clean cancel. The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…), persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE: DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the @@ -208,17 +234,16 @@ hvideotool/ detections (`_refresh_marks` projects `_results` onto row indices; per-pixel deduped so big folders stay cheap). File-list rows are tinted too (`_tag_file`): red = censorship found, green = checked & clean. Both reset on `_invalidate_results`. -- **Cancellation (cooperative, no threads).** A single "■ Стоп" toolbar action (Esc) - cancels the running long op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` - + the Stop button + the progress bar (`total=None` → indeterminate); `_request_cancel` - sets `self._cancel`; the loop-based ops (`_detect_all`, `_load_folder`) and video - extraction (its `progress` cb returns `not self._cancel`) check the flag between - `processEvents` ticks. Single-image restore passes `should_cancel=self._poll_cancel` - (which pumps `processEvents` then returns the flag) into `Restorer.restore`; only - DeepMosaics actually polls it (kills its subprocess + raises `Cancelled`) — cv2 ops are - instant. Entry points guard with `if self._busy: return` (notably `_move_to_collection`, - which mutates `_files` that `_detect_all` iterates). This keeps the synchronous, - single-threaded model — do NOT reintroduce worker threads for cancellation. +- **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, + restore) `_request_cancel` calls `self._job.cancel()`; the worker loop checks + `job.cancelled` between frames and restore polls it via `should_cancel`. The still- + synchronous loops (`_load_folder` listing, import copy, video extraction — its + `progress` cb returns `not self._cancel`) check `self._cancel` between `processEvents` + ticks. Entry points guard with `if self._busy: return` (notably `_move_to_favorites`, + which mutates `_files` that a detect-all job reads — so a snapshot/pending list is used). + `closeEvent` cancels a running job and `waitForDone(3000)` before tearing down. - `image_view.ImageView` draws the image scaled-to-fit plus overlays. Overlay visibility/threshold are applied at paint time. Selecting a row in the detail table calls `set_highlight(i)` — that detection is drawn boldly (even below threshold) and @@ -267,6 +292,10 @@ frame directly. a generic COCO model (e.g. the `yolo11n-seg.pt` in the repo root, which Ultralytics auto-downloads / is the training base), it detects people/objects and maps them to `CensorType.UNKNOWN` → purple boxes that look like noise. This was a real user trap. + **Switching to yolo/combined without a model auto-picks one** via + `MainWindow._auto_find_model()`: it scans `./models/**.pt` and matches only filenames + containing `lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; it + falls back to the "Модель…" file dialog only when nothing suitable is found. - **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low block-reconstruction residual + 2D gradient + contrast) fires on textured real footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing @@ -286,13 +315,23 @@ frame directly. `YoloDetector.__init__`). - **CUDA/torch install is environment-specific.** Don't add torch to core deps; it stays out (the `yolo` extra pulls only Ultralytics) and is installed separately. +- **CPU-only torch must not request CUDA.** A `+cpu` torch build raises "Torch not + compiled with CUDA enabled" the moment something calls `.cuda()`. Both engines guard + for this: `YoloDetector` picks `cuda` only when `torch.cuda.is_available()` (even an + explicit `yolo_device="cuda"` is downgraded to cpu); `DeepMosaicsRestorer._ensure_loaded` + forces `gpu_id="-1"` when CUDA is absent (its vendored `model_util.todevice` / + `data.im2tensor` call `.cuda()` for any `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). + So a wrong/CPU-only torch falls back to CPU instead of crashing. - **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`), otherwise it aliases a buffer that gets freed → garbage/crash. - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths. - Don't reintroduce any generative / ControlNet dependency, nor the removed video - *playback pipeline* (PyAV, worker threads, player). (The new `core/project.py` is an - on-disk layout, not that thread-based "project model".) The one allowed video touch is + *playback pipeline* (PyAV, producer/consumer worker threads, player, project session). + (The new `core/project.py` is an on-disk layout, not that thread-based "project + model".) NOTE: a **single** background `Job` thread for detect/restore (`ui/workers.py`) + IS in scope now (keeps the GUI responsive) — that's different from the rejected + multi-thread video pipeline. The one allowed video touch is `core/video/extract.py` (one-shot decode → a new project's `frames/`, behind "Создать из ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is diff --git a/README.md b/README.md index db36b07..83913f1 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,10 @@ - Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** уверенности прямо в тулбаре — удобно сравнивать. - Выбор файла весов модели кнопкой **«Модель…»**. +- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU» + показывает диагностику (почему GPU не задействован) и команды установки PyTorch с + CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU + (медленнее, но без ошибок). ## Кэш детекций diff --git a/hvideotool/core/detection/yolo.py b/hvideotool/core/detection/yolo.py index e789d85..ad6e895 100644 --- a/hvideotool/core/detection/yolo.py +++ b/hvideotool/core/detection/yolo.py @@ -54,15 +54,20 @@ class YoloDetector(Detector): "(и PyTorch с CUDA отдельно — см. README)." ) from exc - # Resolve the device: explicit override, else CUDA when available. + # Resolve the device: CUDA only when actually available, else CPU. A CPU-only + # torch build raises "Torch not compiled with CUDA enabled" if asked for cuda, + # so we never request it without a working GPU (even on an explicit override). + try: + import torch + + cuda_ok = torch.cuda.is_available() + except Exception: # noqa: BLE001 + cuda_ok = False device = self.cfg.yolo_device if device is None: - try: - import torch - - device = "cuda" if torch.cuda.is_available() else "cpu" - except Exception: # noqa: BLE001 - device = "cpu" + device = "cuda" if cuda_ok else "cpu" + elif "cuda" in str(device) and not cuda_ok: + device = "cpu" self._device = device self._model = YOLO(model_path) diff --git a/hvideotool/core/restore/deepmosaics.py b/hvideotool/core/restore/deepmosaics.py index 850e36e..3f7b2ec 100644 --- a/hvideotool/core/restore/deepmosaics.py +++ b/hvideotool/core/restore/deepmosaics.py @@ -118,6 +118,14 @@ class DeepMosaicsRestorer(Restorer): return if str(_VENDOR) not in sys.path: sys.path.insert(0, str(_VENDOR)) # so vendored `from models/util import …` resolve + + # Fall back to CPU when CUDA isn't available: DeepMosaics calls `.cuda()` + # whenever gpu_id != "-1", which raises "Torch not compiled with CUDA enabled" + # on a CPU-only torch build. + import torch # noqa: E402 + if self._gpu != "-1" and not torch.cuda.is_available(): + self._gpu = "-1" + from models import loadmodel, runmodel # type: ignore # noqa: E402 import util.image_processing as impro # type: ignore # noqa: E402 diff --git a/hvideotool/core/torch_info.py b/hvideotool/core/torch_info.py new file mode 100644 index 0000000..a45c33a --- /dev/null +++ b/hvideotool/core/torch_info.py @@ -0,0 +1,80 @@ +"""Probe the PyTorch / CUDA situation, so the UI can show a device badge. + +Pure (no Qt). ``gather()`` imports torch (slow / heavy) — call it off the GUI +thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and +how to enable the GPU. +""" + +from __future__ import annotations + +# pip index for the CUDA build (matches the README). +CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121" + + +def gather() -> dict: + """Collect torch/CUDA facts. Never raises — missing torch is a valid result.""" + info: dict = { + "installed": False, + "version": None, # torch.__version__ (e.g. "2.12.0+cpu") + "built_cuda": None, # torch.version.cuda (None on a CPU-only build) + "cuda_available": False, + "device_name": None, # the active GPU's name, if any + "import_error": None, + } + try: + import torch + except Exception as exc: # noqa: BLE001 - report any import failure, not just ImportError + info["import_error"] = str(exc) + return info + info["installed"] = True + info["version"] = getattr(torch, "__version__", None) + try: + info["built_cuda"] = torch.version.cuda + except Exception: # noqa: BLE001 + info["built_cuda"] = None + try: + info["cuda_available"] = bool(torch.cuda.is_available()) + except Exception: # noqa: BLE001 + info["cuda_available"] = False + if info["cuda_available"]: + try: + info["device_name"] = torch.cuda.get_device_name(0) + except Exception: # noqa: BLE001 + info["device_name"] = None + return info + + +def device_label(info: dict) -> str: + return "CUDA" if info.get("cuda_available") else "CPU" + + +def reason(info: dict) -> str: + """One-sentence human explanation of the current device choice.""" + if not info.get("installed"): + return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics " + "работают на CPU (классический детектор torch не требует).") + if info.get("cuda_available"): + name = info.get("device_name") or "GPU" + return f"PyTorch использует CUDA: {name}. Вычисления идут на видеокарте." + version = info.get("version") or "?" + built = info.get("built_cuda") + if not built: + return (f"Установлена CPU-сборка PyTorch ({version}) — без поддержки CUDA, " + "поэтому вычисления идут на процессоре (медленно).") + return (f"PyTorch собран с CUDA {built} ({version}), но GPU недоступен: нет " + "NVIDIA-видеокарты, не установлен/устарел драйвер, либо версия CUDA " + "несовместима с драйвером.") + + +def install_hint() -> str: + """Steps to enable the GPU (shown when running on CPU).""" + return ( + "Как включить GPU (NVIDIA):\n" + "1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n" + "2. Переустановите PyTorch со сборкой CUDA:\n\n" + " pip uninstall -y torch torchvision\n" + f" pip install torch torchvision --index-url {CUDA_WHEEL_INDEX}\n\n" + "3. Перезапустите приложение.\n\n" + "Классический детектор работает и без CUDA. На CPU детекция и расцензуривание " + "просто медленнее." + ) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index b102f4a..e0d43f4 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -25,7 +25,7 @@ from __future__ import annotations import shutil from pathlib import Path -from PySide6.QtCore import Qt +from PySide6.QtCore import Qt, QThreadPool from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut from PySide6.QtWidgets import ( QAbstractItemView, @@ -57,7 +57,6 @@ from ..core.detection.factory import build_detector from ..core.detection.types import Detection from ..core.imageio import imread_unicode, imwrite_unicode from ..core.project import PROJECT_FILE, Project -from ..core.restore.base import Cancelled from ..core.restore.factory import build_restorer from ..core.video.extract import extract_frames from ..core.video.frame import Frame @@ -65,6 +64,7 @@ from .extract_dialog import ExtractDialog from .image_view import ImageView from .marker_slider import MarkerSlider from .restore_dialog import RestoreDialog +from .workers import Job _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} _VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)" @@ -89,6 +89,11 @@ class MainWindow(QMainWindow): 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 + self._pool = QThreadPool.globalInstance() + self._job: Job | None = None # the running background job, if any + self._tick_count = 0 # throttles scrubber-mark refreshes during detect-all + self._device_info: dict | None = None # torch/CUDA probe result (for the badge) + self._probe_job: Job | None = None self.setWindowTitle("HVideoTool — инспектор детекции цензуры") self.resize(1180, 720) @@ -97,6 +102,7 @@ class MainWindow(QMainWindow): self._build_central() self._build_statusbar() self._build_menu() + self._probe_device() # determine CUDA/CPU in the background and fill the badge self.statusBar().showMessage("Создайте или откройте проект (Файл)") # ------------------------------------------------------------------ setup @@ -310,17 +316,80 @@ class MainWindow(QMainWindow): self.pos_label.setText(f"{row + 1 if row >= 0 else 0} / {n}") def _build_statusbar(self) -> None: + self.device_badge = QPushButton("⏳ устройство…") + self.device_badge.setFlat(True) + self.device_badge.setCursor(Qt.PointingHandCursor) + self.device_badge.setToolTip("Устройство вычислений (нажмите для подробностей)") + self.device_badge.clicked.connect(self._show_device_info) + self.statusBar().addPermanentWidget(self.device_badge) + self.progress = QProgressBar() self.progress.setMaximumWidth(260) self.progress.setVisible(False) self.statusBar().addPermanentWidget(self.progress) + # ------------------------------------------------------------- device badge + def _probe_device(self) -> None: + """Determine CUDA/CPU off the GUI thread (importing torch is slow).""" + from ..core import torch_info + + job = Job(lambda _job: torch_info.gather()) + self._probe_job = job # keep alive until `done` + job.signals.done.connect(self._set_device_badge) + job.signals.failed.connect(lambda _msg: self._set_device_badge(None)) + self._pool.start(job) + + def _set_device_badge(self, info: dict | None) -> None: + self._probe_job = None + self._device_info = info or {} + if self._device_info.get("cuda_available"): + name = self._device_info.get("device_name") or "GPU" + self.device_badge.setText("⚡ CUDA") + self.device_badge.setToolTip(f"Вычисления на GPU: {name} (нажмите для подробностей)") + self.device_badge.setStyleSheet("QPushButton{color:#16a085; font-weight:bold;}") + else: + self.device_badge.setText("🖥 CPU") + self.device_badge.setToolTip( + "Вычисления на CPU — нажмите, чтобы узнать почему и как включить GPU" + ) + self.device_badge.setStyleSheet("QPushButton{color:#cc8400; font-weight:bold;}") + + def _show_device_info(self) -> None: + from ..core import torch_info + + info = self._device_info if self._device_info else torch_info.gather() + cuda = bool(info.get("cuda_available")) + lines = [ + torch_info.reason(info), + "", + "Диагностика:", + f" • PyTorch: {info.get('version') or 'не установлен'}", + f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}", + f" • CUDA доступна: {'да' if cuda else 'нет'}", + ] + if info.get("device_name"): + lines.append(f" • GPU: {info['device_name']}") + if info.get("import_error"): + lines.append(f" • Ошибка импорта torch: {info['import_error']}") + + box = QMessageBox(self) + box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning) + box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU")) + box.setText("\n".join(lines)) + if not cuda: + box.setInformativeText(torch_info.install_hint()) + box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands + box.exec() + # ------------------------------------------------------------- cancellation def _begin_busy(self, total: int | None = None) -> None: """Enter a cancellable long operation. ``total=None`` => busy spinner.""" self._busy = True self._cancel = False self.stop_action.setEnabled(True) + # Disable inputs that would race a running job (they clear cache / rebuild engines). + self.detector_combo.setEnabled(False) + self.model_action.setEnabled(False) if total is None: self.progress.setRange(0, 0) # indeterminate else: @@ -331,18 +400,55 @@ class MainWindow(QMainWindow): def _end_busy(self) -> None: self._busy = False self.stop_action.setEnabled(False) + self.detector_combo.setEnabled(True) + self.model_action.setEnabled(True) self.progress.setVisible(False) self.progress.setRange(0, 100) # leave it determinate for the next user def _request_cancel(self) -> None: if self._busy: self._cancel = True + if self._job is not None: + self._job.cancel() # stops the background loop at its next check self.statusBar().showMessage("Отмена…") - def _poll_cancel(self) -> bool: - """Cancel hook for core engines: pump the UI so Стоп registers, then report.""" - QApplication.processEvents() - return self._cancel + # ------------------------------------------------------------- background jobs + def _start_job(self, fn, total: int | None, *, on_tick=None, on_done=None) -> None: + """Run ``fn(job)`` on the thread pool; marshal results back to the GUI. + + ``on_tick(payload)`` handles incremental results (GUI thread); ``on_done(result, + cancelled)`` runs when the job finishes. Only one job runs at a time (callers + guard with ``self._busy``). + """ + self._begin_busy(total) + self._tick_count = 0 + job = Job(fn) + self._job = job + if on_tick is not None: + job.signals.tick.connect(on_tick) + job.signals.progress.connect(self._on_job_progress) + job.signals.done.connect(lambda result: self._finish_job(result, on_done)) + job.signals.failed.connect(self._on_job_failed) + self._pool.start(job) + + def _on_job_progress(self, done: int, total: int, message: str) -> None: + if total > 0: + self.progress.setRange(0, total) + self.progress.setValue(done) + if message: + self.statusBar().showMessage(message) + + def _finish_job(self, result, on_done) -> None: + cancelled = self._job.cancelled if self._job is not None else False + self._job = None + self._end_busy() + if on_done is not None: + on_done(result, cancelled) + + def _on_job_failed(self, message: str) -> None: + self._job = None + self._end_busy() + QMessageBox.warning(self, "Ошибка", message) # --------------------------------------------------------------- detector def _make_detector(self): @@ -355,12 +461,34 @@ class MainWindow(QMainWindow): def _on_detector_changed(self, name: str) -> None: self._cfg.detector = name - # YOLO/combined need a model — offer to pick one if missing. + # YOLO/combined need a model. Auto-pick a known one from models/ if we have it; + # only prompt when nothing suitable is found (don't nag when the path is obvious). if name in ("yolo", "combined") and not self._cfg.model_path: - self._choose_model() + found = self._auto_find_model() + if found: + self._cfg.model_path = found + self.statusBar().showMessage(f"Модель найдена автоматически: {found}") + else: + self._choose_model() self._persist_settings() self._invalidate_results() + @staticmethod + def _auto_find_model() -> str | None: + """Find a censorship YOLO model under ./models without prompting. + + Matches LADA/mosaic weights by filename; deliberately ignores generic COCO + models (e.g. yolo11n-seg.pt) that would map objects to purple "noise". + """ + models_dir = Path.cwd() / "models" + if not models_dir.is_dir(): + return None + for p in sorted(models_dir.rglob("*.pt")): + name = p.name.lower() + if "lada" in name or "mosaic" in name: + return str(p) + return None + def _choose_model(self) -> None: start = self._cfg.model_path or str(Path.cwd() / "models") path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)") @@ -681,35 +809,44 @@ class MainWindow(QMainWindow): if self._busy: return path = Path(item.data(Qt.UserRole)) - if str(path) not in self._results: - if self._detect(path) is None: - return - self._refresh_marks() - self._save_results() # only when a detection actually ran - self._show(path) + if str(path) in self._results: + self._show(path) + return + self._detect_one(path, then_show=True) - def _detect(self, path: Path) -> list[Detection] | None: - """Run (or fetch cached) detections for one image. None on failure.""" - key = str(path) - if key in self._results: - return self._results[key] - img = imread_unicode(key) + @staticmethod + def _compute(detector, path: Path) -> list[Detection]: + """Pure read + detect for one image (runs on a worker thread; no Qt).""" + img = imread_unicode(str(path)) if img is None: - self.statusBar().showMessage(f"Не удалось прочитать: {path.name}") - return None - try: - detector = self._make_detector() - except Exception as exc: # noqa: BLE001 - surface config/model errors to the user - QMessageBox.warning(self, "Детектор недоступен", str(exc)) - return None - self.statusBar().showMessage(f"Детекция: {path.name}…") - QApplication.processEvents() + raise RuntimeError(f"Не удалось прочитать: {Path(path).name}") dets = detector.detect(Frame(image=img, index=0, pts=0.0)) dets.sort(key=lambda d: d.score, reverse=True) - self._results[key] = dets - self._tag_file(path, len(dets)) return dets + def _detect_one(self, path: Path, *, then_show: bool) -> None: + """Detect one image on a background thread, then cache/tag/show it.""" + if self._busy: + return + + def fn(job): + return (str(path), self._compute(self._make_detector(), path)) + + def done(result, cancelled): + if result is None: + return + key, dets = result + self._results[key] = dets + self._tag_file(Path(key), len(dets)) + self._refresh_marks() + self._save_results() + if then_show or self._current == Path(key): + self._show(Path(key)) + self.statusBar().showMessage(f"Детекция: {Path(key).name} — {len(dets)} обл.") + + self.statusBar().showMessage(f"Детекция: {path.name}…") + 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).""" self._current = path @@ -721,90 +858,115 @@ class MainWindow(QMainWindow): self._update_restore_actions() def _recompute_current(self) -> None: - """Toolbar/Space: (re)run the detector on the selected frame.""" + """Toolbar/Space: (re)run the detector on the selected frame (background).""" if self._current is None or self._busy: return self._results.pop(str(self._current), None) self._detector_key = None # rebuild the detector so settings changes take effect - if self._detect(self._current) is None: - return - self._refresh_marks() - self._save_results() - self._show(self._current) + self._detect_one(self._current, then_show=True) def _detect_all(self, force: bool = False) -> None: - """Detect the whole folder. ``force`` clears the cache first (full regen); - otherwise already-computed frames are skipped, so it resumes/tops-up.""" + """Detect the whole folder on a background thread. ``force`` clears the cache + first (full regen); otherwise already-computed frames are skipped (resume/top-up). + The GUI stays responsive — results stream in via per-frame ticks.""" if not self._files or self._busy: return if force: self._clear_results() - total = len(self._files) - self._begin_busy(total) - done = 0 - try: - for i, p in enumerate(self._files, 1): - self.progress.setValue(i) - self.statusBar().showMessage(f"Детекция {i}/{total}: {p.name}") - QApplication.processEvents() - if self._cancel: + pending = [p for p in self._files if str(p) not in self._results] + if not pending: + self.statusBar().showMessage("Все кадры уже посчитаны (см. «Все заново»)") + return + total = len(pending) + + def fn(job): + detector = self._make_detector() # built on the worker thread (may raise) + for i, p in enumerate(pending, 1): + if job.cancelled: break - if self._detect(p) is None: - return # detector unavailable — message already shown - done = i - if i % 50 == 0: - self._refresh_marks() # let marks appear progressively - finally: - self._end_busy() - hits = sum(1 for p in self._files if self._results.get(str(p))) - self._refresh_marks() - self._save_results() # persist progress (works for completed and cancelled runs) - if self._cancel: - self.statusBar().showMessage( - f"Отменено на {done}/{total} · детекции на {hits} картинках" - ) - else: - self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок") - if self._current is not None: - self._show(self._current) + try: + dets = self._compute(detector, p) + except RuntimeError: + continue # unreadable image — skip, keep going + job.tick((str(p), dets)) + job.progress(i, total, f"Детекция {i}/{total}: {p.name}") + return None + + def done(_result, cancelled): + hits = sum(1 for p in self._files if self._results.get(str(p))) + self._refresh_marks() + self._save_results() # persist progress (completed or cancelled) + if cancelled: + self.statusBar().showMessage(f"Отменено · детекции на {hits} картинках") + else: + self.statusBar().showMessage( + f"Готово: детекции на {hits} из {len(self._files)} картинок" + ) + if self._current is not None: + self._show(self._current) + + self._start_job(fn, total, on_tick=self._apply_detection, on_done=done) + + def _apply_detection(self, payload) -> None: + """GUI-thread handler for one streamed detect-all result.""" + key, dets = payload + self._results[key] = dets + self._tag_file(Path(key), len(dets)) + self._tick_count += 1 + if self._tick_count % 25 == 0: + self._refresh_marks() # let marks appear progressively (throttled) # ------------------------------------------------------------- restoration def _restore_current(self) -> None: - """Run the restorer on the current frame's detected regions and show it.""" + """Restore the current frame's regions on a background thread, then show it. + + Detections are computed first (in the same job) if not cached. The DeepMosaics + engine polls ``job.cancelled`` so "■ Стоп" stops it promptly.""" if self._current is None or self._busy: return - key = str(self._current) - if key not in self._results and self._detect(self._current) is None: - return - self._refresh_marks() - dets = self._results.get(key) or [] - if not dets: - self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать") - return - img = imread_unicode(key) - if img is None: - return - self._begin_busy() # indeterminate — engine drives the duration - self.statusBar().showMessage(f"Восстановление: {self._current.name}…") - QApplication.processEvents() - try: + path = self._current + key = str(path) + + def fn(job): + dets = self._results.get(key) + if dets is None: + dets = self._compute(self._make_detector(), path) + job.tick(("dets", key, dets)) # cache them on the GUI thread + if not dets: + return ("empty", key) + img = imread_unicode(key) + if img is None: + raise RuntimeError(f"Не удалось прочитать: {path.name}") restorer = self._make_restorer() - restored = restorer.restore(img, dets, should_cancel=self._poll_cancel) - except Cancelled: - self.statusBar().showMessage("Восстановление отменено") - return - except Exception as exc: # noqa: BLE001 - surface model/engine errors - QMessageBox.warning(self, "Ошибка восстановления", str(exc)) - return - finally: - self._end_busy() - self._restored[key] = restored - self._showing_restored = True - self.view.set_image(restored, []) - self._update_restore_actions() - self.statusBar().showMessage( - f"Расцензурено ({restorer.name}): {self._current.name} — {len(dets)} обл." - ) + restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled) + return ("restored", key, restored, len(dets), restorer.name) + + def tick(payload): + if payload[0] == "dets": + _, k, dets = payload + self._results[k] = dets + self._tag_file(Path(k), len(dets)) + self._refresh_marks() + + def done(result, cancelled): + if cancelled: + self.statusBar().showMessage("Восстановление отменено") + return + if result is None: + return + if result[0] == "empty": + self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать") + return + _, k, restored, n, engine = result + self._restored[k] = restored + 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.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name} — {n} обл.") + + self.statusBar().showMessage(f"Восстановление: {path.name}…") + self._start_job(fn, None, on_tick=tick, on_done=done) def _make_restorer(self): key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, @@ -1017,6 +1179,9 @@ class MainWindow(QMainWindow): self._persist_settings() def closeEvent(self, event) -> None: # noqa: N802 - Qt override + if self._job is not None: # stop a running background job before tearing down + self._job.cancel() + self._pool.waitForDone(3000) self._save_results() # persist the detection cache on exit if self._project is not None: self._project.update_from_config(self._cfg) diff --git a/hvideotool/ui/workers.py b/hvideotool/ui/workers.py new file mode 100644 index 0000000..3304a2b --- /dev/null +++ b/hvideotool/ui/workers.py @@ -0,0 +1,67 @@ +"""A tiny background-job helper so heavy work doesn't freeze the GUI. + +The app is otherwise synchronous, but a single ``detector.detect()`` (CPU YOLO) or +a DeepMosaics restore can block the GUI thread for seconds — ``processEvents`` only +runs *between* frames, not *inside* one heavy call. So detection and restoration run +on a ``QThreadPool`` thread via :class:`Job`; results come back to the GUI through +queued Qt signals. + +Contract: the job function ``fn(job)`` runs on a worker thread and may ONLY touch +plain data + the engines (no Qt widgets). It reports progress with ``job.progress``/ +``job.tick`` and checks ``job.cancelled`` to stop early. Its return value is delivered +on the GUI thread via the ``done`` signal; raising :class:`Cancelled` is reported as a +clean cancel (``done`` with ``None``), any other exception via ``failed``. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from PySide6.QtCore import QObject, QRunnable, Signal + +from ..core.restore.base import Cancelled + + +class _Signals(QObject): + progress = Signal(int, int, str) # done, total, message + tick = Signal(object) # incremental payload (delivered on GUI thread) + done = Signal(object) # final result (None if cancelled) + failed = Signal(str) # error message + + +class Job(QRunnable): + """Runs ``fn(job)`` on a thread pool, marshaling progress/result to the GUI.""" + + def __init__(self, fn: Callable[["Job"], Any]) -> None: + super().__init__() + self.setAutoDelete(False) # the GUI keeps a reference until `done`/`failed` + self.signals = _Signals() + self._fn = fn + self._cancelled = False + + # -- called from the GUI thread -- + def cancel(self) -> None: + self._cancelled = True + + @property + def cancelled(self) -> bool: + return self._cancelled + + # -- called from the worker thread by `fn` -- + def progress(self, done: int, total: int, message: str = "") -> None: + self.signals.progress.emit(done, total, message) + + def tick(self, payload: Any) -> None: + self.signals.tick.emit(payload) + + # -- thread entry point -- + def run(self) -> None: # noqa: D401 - QRunnable override + try: + result = self._fn(self) + except Cancelled: + self.signals.done.emit(None) + except Exception as exc: # noqa: BLE001 - surface engine/model errors to the GUI + self.signals.failed.emit(str(exc)) + else: + self.signals.done.emit(result)