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)
+            )