Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.

2026-06-07 06:16:05 +03:00
parent cc518cc3e6
commit 9c471ca701
17 changed files with 798 additions and 676 deletions
@@ -26,18 +26,29 @@ from types import SimpleNamespace
 import numpy as np

 from ..detection.types import Detection
-from .base import CancelCheck, Cancelled, Restorer
+from .base import (
+    CancelCheck,
+    Cancelled,
+    DetGetter,
+    FrameGetter,
+    ResultSink,
+    Restorer,
+)

 _VENDOR = Path(__file__).parent / "_deepmosaics"
 # Default place to drop DeepMosaics clean weights (gitignored — see models/).
 DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"


-def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
-    """Find usable per-frame clean models: (display_name, full_path).
+def discover_models(
+    extra_dir: str | None = None, include_video: bool = False
+) -> list[tuple[str, str]]:
+    """Find usable clean models: (display_name, full_path).

    Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given)
-    for ``clean_*.pth``. The video model is skipped — it can't run per-frame.
+    for ``clean_*.pth``. By default the video model is skipped — it can't run
+    per-frame; pass ``include_video=True`` for the temporal (BVDNet) engine, which
+    *needs* ``clean_*_video.pth``.
    """
    dirs = [DEFAULT_WEIGHTS_DIR]
    if extra_dir:
@@ -48,13 +59,23 @@ def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
        if not d.is_dir():
            continue
        for p in sorted(d.glob("clean_*.pth")):
-            if "video" in p.name.lower() or p.name in seen:
+            if p.name in seen:
+                continue
+            if "video" in p.name.lower() and not include_video:
                continue
            seen.add(p.name)
            out.append((p.stem, str(p)))
    return out


+def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
+    """Locate ``mosaic_position.pth`` (the BiSeNet mosaic locator) for ``model``."""
+    candidates = [model.parent / "mosaic_position.pth"]
+    if dm_dir:
+        candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
+    return next((p for p in candidates if p.is_file()), None)
+
+
 def _netg_kind(model_name: str) -> str:
    """Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
    n = model_name.lower()
@@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer):
                )
        model = Path(model_path)
        self._netg = _netg_kind(model.name)  # raises on a video model
-        pos = self._find_mosaic_position(model, deepmosaics_dir)
+        pos = _find_mosaic_position(model, deepmosaics_dir)
        if pos is None:
            raise ValueError(
                "Рядом с clean-моделью не найден mosaic_position.pth.\n"
@@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer):
        self._gpu = gpu_id
        self._loaded = False  # models loaded lazily on first restore

-    @staticmethod
-    def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
-        candidates = [model.parent / "mosaic_position.pth"]
-        if dm_dir:
-            candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
-        return next((p for p in candidates if p.is_file()), None)
-
    @property
    def name(self) -> str:
        return f"DeepMosaics(gpu={self._gpu})"
@@ -167,3 +181,147 @@ class DeepMosaicsRestorer(Restorer):
        img_mosaic = work[y - size:y + size, x - size:x + size]
        img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
        return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather)
