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,105 @@
+"""Configure the restoration ("расцензурить") engine.
+
+inpaint — no setup. deepmosaics — point at a user-installed DeepMosaics folder and
+its clean weights; a CUDA GPU is strongly recommended (set GPU id, -1 = CPU/slow).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from PySide6.QtWidgets import (
+    QComboBox,
+    QDialog,
+    QDialogButtonBox,
+    QFileDialog,
+    QFormLayout,
+    QHBoxLayout,
+    QLabel,
+    QLineEdit,
+    QPushButton,
+    QWidget,
+)
+
+from ..config import AppConfig
+
+
+class RestoreDialog(QDialog):
+    def __init__(self, config: AppConfig, parent: QWidget | None = None) -> None:
+        super().__init__(parent)
+        self._cfg = config
+        self.setWindowTitle("Движок восстановления")
+        self.setMinimumWidth(520)
+
+        self.engine = QComboBox()
+        self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint")
+        self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics")
+        self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0)
+        self.engine.currentIndexChanged.connect(self._sync)
+
+        self.dm_dir = QLineEdit(config.dm_dir or "")
+        self.dm_model = QLineEdit(config.dm_model or "")
+        self.dm_python = QLineEdit(config.dm_python or "")
+        self.dm_python.setPlaceholderText("по умолчанию — python текущего venv")
+        self.dm_gpu = QLineEdit(config.dm_gpu or "0")
+        self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)")
+
+        form = QFormLayout(self)
+        form.addRow("Движок:", self.engine)
+        form.addRow("Папка DeepMosaics:", self._with_browse(self.dm_dir, self._browse_dir))
+        form.addRow("Веса (clean_*.pth):", self._with_browse(self.dm_model, self._browse_model))
+        form.addRow("Python для DeepMosaics:", self._with_browse(self.dm_python, self._browse_python))
+        form.addRow("GPU id:", self.dm_gpu)
+        hint = QLabel(
+            "DeepMosaics ставится отдельно (git clone + зависимости). Рядом с весами "
+            "clean_*.pth должен лежать mosaic_position.pth.\n"
+            "Для покадрового режима берите clean_youknow_resnet_9blocks.pth — "
+            "видеомодель clean_youknow_video.pth покадрово не работает. См. README."
+        )
+        hint.setWordWrap(True)
+        form.addRow(hint)
+
+        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
+        buttons.accepted.connect(self.accept)
+        buttons.rejected.connect(self.reject)
+        form.addRow(buttons)
+        self._sync()
+
+    def _with_browse(self, line: QLineEdit, slot) -> QWidget:
+        w = QWidget()
+        h = QHBoxLayout(w)
+        h.setContentsMargins(0, 0, 0, 0)
+        h.addWidget(line, 1)
+        btn = QPushButton("…")
+        btn.setMaximumWidth(32)
+        btn.clicked.connect(slot)
+        h.addWidget(btn)
+        return w
+
+    def _sync(self) -> None:
+        is_dm = self.engine.currentData() == "deepmosaics"
+        for w in (self.dm_dir, self.dm_model, self.dm_python, self.dm_gpu):
+            w.setEnabled(is_dm)
+
+    def _browse_dir(self) -> None:
+        d = QFileDialog.getExistingDirectory(self, "Папка DeepMosaics", self.dm_dir.text())
+        if d:
+            self.dm_dir.setText(d)
+
+    def _browse_model(self) -> None:
+        start = self.dm_model.text() or self.dm_dir.text()
+        p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", start, "Веса (*.pth);;Все файлы (*.*)")
+        if p:
+            self.dm_model.setText(p)
+
+    def _browse_python(self) -> None:
+        p, _ = QFileDialog.getOpenFileName(self, "Python для DeepMosaics", self.dm_python.text(), "python (*.exe);;Все файлы (*.*)")
+        if p:
+            self.dm_python.setText(p)
+
+    def apply_to_config(self) -> None:
+        self._cfg.restorer = self.engine.currentData()
+        self._cfg.dm_dir = self.dm_dir.text().strip() or None
+        self._cfg.dm_model = self.dm_model.text().strip() or None
+        self._cfg.dm_python = self.dm_python.text().strip() or None
+        self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0"