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:
@@ -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"]),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user