Refactor HVideoTool's configuration and project management: updated project handling in core/project.py, streamlined restoration settings in config.py, and improved documentation in CLAUDE.md. Removed unused parameters and enhanced type hints for better clarity.

This commit is contained in:
Leonid Pershin
2026-06-07 06:33:00 +03:00
parent 9c471ca701
commit ac02ca27a8
16 changed files with 81 additions and 105 deletions
+3 -3
View File
@@ -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"]),
+1 -3
View File
@@ -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:
+2 -3
View File
@@ -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
+27 -25
View File
@@ -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
+3 -3
View File
@@ -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(
+6 -10
View File
@@ -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])
+2 -2
View File
@@ -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
+2 -3
View File
@@ -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)