diff --git a/CLAUDE.md b/CLAUDE.md index aa494c5..bf887e8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -160,9 +160,9 @@ hvideotool/ команду установки" (`_copy_install_command` → the recommended cu121/cu118 command, picked by `recommend_channel` from the driver's CUDA) and "Проверить заново" (re-runs `_probe_device`). `core/torch_info.py` is pure (no Qt); subprocess uses `CREATE_NO_WINDOW` on Windows. -- **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: "Создать проект…" +- **Projects (`core/project.py`).** `MainWindow._project` is the open `Project`; its + frames come from `project.frames_dir` (navigation/cache/tags work off `_files`). Entry + points: "Создать проект…" (`_create_project`), "Открыть проект…" (`_open_project_dialog`), "Импортировать папку как проект…" (`_import_folder_as_project` — copies images into a new project's `frames/`, carries over an old `.hvideotool_detections.json` sidecar if present), and diff --git a/hvideotool/config.py b/hvideotool/config.py index 0725d09..3f911c0 100644 --- a/hvideotool/config.py +++ b/hvideotool/config.py @@ -49,9 +49,8 @@ class AppConfig: # --- restoration ("расцензурить") --- restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video" - dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py) + dm_dir: str | None = None # optional extra dir to search for mosaic_position.pth dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth) - dm_python: str | None = None # python exe for DeepMosaics (None = current) dm_gpu: str = "0" # CUDA device id, "-1" for CPU diff --git a/hvideotool/core/detection/types.py b/hvideotool/core/detection/types.py index 2c9a243..15251c3 100644 --- a/hvideotool/core/detection/types.py +++ b/hvideotool/core/detection/types.py @@ -3,10 +3,10 @@ from __future__ import annotations from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum -class CensorType(str, Enum): +class CensorType(StrEnum): """Kind of already-applied censorship a detection represents.""" MOSAIC = "mosaic" @@ -33,7 +33,7 @@ class Detection: } @classmethod - def from_dict(cls, data: dict) -> "Detection": + def from_dict(cls, data: dict) -> Detection: return cls( type=CensorType(data["type"]), score=float(data["score"]), diff --git a/hvideotool/core/detection/yolo.py b/hvideotool/core/detection/yolo.py index 7b7f4c2..82e6402 100644 --- a/hvideotool/core/detection/yolo.py +++ b/hvideotool/core/detection/yolo.py @@ -16,8 +16,6 @@ from __future__ import annotations import os -import numpy as np - from ...config import DetectionConfig from ..video.frame import Frame from .base import Detector @@ -61,7 +59,7 @@ class YoloDetector(Detector): import torch cuda_ok = torch.cuda.is_available() - except Exception: # noqa: BLE001 + except Exception: cuda_ok = False device = self.cfg.yolo_device if device is None: diff --git a/hvideotool/core/project.py b/hvideotool/core/project.py index 0fd6b34..94147ee 100644 --- a/hvideotool/core/project.py +++ b/hvideotool/core/project.py @@ -41,7 +41,6 @@ _SETTING_KEYS = ( "restorer", "dm_dir", "dm_model", - "dm_python", "dm_gpu", ) @@ -93,7 +92,7 @@ class Project: name: str | None = None, settings: dict | None = None, source: str | None = None, - ) -> "Project": + ) -> Project: """Create a new project folder (with ``frames/``) and write ``project.json``.""" root = Path(root) proj = cls( @@ -108,7 +107,7 @@ class Project: return proj @classmethod - def load(cls, path: Path | str) -> "Project": + def load(cls, path: Path | str) -> Project: """Load a project from its folder or directly from its ``project.json``.""" path = Path(path) root = path.parent if path.name == PROJECT_FILE else path diff --git a/hvideotool/core/restore/deepmosaics.py b/hvideotool/core/restore/deepmosaics.py index 991d13d..9813dc4 100644 --- a/hvideotool/core/restore/deepmosaics.py +++ b/hvideotool/core/restore/deepmosaics.py @@ -1,20 +1,22 @@ -"""DeepMosaics restorer — real generative mosaic removal, in-process. +"""DeepMosaics restorers — real generative mosaic removal, in-process. -The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see -its NOTICE/LICENSE). We load the models **once** and run the per-frame clean path -in-process — far faster than spawning a subprocess per frame (which reloaded the -models every time). Only the model *weights* are user-supplied. +The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see its +NOTICE/LICENSE). We load the models **once** and run in-process — far faster than +spawning a subprocess per frame (which reloaded the models every time). Only the model +*weights* are user-supplied. Both engines locate the mosaic themselves (BiSeNet +``mosaic_position.pth``); detections are not passed to them. -Per-frame clean = DeepMosaics' ``cleanmosaic_img_server`` logic, reimplemented -here (so we don't pull in their video/ffmpeg modules): - locate mosaic (BiSeNet ``mosaic_position.pth``) → run the clean generator on the - crop → feather it back. DeepMosaics finds the mosaic itself; our detections are - used for navigation, not passed to it. +Two engines: + - :class:`DeepMosaicsRestorer` (per-frame): reproduces ``cleanmosaic_img_server`` — + locate mosaic → run the image generator on the crop → feather it back. Image weights + ``clean_youknow_resnet_9blocks.pth``. + - :class:`DeepMosaicsVideoRestorer` (temporal/BVDNet): reproduces + ``cleanmosaic_video_fusion`` — a window of neighbouring frames + recurrence, for + temporal coherence. Video weights ``clean_youknow_video.pth``; needs a contiguous + sequence (see :meth:`Restorer.restore_sequence`). -Setup (see README → Восстановление): download the **image** clean weights -``clean_youknow_resnet_9blocks.pth`` + ``mosaic_position.pth`` into one folder and -point the app at the clean-model file. The video model ``clean_youknow_video.pth`` -(BVDNet) needs neighbour frames and does NOT work per-frame. +Setup (see README → Восстановление): drop the chosen ``clean_*.pth`` + ``mosaic_position.pth`` +into one folder (``models/deepmosaics``) and pick it in the restore dialog. """ from __future__ import annotations @@ -31,8 +33,8 @@ from .base import ( Cancelled, DetGetter, FrameGetter, - ResultSink, Restorer, + ResultSink, ) _VENDOR = Path(__file__).parent / "_deepmosaics" @@ -94,9 +96,8 @@ def _netg_kind(model_name: str) -> str: class DeepMosaicsRestorer(Restorer): def __init__( self, - deepmosaics_dir: str | None, # kept for factory/config compatibility (weights hint) + deepmosaics_dir: str | None, # optional extra dir to find mosaic_position.pth model_path: str | None, - python_exe: str | None = None, # unused now (in-process) gpu_id: str = "0", ) -> None: if not model_path or not Path(model_path).is_file(): @@ -136,12 +137,13 @@ class DeepMosaicsRestorer(Restorer): # 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 + import torch 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 + import util.image_processing as impro # type: ignore + + from models import loadmodel, runmodel # type: ignore self._runmodel = runmodel self._impro = impro @@ -206,7 +208,6 @@ class DeepMosaicsVideoRestorer(Restorer): self, deepmosaics_dir: str | None, model_path: str | None, - python_exe: str | None = None, # unused (in-process); kept for factory parity gpu_id: str = "0", ) -> None: chosen: Path | None = None @@ -244,13 +245,14 @@ class DeepMosaicsVideoRestorer(Restorer): if str(_VENDOR) not in sys.path: sys.path.insert(0, str(_VENDOR)) - import torch # noqa: E402 + import torch if self._gpu != "-1" and not torch.cuda.is_available(): self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why) - from models import loadmodel, runmodel # type: ignore # noqa: E402 - import util.data as data # type: ignore # noqa: E402 - import util.image_processing as impro # type: ignore # noqa: E402 + import util.data as data # type: ignore + import util.image_processing as impro # type: ignore + + from models import loadmodel, runmodel # type: ignore self._torch = torch self._runmodel = runmodel diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py index 00e3e75..4fbfd0a 100644 --- a/hvideotool/core/restore/factory.py +++ b/hvideotool/core/restore/factory.py @@ -20,20 +20,20 @@ if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed) from ...config import AppConfig -def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer: +def build_restorer(name: str = "deepmosaics", config: AppConfig | None = None) -> Restorer: if config is None: raise ValueError("Для DeepMosaics нужны настройки (config).") if name == "deepmosaics": from .deepmosaics import DeepMosaicsRestorer return DeepMosaicsRestorer( - config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu + config.dm_dir, config.dm_model, config.dm_gpu ) if name == "deepmosaics_video": from .deepmosaics import DeepMosaicsVideoRestorer return DeepMosaicsVideoRestorer( - config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu + config.dm_dir, config.dm_model, config.dm_gpu ) if name == "lada": raise ValueError( diff --git a/hvideotool/core/torch_info.py b/hvideotool/core/torch_info.py index eb16221..b392596 100644 --- a/hvideotool/core/torch_info.py +++ b/hvideotool/core/torch_info.py @@ -44,7 +44,7 @@ def _run_nvidia_smi() -> dict: out["gpus"].append(parts[0]) if len(parts) > 1 and parts[1]: out["driver_version"] = parts[1] - except Exception: # noqa: BLE001 - any failure => "not found" + except Exception: return out # Max CUDA version the driver supports (only in the plain header). try: @@ -54,7 +54,7 @@ def _run_nvidia_smi() -> dict: m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout) if m: out["cuda_driver"] = m.group(1) - except Exception: # noqa: BLE001 + except Exception: pass return out @@ -76,23 +76,23 @@ def gather() -> dict: } try: import torch - except Exception as exc: # noqa: BLE001 - report any import failure + except Exception as exc: info["import_error"] = str(exc) else: info["installed"] = True info["version"] = getattr(torch, "__version__", None) try: info["built_cuda"] = torch.version.cuda - except Exception: # noqa: BLE001 + except Exception: info["built_cuda"] = None try: info["cuda_available"] = bool(torch.cuda.is_available()) - except Exception: # noqa: BLE001 + except Exception: info["cuda_available"] = False if info["cuda_available"]: try: info["device_name"] = torch.cuda.get_device_name(0) - except Exception: # noqa: BLE001 + except Exception: info["device_name"] = None smi = _run_nvidia_smi() @@ -103,10 +103,6 @@ def gather() -> dict: return info -def device_label(info: dict) -> str: - return "CUDA" if info.get("cuda_available") else "CPU" - - def _ver_tuple(v: str | None) -> tuple[int, ...]: try: return tuple(int(x) for x in str(v).split(".")[:2]) diff --git a/hvideotool/core/video/extract.py b/hvideotool/core/video/extract.py index 6462967..71d9c6e 100644 --- a/hvideotool/core/video/extract.py +++ b/hvideotool/core/video/extract.py @@ -18,8 +18,8 @@ from __future__ import annotations import shutil import subprocess +from collections.abc import Callable from pathlib import Path -from typing import Callable import cv2 @@ -38,7 +38,7 @@ def _find_ffmpeg() -> str | None: import imageio_ffmpeg return imageio_ffmpeg.get_ffmpeg_exe() - except Exception: # noqa: BLE001 - package missing or no bundled binary + except Exception: return None diff --git a/hvideotool/core/video/frame.py b/hvideotool/core/video/frame.py index 4c58151..073a97e 100644 --- a/hvideotool/core/video/frame.py +++ b/hvideotool/core/video/frame.py @@ -1,4 +1,4 @@ -"""Decoded video frame passed from the reader to the detector.""" +"""Image passed to the detector — the ``Detector.detect`` input type.""" from __future__ import annotations @@ -10,5 +10,4 @@ import numpy as np @dataclass class Frame: image: np.ndarray # BGR, HxWx3, uint8 (OpenCV convention) - index: int # 0-based frame counter since the last open/seek - pts: float # presentation timestamp, seconds + index: int = 0 # 0-based position in the sequence (informational) diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py index f533f27..008c0fb 100644 --- a/hvideotool/settings_store.py +++ b/hvideotool/settings_store.py @@ -39,7 +39,7 @@ def apply(config: AppConfig) -> None: config.default_threshold = float(data["threshold"]) if data.get("restorer"): config.restorer = data["restorer"] - for key in ("dm_dir", "dm_model", "dm_python", "dm_gpu"): + for key in ("dm_dir", "dm_model", "dm_gpu"): if key in data: setattr(config, key, data[key]) @@ -54,7 +54,6 @@ def save(config: AppConfig) -> None: restorer=config.restorer, dm_dir=config.dm_dir, dm_model=config.dm_model, - dm_python=config.dm_python, dm_gpu=config.dm_gpu, ) _write(data) diff --git a/hvideotool/ui/image_view.py b/hvideotool/ui/image_view.py index ae8863a..61c12ca 100644 --- a/hvideotool/ui/image_view.py +++ b/hvideotool/ui/image_view.py @@ -1,9 +1,9 @@ """Widget that renders an image and draws detection overlays. -Overlay visibility and the confidence threshold are applied at paint time, so -toggling them is instant. One detection can be *highlighted* (selected in the -detail table) — it is drawn boldly even if below the threshold, while the others -dim, so the user can inspect exactly what the detector found. +The confidence threshold is applied at paint time, so changing it is instant. One +detection can be *highlighted* (selected in the detail table) — it is drawn boldly +even if below the threshold, while the others dim, so the user can inspect exactly +what the detector found. """ from __future__ import annotations @@ -24,7 +24,6 @@ class ImageView(QWidget): self._cfg = overlay_cfg self._qimage: QImage | None = None self._dets: list[Detection] = [] - self._overlay_enabled = True self._threshold = 0.0 self._highlight: int | None = None self.setMinimumSize(480, 360) @@ -42,10 +41,6 @@ class ImageView(QWidget): self._highlight = None self.update() - def set_overlay_enabled(self, enabled: bool) -> None: - self._overlay_enabled = enabled - self.update() - def set_threshold(self, threshold: float) -> None: self._threshold = threshold self.update() @@ -59,13 +54,13 @@ class ImageView(QWidget): r, g, b = self._cfg.colors.get(ctype.value, (255, 0, 0)) return QColor(r, g, b) - def paintEvent(self, event) -> None: # noqa: N802 - Qt signature + def paintEvent(self, event) -> None: painter = QPainter(self) painter.fillRect(self.rect(), QColor(18, 18, 18)) if self._qimage is None: painter.setPen(QColor(160, 160, 160)) - painter.drawText(self.rect(), Qt.AlignCenter, "Откройте папку с картинками (Файл → Открыть папку…)") + painter.drawText(self.rect(), Qt.AlignCenter, "Откройте проект (Файл → Открыть проект…)") painter.end() return @@ -77,7 +72,7 @@ class ImageView(QWidget): painter.setRenderHint(QPainter.SmoothPixmapTransform, True) painter.drawImage(QRectF(ox, oy, dw, dh), self._qimage) - if self._overlay_enabled and self._dets: + if self._dets: painter.setRenderHint(QPainter.Antialiasing, True) for i, d in enumerate(self._dets): highlighted = i == self._highlight diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index 4fca337..1801017 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -5,23 +5,24 @@ A project is a folder (``project.json`` + ``frames/`` + ``detections.json`` + model, threshold, restore engine) live in ``project.json``; the global ``settings.json`` only seeds defaults for new projects. -Layout: a toolbar (new/open project · from-video · detector · model · calc-frame · -detect-all · threshold), then a splitter with three panes — left: collection -controls + the file list; center: the image with overlays; right: a detail table -of every detection. Collection controls sit by the file list (they act on its -selection), keeping the toolbar to detection/entry actions only. +Layout: a toolbar (new/open project · from-video · model · calc-frame · detect-all · +restore · threshold), then a splitter with three panes — left: collection controls + +the file list; center: the image with overlays; right: a detail table of every +detection. Collection controls sit by the file list (they act on its selection), +keeping the toolbar to detection/entry actions only. -Viewing and detecting are decoupled, so browsing a big project stays instant even -with a slow (CPU) detector: +Detection is YOLO-only; restoration is DeepMosaics-only. Viewing and detecting are +decoupled, so browsing a big project stays instant even with a slow (CPU) detector: - selecting a file just **shows** it (with its cached result, if any); - **double-clicking** a file, or "Рассчитать кадр", runs the detector on it; - "Детектировать все" runs the whole project. Both project loading and detect-all show a progress bar. Results are cached in the -project; switching detector/model clears the cache. +project; switching the model clears the cache. """ from __future__ import annotations +import contextlib import shutil from pathlib import Path @@ -77,13 +78,12 @@ class MainWindow(QMainWindow): self._detector = None self._detector_key = None self._project: Project | None = None # the open project (None until one is opened) - self._folder: Path | None = None # == project.frames_dir while a project is open self._files: list[Path] = [] self._results: dict[str, list[Detection]] = {} # path -> detections (cache) self._current: Path | None = None self._restorer = None # un-censor engine, built lazily from config self._restorer_key = None - self._restored: dict[str, "object"] = {} # path -> restored image (BGR ndarray) + self._restored: dict[str, object] = {} # path -> restored image (BGR ndarray) self._showing_restored = False self._nav_sync = False # guard against slider<->list signal loops self._busy = False # a long operation is running @@ -547,7 +547,7 @@ class MainWindow(QMainWindow): if Project.is_project(p): try: self._open_project(Project.load(p)) - except (OSError, ValueError) as exc: # noqa: BLE001 - surface to the user + except (OSError, ValueError) as exc: QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}") elif p.is_dir(): QMessageBox.information( @@ -562,10 +562,8 @@ class MainWindow(QMainWindow): """On startup, reopen the last project if it still exists (best-effort).""" last = settings_store.last_project() if last and Project.is_project(last): - try: + with contextlib.suppress(OSError, ValueError): self._open_project(Project.load(last)) - except (OSError, ValueError): - pass def _new_project_root(self, default_name: str = "") -> Path | None: """Prompt for a parent dir + name; return a fresh (empty) project root or None.""" @@ -662,10 +660,8 @@ class MainWindow(QMainWindow): # Carry over an old sidecar detection cache (basename-keyed) if present. old_sidecar = src / ".hvideotool_detections.json" if old_sidecar.is_file(): - try: + with contextlib.suppress(OSError): shutil.copy2(str(old_sidecar), str(project.cache_path)) - except OSError: - pass self.statusBar().showMessage(f"Импортировано {copied} картинок → {project.name}") self._open_project(project) @@ -764,7 +760,7 @@ class MainWindow(QMainWindow): str(video), str(out), step=step, keyframes_only=keyframes_only, max_dim=max_dim, progress=cb, ) - except Exception as exc: # noqa: BLE001 - surface decode errors to the user + except Exception as exc: QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}") return finally: @@ -788,7 +784,6 @@ class MainWindow(QMainWindow): self.statusBar().showMessage(f"Сканирую папку: {folder}…") QApplication.processEvents() files = sorted(p for p in folder.iterdir() if p.suffix.lower() in _IMAGE_EXTS) - self._folder = folder self._files = files self._results.clear() self._current = None @@ -850,7 +845,7 @@ class MainWindow(QMainWindow): img = imread_unicode(str(path)) if img is None: raise RuntimeError(f"Не удалось прочитать: {Path(path).name}") - dets = detector.detect(Frame(image=img, index=0, pts=0.0)) + dets = detector.detect(Frame(image=img)) dets.sort(key=lambda d: d.score, reverse=True) return dets @@ -1057,8 +1052,7 @@ class MainWindow(QMainWindow): self._start_job(fn, total, on_done=done) def _make_restorer(self): - key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, - self._cfg.dm_python, self._cfg.dm_gpu) + key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, self._cfg.dm_gpu) if key != self._restorer_key: self._restorer = build_restorer(self._cfg.restorer, self._cfg) # may raise self._restorer_key = key @@ -1266,7 +1260,7 @@ class MainWindow(QMainWindow): self.view.set_threshold(value) self._persist_settings() - def closeEvent(self, event) -> None: # noqa: N802 - Qt override + def closeEvent(self, event) -> None: if self._job is not None: # stop a running background job before tearing down self._job.cancel() self._pool.waitForDone(3000) diff --git a/hvideotool/ui/marker_slider.py b/hvideotool/ui/marker_slider.py index 61b1409..9a998ac 100644 --- a/hvideotool/ui/marker_slider.py +++ b/hvideotool/ui/marker_slider.py @@ -27,11 +27,6 @@ class MarkerSlider(QSlider): self._marks = marks self.update() - def clear_marks(self) -> None: - if self._marks: - self._marks = set() - self.update() - def paintEvent(self, event) -> None: super().paintEvent(event) if not self._marks or self.maximum() <= self.minimum(): diff --git a/hvideotool/ui/workers.py b/hvideotool/ui/workers.py index 3304a2b..ddbf456 100644 --- a/hvideotool/ui/workers.py +++ b/hvideotool/ui/workers.py @@ -33,7 +33,7 @@ class _Signals(QObject): class Job(QRunnable): """Runs ``fn(job)`` on a thread pool, marshaling progress/result to the GUI.""" - def __init__(self, fn: Callable[["Job"], Any]) -> None: + def __init__(self, fn: Callable[[Job], Any]) -> None: super().__init__() self.setAutoDelete(False) # the GUI keeps a reference until `done`/`failed` self.signals = _Signals() @@ -56,12 +56,12 @@ class Job(QRunnable): self.signals.tick.emit(payload) # -- thread entry point -- - def run(self) -> None: # noqa: D401 - QRunnable override + def run(self) -> None: 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 + except Exception as exc: self.signals.failed.emit(str(exc)) else: self.signals.done.emit(result) diff --git a/pyproject.toml b/pyproject.toml index 11d2463..5e2a95d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,9 +19,9 @@ dependencies = [ ] [project.optional-dependencies] -# Trained-model detector (future work). PyTorch must be installed separately -# with the correct CUDA build — see README. Installing this extra only pulls in -# Ultralytics; it does NOT install torch. +# YOLO detector (Ultralytics). PyTorch must be installed separately with the correct +# CUDA build — see README. This extra only pulls in Ultralytics; it does NOT install +# torch. Both detection (YOLO) and restoration (DeepMosaics) need torch at runtime. yolo = ["ultralytics>=8.0"] [project.scripts]