Implement DeepMosaics restoration engine in HVideoTool: updated configuration options, integrated into the UI, and enhanced documentation in README and CLAUDE.md. The restoration process now supports both inpainting and generative models, with necessary setup instructions included.

2026-06-07 04:02:25 +03:00
parent ddc8543647
commit 7f0121b7df
9 changed files with 413 additions and 27 deletions
@@ -0,0 +1,107 @@
+"""DeepMosaics restorer — real generative mosaic removal.
+
+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.
+
+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).
+
+NOTE: DeepMosaics finds the mosaic itself; our detections are used for navigation,
+not passed to it.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+from ..detection.types import Detection
+from ..imageio import imread_unicode, imwrite_unicode
+from .base import Restorer
+
+_IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp"}
+
+
+class DeepMosaicsRestorer(Restorer):
+    def __init__(
+        self,
+        deepmosaics_dir: str | None,
+        model_path: str | None,
+        python_exe: str | None = None,
+        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():
+            raise ValueError(
+                "Не найдены веса DeepMosaics (clean_*.pth).\n"
+                "Скачайте clean_youknow_video.pth + mosaic_position.pth в одну папку. См. README."
+            )
+        self._dir = Path(deepmosaics_dir)
+        self._model = model_path
+        self._python = python_exe or sys.executable
+        self._gpu = gpu_id
+
+    @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)
+
+            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
+
+            # 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)
+            )
@@ -1,22 +1,34 @@
-"""Restorer factory: build a Restorer by name.
+"""Restorer factory: build a Restorer from the app config.

-Currently only the cv2 inpaint baseline is wired. Generative engines
-(DeepMosaics / LADA BasicVSR++) are placeholders — they need model weights and a
-CUDA GPU, and raise a clear, actionable error until integrated. See README.
+- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct).
+- ``deepmosaics``: real generative mosaic removal via a user-installed DeepMosaics
+  (subprocess). Needs the DeepMosaics folder + clean weights + (ideally) a CUDA GPU.
 """

 from __future__ import annotations

+from typing import TYPE_CHECKING
+
 from .base import Restorer
 from .inpaint import InpaintRestorer

+if TYPE_CHECKING:  # avoid importing AppConfig at runtime here (not needed)
+    from ...config import AppConfig

-def build_restorer(name: str = "inpaint", model_path: str | None = None) -> Restorer:
+
+def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer:
    if name == "inpaint":
        return InpaintRestorer()
-    if name in ("deepmosaics", "lada"):
+    if name == "deepmosaics":
+        from .deepmosaics import DeepMosaicsRestorer
+
+        if config is None:
+            raise ValueError("Для DeepMosaics нужны настройки (config).")
+        return DeepMosaicsRestorer(
+            config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
+        )
+    if name == "lada":
        raise ValueError(
-            "Генеративное восстановление пока не подключено.\n"
-            "Нужна модель (DeepMosaics / LADA) и GPU (CUDA). См. README → Восстановление."
+            "Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README."
        )
    raise ValueError(f"Неизвестный режим восстановления: {name!r}")