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:
@@ -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
|
||||
`<stem>_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
|
||||
|
||||
@@ -67,6 +67,10 @@
|
||||
- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога**
|
||||
уверенности прямо в тулбаре — удобно сравнивать.
|
||||
- Выбор файла весов модели кнопкой **«Модель…»**.
|
||||
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
|
||||
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
|
||||
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU
|
||||
(медленнее, но без ошибок).
|
||||
|
||||
## Кэш детекций
|
||||
|
||||
|
||||
@@ -54,14 +54,19 @@ class YoloDetector(Detector):
|
||||
"(и PyTorch с CUDA отдельно — см. README)."
|
||||
) from exc
|
||||
|
||||
# Resolve the device: explicit override, else CUDA when available.
|
||||
device = self.cfg.yolo_device
|
||||
if device is None:
|
||||
# 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
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
cuda_ok = torch.cuda.is_available()
|
||||
except Exception: # noqa: BLE001
|
||||
cuda_ok = False
|
||||
device = self.cfg.yolo_device
|
||||
if device is None:
|
||||
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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 детекция и расцензуривание "
|
||||
"просто медленнее."
|
||||
)
|
||||
+247
-82
@@ -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:
|
||||
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
|
||||
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()
|
||||
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 (works for completed and cancelled runs)
|
||||
if self._cancel:
|
||||
self.statusBar().showMessage(
|
||||
f"Отменено на {done}/{total} · детекции на {hits} картинках"
|
||||
)
|
||||
self._save_results() # persist progress (completed or cancelled)
|
||||
if cancelled:
|
||||
self.statusBar().showMessage(f"Отменено · детекции на {hits} картинках")
|
||||
else:
|
||||
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
||||
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 []
|
||||
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:
|
||||
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
|
||||
return
|
||||
return ("empty", key)
|
||||
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:
|
||||
raise RuntimeError(f"Не удалось прочитать: {path.name}")
|
||||
restorer = self._make_restorer()
|
||||
restored = restorer.restore(img, dets, should_cancel=self._poll_cancel)
|
||||
except Cancelled:
|
||||
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
|
||||
except Exception as exc: # noqa: BLE001 - surface model/engine errors
|
||||
QMessageBox.warning(self, "Ошибка восстановления", str(exc))
|
||||
if result is None:
|
||||
return
|
||||
finally:
|
||||
self._end_busy()
|
||||
self._restored[key] = restored
|
||||
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"Расцензурено ({restorer.name}): {self._current.name} — {len(dets)} обл."
|
||||
)
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user