+
+
+class DeepMosaicsVideoRestorer(Restorer):
+    """Temporal DeepMosaics (BVDNet) — un-censors using *neighbouring* frames.
+
+    Reproduces DeepMosaics' ``cleanmosaic_video_fusion`` per target frame: for frame
+    ``i`` it feeds the network a temporal window of ``T`` frames (sampled at step ``S``
+    around ``i``) plus its own previous output (recurrent), so the reconstruction is
+    temporally coherent. Because of that recurrence the frames MUST be processed in
+    order over a contiguous range — see :meth:`restore_sequence` (the batch run).
+
+    Needs the **video** weights ``clean_youknow_video.pth`` + ``mosaic_position.pth``
+    (beside it). Single-frame :meth:`restore` degrades to a window of the same frame.
+    """
+
+    temporal = True
+
+    # DeepMosaics fusion window: N before/after at step S → T = 2N+1 frames, INPUT_SIZE px.
+    _N, _T, _S = 2, 5, 3
+    _INPUT_SIZE = 256
+
+    def __init__(
+        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
+        if model_path and Path(model_path).is_file() and "video" in Path(model_path).name.lower():
+            chosen = Path(model_path)
+        else:  # configured model missing or not a video model → auto-pick a video model
+            vids = [p for _n, p in discover_models(include_video=True) if "video" in Path(p).name.lower()]
+            if vids:
+                chosen = Path(vids[0])
+        if chosen is None:
+            raise ValueError(
+                "Не найдены веса видеомодели DeepMosaics (clean_*_video.pth).\n"
+                "Положите clean_youknow_video.pth + mosaic_position.pth в models/deepmosaics "
+                "(или выберите в «Восстановление…»). См. README."
+            )
+        pos = _find_mosaic_position(chosen, deepmosaics_dir)
+        if pos is None:
+            raise ValueError(
+                "Рядом с видеомоделью не найден mosaic_position.pth.\n"
+                "Положите mosaic_position.pth в ту же папку, что и clean_*_video.pth. См. README."
+            )
+        self._model = str(chosen)
+        self._pos = str(pos)
+        self._gpu = gpu_id
+        self._loaded = False
+
+    @property
+    def name(self) -> str:
+        return f"DeepMosaicsVideo(gpu={self._gpu})"
+
+    # ------------------------------------------------------------------ engine
+    def _ensure_loaded(self) -> None:
+        if self._loaded:
+            return
+        if str(_VENDOR) not in sys.path:
+            sys.path.insert(0, str(_VENDOR))
+
+        import torch  # noqa: E402
+        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
+
+        self._torch = torch
+        self._runmodel = runmodel
+        self._data = data
+        self._impro = impro
+        self._opt = SimpleNamespace(
+            gpu_id=self._gpu,
+            model_path=self._model,
+            mosaic_position_model_path=self._pos,
+            mask_threshold=64,
+            all_mosaic_area=False,
+            ex_mult=1.5,
+            no_feather=False,
+        )
+        self._netM = loadmodel.bisenet(self._opt, "mosaic")
+        self._netG = loadmodel.video(self._opt)  # BVDNet
+        self._loaded = True
+
+    def restore(
+        self,
+        image: np.ndarray,
+        detections: list[Detection],
+        should_cancel: CancelCheck | None = None,
+    ) -> np.ndarray:
+        """Single-frame restore — no neighbours, so the window is the same frame."""
+        out: dict[int, np.ndarray] = {}
+        self.restore_sequence(
+            1,
+            lambda _i: image,
+            lambda _i: detections,
+            lambda i, r: out.__setitem__(i, r),
+            should_cancel,
+        )
+        return out.get(0, image.copy())
+
+    def restore_sequence(
+        self,
+        count: int,
+        get_frame: FrameGetter,
+        get_dets: DetGetter,
+        emit: ResultSink,
+        should_cancel: CancelCheck | None = None,
+    ) -> None:
+        self._ensure_loaded()
+        torch, data, impro, opt = self._torch, self._data, self._impro, self._opt
+        N, T, S, SZ = self._N, self._T, self._S, self._INPUT_SIZE
+
+        previous = None  # recurrent state: the network's previous output (a tensor)
+        for i in range(count):
+            if should_cancel is not None and should_cancel():
+                raise Cancelled("Восстановление отменено")
+            img_origin = get_frame(i)
+            x, y, size, mask = self._runmodel.get_mosaic_position(img_origin, self._netM, opt)
+            if size <= 50:
+                emit(i, img_origin.copy())  # no mosaic here; recurrence carries over
+                continue
+
+            stream = []
+            for k in range(T):
+                j = min(max(i + (k - N) * S, 0), count - 1)  # clamp window to range edges
+                frame = img_origin if j == i else get_frame(j)
+                crop = frame[y - size:y + size, x - size:x + size]
+                stream.append(impro.resize(crop, SZ)[:, :, ::-1])  # BGR→RGB, SZ×SZ
+
+            if previous is None:  # seed recurrence with the (centre) input crop
+                previous = data.im2tensor(stream[N], bgr2rgb=False, gpu_id=opt.gpu_id)
+
+            arr = np.array(stream).reshape(1, T, SZ, SZ, 3).transpose((0, 4, 1, 2, 3))
+            tensor = data.to_tensor(data.normalize(arr), gpu_id=opt.gpu_id)
+            with torch.no_grad():
+                pred = self._netG(tensor, previous)
+            previous = pred
+            img_fake = data.tensor2im(pred, rgb2bgr=True)
+            emit(i, impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather))