Refactor HVideoTool to support project-based workflow: introduced project management features, updated UI for project handling, and enhanced documentation in README and CLAUDE.md. The tool now organizes images and settings into projects, improving usability and detection caching.

2026-06-07 04:53:27 +03:00
parent 7f0121b7df
commit e27dfdf518
27 changed files with 2998 additions and 387 deletions
@@ -1,107 +1,161 @@
-"""DeepMosaics restorer — real generative mosaic removal.
+"""DeepMosaics restorer — real generative mosaic removal, in-process.

-Rather than vendoring DeepMosaics' GPL network code (which must match the exact
-checkpoint), we drive a **user-installed** DeepMosaics (https://github.com/HypoX64/DeepMosaics)
-as a subprocess: write the frame to a temp file, run ``deepmosaic.py --mode clean``,
-read the cleaned image back. This reuses their tested pipeline (incl. their own
-mosaic locator ``mosaic_position.pth``) and respects the GPL boundary.
+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.

-Setup the user must do once (see README → Восстановление):
-1. ``git clone https://github.com/HypoX64/DeepMosaics`` and install its deps.
-2. Download clean weights (e.g. ``clean_youknow_video.pth``) AND ``mosaic_position.pth``
-   into one folder.
-3. In the app: Восстановление… → engine "deepmosaics", set the DeepMosaics folder
-   and the clean-model path (a CUDA GPU is strongly recommended).
+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.

-NOTE: DeepMosaics finds the mosaic itself; our detections are used for navigation,
-not passed to it.
+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.
 """

 from __future__ import annotations

-import subprocess
 import sys
-import tempfile
 from pathlib import Path
+from types import SimpleNamespace

 import numpy as np

 from ..detection.types import Detection
-from ..imageio import imread_unicode, imwrite_unicode
-from .base import Restorer
+from .base import CancelCheck, Cancelled, Restorer

-_IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp"}
+_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).
+
+    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.
+    """
+    dirs = [DEFAULT_WEIGHTS_DIR]
+    if extra_dir:
+        dirs.insert(0, Path(extra_dir))
+    out: list[tuple[str, str]] = []
+    seen: set[str] = set()
+    for d in dirs:
+        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:
+                continue
+            seen.add(p.name)
+            out.append((p.stem, str(p)))
+    return out
+
+
+def _netg_kind(model_name: str) -> str:
+    """Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
+    n = model_name.lower()
+    if "video" in n:
+        raise ValueError(
+            "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен соседний "
+            "кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
+        )
+    if "unet_128" in n:
+        return "unet_128"
+    if "hd" in n:
+        return "HD"
+    return "resnet_9blocks"


 class DeepMosaicsRestorer(Restorer):
    def __init__(
        self,
-        deepmosaics_dir: str | None,
+        deepmosaics_dir: str | None,   # kept for factory/config compatibility (weights hint)
        model_path: str | None,
-        python_exe: str | None = None,
+        python_exe: str | None = None,  # unused now (in-process)
        gpu_id: str = "0",
    ) -> None:
-        if not deepmosaics_dir or not (Path(deepmosaics_dir) / "deepmosaic.py").is_file():
-            raise ValueError(
-                "Не указана папка DeepMosaics (с deepmosaic.py).\n"
-                "Установите DeepMosaics и укажите её в «Восстановление…». См. README."
-            )
        if not model_path or not Path(model_path).is_file():
+            discovered = discover_models()  # fall back to a bundled model
+            if discovered:
+                model_path = discovered[0][1]
+            else:
+                raise ValueError(
+                    "Не найдены веса DeepMosaics (clean_*.pth).\n"
+                    "Положите clean_youknow_resnet_9blocks.pth + mosaic_position.pth в "
+                    "models/deepmosaics (или выберите в «Восстановление…»). См. README."
+                )
+        model = Path(model_path)
+        self._netg = _netg_kind(model.name)  # raises on a video model
+        pos = self._find_mosaic_position(model, deepmosaics_dir)
+        if pos is None:
            raise ValueError(
-                "Не найдены веса DeepMosaics (clean_*.pth).\n"
-                "Скачайте clean_youknow_video.pth + mosaic_position.pth в одну папку. См. README."
+                "Рядом с clean-моделью не найден mosaic_position.pth.\n"
+                "Положите mosaic_position.pth в ту же папку, что и clean_*.pth. См. README."
            )
-        self._dir = Path(deepmosaics_dir)
-        self._model = model_path
-        self._python = python_exe or sys.executable
+        self._model = str(model)
+        self._pos = str(pos)
        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})"

-    def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray:
-        with tempfile.TemporaryDirectory(prefix="hvt_dm_") as tmp:
-            tmpd = Path(tmp)
-            src = tmpd / "frame.jpg"
-            result_dir = tmpd / "result"
-            result_dir.mkdir()
-            imwrite_unicode(str(src), image)
+    # ------------------------------------------------------------------ engine
+    def _ensure_loaded(self) -> None:
+        if self._loaded:
+            return
+        if str(_VENDOR) not in sys.path:
+            sys.path.insert(0, str(_VENDOR))  # so vendored `from models/util import …` resolve
+        from models import loadmodel, runmodel  # type: ignore  # noqa: E402
+        import util.image_processing as impro    # type: ignore  # noqa: E402

-            cmd = [
-                self._python, "deepmosaic.py",
-                "--media_path", str(src),
-                "--model_path", str(self._model),
-                "--mode", "clean",
-                "--result_dir", str(result_dir),
-                "--temp_dir", str(tmpd / "dmtmp"),
-                "--gpu_id", str(self._gpu),
-                "--no_preview",
-            ]
-            proc = subprocess.run(
-                cmd, cwd=str(self._dir),
-                stdin=subprocess.DEVNULL,            # so DeepMosaics' error input() can't hang
-                capture_output=True, text=True,
-            )
-            outputs = [p for p in result_dir.iterdir() if p.suffix.lower() in _IMG_EXTS]
-            if outputs:
-                newest = max(outputs, key=lambda p: p.stat().st_mtime)
-                restored = imread_unicode(str(newest))
-                if restored is None:
-                    raise RuntimeError("Не удалось прочитать результат DeepMosaics.")
-                return restored
+        self._runmodel = runmodel
+        self._impro = impro
+        self._opt = SimpleNamespace(
+            gpu_id=self._gpu,
+            netG=self._netg,
+            model_path=self._model,
+            mosaic_position_model_path=self._pos,
+            mask_threshold=64,
+            all_mosaic_area=False,
+            ex_mult=1.5,
+            no_feather=False,
+            traditional=False,
+        )
+        self._netM = loadmodel.bisenet(self._opt, "mosaic")
+        self._netG = loadmodel.pix2pix(self._opt)
+        self._loaded = True

-            # No output file — figure out why.
-            log = (proc.stderr or "") + (proc.stdout or "")
-            if "BVDNet.forward()" in log or "argument: 'previous'" in log:
-                raise RuntimeError(
-                    "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен "
-                    "соседний кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
-                )
-            if proc.returncode == 0:
-                # DeepMosaics ran fine but found no mosaic to clean — keep the frame as is.
-                return image.copy()
-            tail = log.strip().splitlines()[-6:]
-            raise RuntimeError(
-                f"DeepMosaics не вернул результат (код {proc.returncode}).\n" + "\n".join(tail)
-            )
+    def restore(
+        self,
+        image: np.ndarray,
+        detections: list[Detection],
+        should_cancel: CancelCheck | None = None,
+    ) -> np.ndarray:
+        if should_cancel is not None and should_cancel():
+            raise Cancelled("Восстановление отменено")
+        self._ensure_loaded()
+        rm, impro, opt = self._runmodel, self._impro, self._opt
+
+        # DeepMosaics' cleanmosaic_img_server, faithfully reproduced.
+        x, y, size, mask = rm.get_mosaic_position(image, self._netM, opt)
+        if size <= 100:
+            return image.copy()  # no mosaic located — leave the frame untouched
+        if should_cancel is not None and should_cancel():
+            raise Cancelled("Восстановление отменено")
+        work = image.copy()
+        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)