Enhance HVideoTool with background processing and device detection: implemented a background job system for detection and restoration tasks, updated the UI to display CUDA/CPU status, and improved device diagnostics. Documentation in README and CLAUDE.md reflects these changes.

This commit is contained in:
Leonid Pershin
2026-06-07 05:24:06 +03:00
parent e27dfdf518
commit cc518cc3e6
7 changed files with 497 additions and 129 deletions
+63 -24
View File
@@ -85,15 +85,18 @@ torch present.
## Architecture (as implemented) ## Architecture (as implemented)
> This reflects the actual code on disk. It is a synchronous, single-threaded GUI app > This reflects the actual code on disk. The GUI is mostly synchronous, but the
> — no worker threads. Work is organized into **projects** (`core/project.py`): a > **heavy compute (detection + restoration) runs on a background thread** so the UI
> project folder holds `project.json` (metadata + per-project settings), `frames/` (the > stays responsive — see `ui/workers.py` and the "Background jobs" bullet (this reverses
> images), `detections.json` (the detection cache, at the project root — no longer a > the earlier "no worker threads" rule; `processEvents` can't unfreeze a single multi-
> sidecar next to the images), and `collections/Избранное` (the favorites collection). The only > second `detector.detect()`/DeepMosaics call). Work is organized into **projects**
> video touch is a one-shot "Создать из ролика…" that creates a new project and decodes > (`core/project.py`): a project folder holds `project.json` (metadata + per-project
> a clip into its `frames/` via the **ffmpeg CLI** (cv2.VideoCapture fallback; NOT > settings), `frames/` (the images), `detections.json` (the detection cache, at the
> PyAV). Detection runs on the GUI thread (lazily per image, or via "Детектировать > project root — no longer a sidecar next to the images), and `collections/Избранное`
> все"). When code and this file disagree, trust the code. > (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/ hvideotool/
@@ -103,9 +106,11 @@ hvideotool/
├── settings_store.py # new-project DEFAULTS + last/recent projects to ~/HVideoTool/settings.json ├── settings_store.py # new-project DEFAULTS + last/recent projects to ~/HVideoTool/settings.json
├── ui/ ├── ui/
│ ├── main_window.py # the whole UI: toolbar + [file list | image view | detail table] │ ├── 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 │ └── image_view.py # renders an image + draws polygon/bbox overlays (QPainter); can highlight one
└── core/ └── core/
├── imageio.py # unicode-safe imread/imwrite (np.fromfile + imdecode) ├── 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 ├── project.py # Project: layout (project.json/frames/detections.json/collections) + per-project settings
├── video/ ├── video/
│ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ ├── 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` - `MainWindow` holds the config, builds the detector lazily via `build_detector`
(cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path (cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path
-> list[Detection]]` as the detection cache. -> 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`; - **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 `_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: "Создать проект…" (navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…"
@@ -182,7 +205,9 @@ hvideotool/
training/example set while inspecting detections. training/example set while inspecting detections.
- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via - **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via
`build_restorer`) on the current frame's detections (computing them first if needed), `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 результат" toggles (`_showing_restored`); "Сохранить результат" writes
`<stem>_restored.jpg` beside the frame. The baseline `<stem>_restored.jpg` beside the frame. The baseline
is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run 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), `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** ~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 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` (Файл → Движок восстановления…), The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…),
persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE: persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE:
DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the 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 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 = so big folders stay cheap). File-list rows are tinted too (`_tag_file`): red =
censorship found, green = checked & clean. Both reset on `_invalidate_results`. censorship found, green = checked & clean. Both reset on `_invalidate_results`.
- **Cancellation (cooperative, no threads).** A single "■ Стоп" toolbar action (Esc) - **Cancellation (cooperative).** A single "■ Стоп" toolbar action (Esc) cancels the
cancels the running long op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button +
+ the Stop button + the progress bar (`total=None` → indeterminate); `_request_cancel` the progress bar (`total=None` → indeterminate). For **background jobs** (detection,
sets `self._cancel`; the loop-based ops (`_detect_all`, `_load_folder`) and video restore) `_request_cancel` calls `self._job.cancel()`; the worker loop checks
extraction (its `progress` cb returns `not self._cancel`) check the flag between `job.cancelled` between frames and restore polls it via `should_cancel`. The still-
`processEvents` ticks. Single-image restore passes `should_cancel=self._poll_cancel` synchronous loops (`_load_folder` listing, import copy, video extraction — its
(which pumps `processEvents` then returns the flag) into `Restorer.restore`; only `progress` cb returns `not self._cancel`) check `self._cancel` between `processEvents`
DeepMosaics actually polls it (kills its subprocess + raises `Cancelled`) — cv2 ops are ticks. Entry points guard with `if self._busy: return` (notably `_move_to_favorites`,
instant. Entry points guard with `if self._busy: return` (notably `_move_to_collection`, which mutates `_files` that a detect-all job reads — so a snapshot/pending list is used).
which mutates `_files` that `_detect_all` iterates). This keeps the synchronous, `closeEvent` cancels a running job and `waitForDone(3000)` before tearing down.
single-threaded model — do NOT reintroduce worker threads for cancellation.
- `image_view.ImageView` draws the image scaled-to-fit plus overlays. Overlay - `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 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 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 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 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. `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 - **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low
block-reconstruction residual + 2D gradient + contrast) fires on textured real block-reconstruction residual + 2D gradient + contrast) fires on textured real
footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing
@@ -286,13 +315,23 @@ frame directly.
`YoloDetector.__init__`). `YoloDetector.__init__`).
- **CUDA/torch install is environment-specific.** Don't add torch to core deps; it - **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. 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`), - **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`),
otherwise it aliases a buffer that gets freed → garbage/crash. otherwise it aliases a buffer that gets freed → garbage/crash.
- **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images —
`cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths. `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths.
- Don't reintroduce any generative / ControlNet dependency, nor the removed video - Don't reintroduce any generative / ControlNet dependency, nor the removed video
*playback pipeline* (PyAV, worker threads, player). (The new `core/project.py` is an *playback pipeline* (PyAV, producer/consumer worker threads, player, project session).
on-disk layout, not that thread-based "project model".) The one allowed video touch is (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 "Создать `core/video/extract.py` (one-shot decode → a new project's `frames/`, behind "Создать
из ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by из ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by
the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is
+4
View File
@@ -67,6 +67,10 @@
- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** - Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога**
уверенности прямо в тулбаре — удобно сравнивать. уверенности прямо в тулбаре — удобно сравнивать.
- Выбор файла весов модели кнопкой **«Модель…»**. - Выбор файла весов модели кнопкой **«Модель…»**.
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU
(медленнее, но без ошибок).
## Кэш детекций ## Кэш детекций
+12 -7
View File
@@ -54,15 +54,20 @@ class YoloDetector(Detector):
"(и PyTorch с CUDA отдельно — см. README)." "(и PyTorch с CUDA отдельно — см. README)."
) from exc ) 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 device = self.cfg.yolo_device
if device is None: if device is None:
try: device = "cuda" if cuda_ok else "cpu"
import torch elif "cuda" in str(device) and not cuda_ok:
device = "cpu"
device = "cuda" if torch.cuda.is_available() else "cpu"
except Exception: # noqa: BLE001
device = "cpu"
self._device = device self._device = device
self._model = YOLO(model_path) self._model = YOLO(model_path)
+8
View File
@@ -118,6 +118,14 @@ class DeepMosaicsRestorer(Restorer):
return return
if str(_VENDOR) not in sys.path: if str(_VENDOR) not in sys.path:
sys.path.insert(0, str(_VENDOR)) # so vendored `from models/util import …` resolve 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 from models import loadmodel, runmodel # type: ignore # noqa: E402
import util.image_processing as impro # type: ignore # noqa: E402 import util.image_processing as impro # type: ignore # noqa: E402
+80
View File
@@ -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 детекция и расцензуривание "
"просто медленнее."
)
+263 -98
View File
@@ -25,7 +25,7 @@ from __future__ import annotations
import shutil import shutil
from pathlib import Path 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.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
from PySide6.QtWidgets import ( from PySide6.QtWidgets import (
QAbstractItemView, QAbstractItemView,
@@ -57,7 +57,6 @@ from ..core.detection.factory import build_detector
from ..core.detection.types import Detection from ..core.detection.types import Detection
from ..core.imageio import imread_unicode, imwrite_unicode from ..core.imageio import imread_unicode, imwrite_unicode
from ..core.project import PROJECT_FILE, Project from ..core.project import PROJECT_FILE, Project
from ..core.restore.base import Cancelled
from ..core.restore.factory import build_restorer from ..core.restore.factory import build_restorer
from ..core.video.extract import extract_frames from ..core.video.extract import extract_frames
from ..core.video.frame import Frame from ..core.video.frame import Frame
@@ -65,6 +64,7 @@ from .extract_dialog import ExtractDialog
from .image_view import ImageView from .image_view import ImageView
from .marker_slider import MarkerSlider from .marker_slider import MarkerSlider
from .restore_dialog import RestoreDialog from .restore_dialog import RestoreDialog
from .workers import Job
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)" _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._nav_sync = False # guard against slider<->list signal loops
self._busy = False # a long operation is running self._busy = False # a long operation is running
self._cancel = False # the user asked to stop it self._cancel = False # the user asked to stop it
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.setWindowTitle("HVideoTool — инспектор детекции цензуры")
self.resize(1180, 720) self.resize(1180, 720)
@@ -97,6 +102,7 @@ class MainWindow(QMainWindow):
self._build_central() self._build_central()
self._build_statusbar() self._build_statusbar()
self._build_menu() self._build_menu()
self._probe_device() # determine CUDA/CPU in the background and fill the badge
self.statusBar().showMessage("Создайте или откройте проект (Файл)") self.statusBar().showMessage("Создайте или откройте проект (Файл)")
# ------------------------------------------------------------------ setup # ------------------------------------------------------------------ setup
@@ -310,17 +316,80 @@ class MainWindow(QMainWindow):
self.pos_label.setText(f"{row + 1 if row >= 0 else 0} / {n}") self.pos_label.setText(f"{row + 1 if row >= 0 else 0} / {n}")
def _build_statusbar(self) -> None: 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 = QProgressBar()
self.progress.setMaximumWidth(260) self.progress.setMaximumWidth(260)
self.progress.setVisible(False) self.progress.setVisible(False)
self.statusBar().addPermanentWidget(self.progress) 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 # ------------------------------------------------------------- cancellation
def _begin_busy(self, total: int | None = None) -> None: def _begin_busy(self, total: int | None = None) -> None:
"""Enter a cancellable long operation. ``total=None`` => busy spinner.""" """Enter a cancellable long operation. ``total=None`` => busy spinner."""
self._busy = True self._busy = True
self._cancel = False self._cancel = False
self.stop_action.setEnabled(True) 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: if total is None:
self.progress.setRange(0, 0) # indeterminate self.progress.setRange(0, 0) # indeterminate
else: else:
@@ -331,18 +400,55 @@ class MainWindow(QMainWindow):
def _end_busy(self) -> None: def _end_busy(self) -> None:
self._busy = False self._busy = False
self.stop_action.setEnabled(False) self.stop_action.setEnabled(False)
self.detector_combo.setEnabled(True)
self.model_action.setEnabled(True)
self.progress.setVisible(False) self.progress.setVisible(False)
self.progress.setRange(0, 100) # leave it determinate for the next user self.progress.setRange(0, 100) # leave it determinate for the next user
def _request_cancel(self) -> None: def _request_cancel(self) -> None:
if self._busy: if self._busy:
self._cancel = True self._cancel = True
if self._job is not None:
self._job.cancel() # stops the background loop at its next check
self.statusBar().showMessage("Отмена…") self.statusBar().showMessage("Отмена…")
def _poll_cancel(self) -> bool: # ------------------------------------------------------------- background jobs
"""Cancel hook for core engines: pump the UI so Стоп registers, then report.""" def _start_job(self, fn, total: int | None, *, on_tick=None, on_done=None) -> None:
QApplication.processEvents() """Run ``fn(job)`` on the thread pool; marshal results back to the GUI.
return self._cancel
``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 # --------------------------------------------------------------- detector
def _make_detector(self): def _make_detector(self):
@@ -355,12 +461,34 @@ class MainWindow(QMainWindow):
def _on_detector_changed(self, name: str) -> None: def _on_detector_changed(self, name: str) -> None:
self._cfg.detector = name 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: 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._persist_settings()
self._invalidate_results() 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: def _choose_model(self) -> None:
start = self._cfg.model_path or str(Path.cwd() / "models") start = self._cfg.model_path or str(Path.cwd() / "models")
path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)") path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)")
@@ -681,35 +809,44 @@ class MainWindow(QMainWindow):
if self._busy: if self._busy:
return return
path = Path(item.data(Qt.UserRole)) path = Path(item.data(Qt.UserRole))
if str(path) not in self._results: if str(path) in self._results:
if self._detect(path) is None: self._show(path)
return return
self._refresh_marks() self._detect_one(path, then_show=True)
self._save_results() # only when a detection actually ran
self._show(path)
def _detect(self, path: Path) -> list[Detection] | None: @staticmethod
"""Run (or fetch cached) detections for one image. None on failure.""" def _compute(detector, path: Path) -> list[Detection]:
key = str(path) """Pure read + detect for one image (runs on a worker thread; no Qt)."""
if key in self._results: img = imread_unicode(str(path))
return self._results[key]
img = imread_unicode(key)
if img is None: if img is None:
self.statusBar().showMessage(f"Не удалось прочитать: {path.name}") raise RuntimeError(f"Не удалось прочитать: {Path(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()
dets = detector.detect(Frame(image=img, index=0, pts=0.0)) dets = detector.detect(Frame(image=img, index=0, pts=0.0))
dets.sort(key=lambda d: d.score, reverse=True) dets.sort(key=lambda d: d.score, reverse=True)
self._results[key] = dets
self._tag_file(path, len(dets))
return 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: def _show(self, path: Path) -> None:
"""Display the image with its cached detections (does not run the detector).""" """Display the image with its cached detections (does not run the detector)."""
self._current = path self._current = path
@@ -721,90 +858,115 @@ class MainWindow(QMainWindow):
self._update_restore_actions() self._update_restore_actions()
def _recompute_current(self) -> None: 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: if self._current is None or self._busy:
return return
self._results.pop(str(self._current), None) self._results.pop(str(self._current), None)
self._detector_key = None # rebuild the detector so settings changes take effect self._detector_key = None # rebuild the detector so settings changes take effect
if self._detect(self._current) is None: self._detect_one(self._current, then_show=True)
return
self._refresh_marks()
self._save_results()
self._show(self._current)
def _detect_all(self, force: bool = False) -> None: def _detect_all(self, force: bool = False) -> None:
"""Detect the whole folder. ``force`` clears the cache first (full regen); """Detect the whole folder on a background thread. ``force`` clears the cache
otherwise already-computed frames are skipped, so it resumes/tops-up.""" 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: if not self._files or self._busy:
return return
if force: if force:
self._clear_results() self._clear_results()
total = len(self._files) pending = [p for p in self._files if str(p) not in self._results]
self._begin_busy(total) if not pending:
done = 0 self.statusBar().showMessage("Все кадры уже посчитаны (см. «Все заново»)")
try: return
for i, p in enumerate(self._files, 1): total = len(pending)
self.progress.setValue(i)
self.statusBar().showMessage(f"Детекция {i}/{total}: {p.name}") def fn(job):
QApplication.processEvents() detector = self._make_detector() # built on the worker thread (may raise)
if self._cancel: for i, p in enumerate(pending, 1):
if job.cancelled:
break break
if self._detect(p) is None: try:
return # detector unavailable — message already shown dets = self._compute(detector, p)
done = i except RuntimeError:
if i % 50 == 0: continue # unreadable image — skip, keep going
self._refresh_marks() # let marks appear progressively job.tick((str(p), dets))
finally: job.progress(i, total, f"Детекция {i}/{total}: {p.name}")
self._end_busy() return None
hits = sum(1 for p in self._files if self._results.get(str(p)))
self._refresh_marks() def done(_result, cancelled):
self._save_results() # persist progress (works for completed and cancelled runs) hits = sum(1 for p in self._files if self._results.get(str(p)))
if self._cancel: self._refresh_marks()
self.statusBar().showMessage( self._save_results() # persist progress (completed or cancelled)
f"Отменено на {done}/{total} · детекции на {hits} картинках" if cancelled:
) self.statusBar().showMessage(f"Отменено · детекции на {hits} картинках")
else: else:
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок") self.statusBar().showMessage(
if self._current is not None: f"Готово: детекции на {hits} из {len(self._files)} картинок"
self._show(self._current) )
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 # ------------------------------------------------------------- restoration
def _restore_current(self) -> None: 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: if self._current is None or self._busy:
return return
key = str(self._current) path = self._current
if key not in self._results and self._detect(self._current) is None: key = str(path)
return
self._refresh_marks() def fn(job):
dets = self._results.get(key) or [] dets = self._results.get(key)
if not dets: if dets is None:
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать") dets = self._compute(self._make_detector(), path)
return job.tick(("dets", key, dets)) # cache them on the GUI thread
img = imread_unicode(key) if not dets:
if img is None: return ("empty", key)
return img = imread_unicode(key)
self._begin_busy() # indeterminate — engine drives the duration if img is None:
self.statusBar().showMessage(f"Восстановление: {self._current.name}") raise RuntimeError(f"Не удалось прочитать: {path.name}")
QApplication.processEvents()
try:
restorer = self._make_restorer() restorer = self._make_restorer()
restored = restorer.restore(img, dets, should_cancel=self._poll_cancel) restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
except Cancelled: return ("restored", key, restored, len(dets), restorer.name)
self.statusBar().showMessage("Восстановление отменено")
return def tick(payload):
except Exception as exc: # noqa: BLE001 - surface model/engine errors if payload[0] == "dets":
QMessageBox.warning(self, "Ошибка восстановления", str(exc)) _, k, dets = payload
return self._results[k] = dets
finally: self._tag_file(Path(k), len(dets))
self._end_busy() self._refresh_marks()
self._restored[key] = restored
self._showing_restored = True def done(result, cancelled):
self.view.set_image(restored, []) if cancelled:
self._update_restore_actions() self.statusBar().showMessage("Восстановление отменено")
self.statusBar().showMessage( return
f"Расцензурено ({restorer.name}): {self._current.name}{len(dets)} обл." 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): def _make_restorer(self):
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,
@@ -1017,6 +1179,9 @@ class MainWindow(QMainWindow):
self._persist_settings() self._persist_settings()
def closeEvent(self, event) -> None: # noqa: N802 - Qt override 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 self._save_results() # persist the detection cache on exit
if self._project is not None: if self._project is not None:
self._project.update_from_config(self._cfg) self._project.update_from_config(self._cfg)
+67
View File
@@ -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)