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.

hvideotool/config.py

@@ -65,3 +65,10 @@ class AppConfig:
     detector: str = "classic"          # "classic" | "yolo" | "combined"
     model_path: str | None = None      # weights path, used by the YOLO detector
     default_threshold: float = 0.20    # initial overlay confidence threshold
+
+    # --- restoration ("расцензурить") ---
+    restorer: str = "inpaint"          # "inpaint" | "deepmosaics"
+    dm_dir: str | None = None          # DeepMosaics repo dir (contains deepmosaic.py)
+    dm_model: str | None = None        # DeepMosaics clean weights (clean_*.pth)
+    dm_python: str | None = None       # python exe for DeepMosaics (None = current)
+    dm_gpu: str = "0"                  # CUDA device id, "-1" for CPU

hvideotool/core/restore/deepmosaics.py

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

hvideotool/core/restore/factory.py

@@ -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}")

hvideotool/settings_store.py

@@ -35,6 +35,11 @@ def apply(config: AppConfig) -> None:
         config.model_path = data["model_path"]
     if "threshold" in data:
         config.default_threshold = float(data["threshold"])
+    if data.get("restorer"):
+        config.restorer = data["restorer"]
+    for key in ("dm_dir", "dm_model", "dm_python", "dm_gpu"):
+        if key in data:
+            setattr(config, key, data[key])
 
 
 def save(config: AppConfig) -> None:
@@ -44,6 +49,11 @@ def save(config: AppConfig) -> None:
         detector=config.detector,
         model_path=config.model_path,
         threshold=config.default_threshold,
+        restorer=config.restorer,
+        dm_dir=config.dm_dir,
+        dm_model=config.dm_model,
+        dm_python=config.dm_python,
+        dm_gpu=config.dm_gpu,
     )
     _write(data)
 

hvideotool/ui/main_window.py

@@ -21,7 +21,7 @@ import shutil
 from pathlib import Path
 
 from PySide6.QtCore import Qt
-from PySide6.QtGui import QAction, QKeySequence, QShortcut
+from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
 from PySide6.QtWidgets import (
     QAbstractItemView,
     QApplication,
@@ -38,7 +38,6 @@ from PySide6.QtWidgets import (
     QMessageBox,
     QProgressBar,
     QPushButton,
-    QSlider,
     QSplitter,
     QTableWidget,
     QTableWidgetItem,
@@ -56,6 +55,8 @@ from ..core.video.extract import extract_frames
 from ..core.video.frame import Frame
 from .extract_dialog import ExtractDialog
 from .image_view import ImageView
+from .marker_slider import MarkerSlider
+from .restore_dialog import RestoreDialog
 
 _IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
 _VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
@@ -73,7 +74,8 @@ class MainWindow(QMainWindow):
         self._results: dict[str, list[Detection]] = {}  # path -> detections (cache)
         self._current: Path | None = None
         self._collection: Path | None = None  # active destination folder for moves
-        self._restorer = build_restorer("inpaint")  # un-censor engine (baseline)
+        self._restorer = None           # un-censor engine, built lazily from config
+        self._restorer_key = None
         self._restored: dict[str, "object"] = {}    # path -> restored image (BGR ndarray)
         self._showing_restored = False
         self._nav_sync = False  # guard against slider<->list signal loops
@@ -97,6 +99,8 @@ class MainWindow(QMainWindow):
         file_menu.addAction("Рассчитать кадр", self._recompute_current).setShortcut("Space")
         file_menu.addAction("Детектировать все", self._detect_all)
         file_menu.addSeparator()
+        file_menu.addAction("Движок восстановления…", self._open_restore_settings)
+        file_menu.addSeparator()
         file_menu.addAction("Создать коллекцию…", self._create_collection)
         file_menu.addAction("В коллекцию", self._move_to_collection).setShortcut("Ctrl+M")
         file_menu.addSeparator()
@@ -223,10 +227,10 @@ class MainWindow(QMainWindow):
         self.next_btn.setToolTip("Следующий кадр (.)")
         self.next_btn.clicked.connect(lambda: self._step(1))
 
-        self.frame_slider = QSlider(Qt.Horizontal)
+        self.frame_slider = MarkerSlider(Qt.Horizontal)
         self.frame_slider.setMinimum(0)
         self.frame_slider.setMaximum(0)
-        self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →)")
+        self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
         self.frame_slider.valueChanged.connect(self._on_slider)
 
         self.pos_label = QLabel("0 / 0")
@@ -329,7 +333,10 @@ class MainWindow(QMainWindow):
         self._detector_key = None
         self._results.clear()
         for i in range(self.file_list.count()):
-            self.file_list.item(i).setText(self.file_list.item(i).data(Qt.UserRole + 1))
+            item = self.file_list.item(i)
+            item.setText(item.data(Qt.UserRole + 1))
+            item.setBackground(QBrush())
+        self._refresh_marks()
         if self._current is not None:
             self._show(self._current)
 
@@ -416,6 +423,7 @@ class MainWindow(QMainWindow):
         self.file_list.setUpdatesEnabled(True)
         self.file_list.blockSignals(False)
         self.progress.setVisible(False)
+        self._refresh_marks()
 
         self._refresh_collections()
         if not files:
@@ -436,6 +444,7 @@ class MainWindow(QMainWindow):
         path = Path(item.data(Qt.UserRole))
         if str(path) not in self._results and self._detect(path) is None:
             return
+        self._refresh_marks()
         self._show(path)
 
     def _detect(self, path: Path) -> list[Detection] | None:
@@ -478,6 +487,7 @@ class MainWindow(QMainWindow):
         self._detector_key = None  # rebuild the detector so settings changes take effect
         if self._detect(self._current) is None:
             return
+        self._refresh_marks()
         self._show(self._current)
 
     def _detect_all(self) -> None:
@@ -496,6 +506,7 @@ class MainWindow(QMainWindow):
         finally:
             self.progress.setVisible(False)
         hits = sum(1 for p in self._files if self._results.get(str(p)))
+        self._refresh_marks()
         self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
         if self._current is not None:
             self._show(self._current)
@@ -508,6 +519,7 @@ class MainWindow(QMainWindow):
         key = str(self._current)
         if key not in self._results and self._detect(self._current) is None:
             return
+        self._refresh_marks()
         dets = self._results.get(key) or []
         if not dets:
             self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
@@ -518,7 +530,8 @@ class MainWindow(QMainWindow):
         self.statusBar().showMessage(f"Восстановление: {self._current.name}…")
         QApplication.processEvents()
         try:
-            restored = self._restorer.restore(img, dets)
+            restorer = self._make_restorer()
+            restored = restorer.restore(img, dets)
         except Exception as exc:  # noqa: BLE001 - surface model/engine errors
             QMessageBox.warning(self, "Ошибка восстановления", str(exc))
             return
@@ -527,9 +540,26 @@ class MainWindow(QMainWindow):
         self.view.set_image(restored, [])
         self._update_restore_actions()
         self.statusBar().showMessage(
-            f"Расцензурено ({self._restorer.name}): {self._current.name} — {len(dets)} обл."
+            f"Расцензурено ({restorer.name}): {self._current.name} — {len(dets)} обл."
         )
 
+    def _make_restorer(self):
+        key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,
+               self._cfg.dm_python, self._cfg.dm_gpu)
+        if key != self._restorer_key:
+            self._restorer = build_restorer(self._cfg.restorer, self._cfg)  # may raise
+            self._restorer_key = key
+        return self._restorer
+
+    def _open_restore_settings(self) -> None:
+        dlg = RestoreDialog(self._cfg, self)
+        if dlg.exec() != QDialog.Accepted:
+            return
+        dlg.apply_to_config()
+        settings_store.save(self._cfg)
+        self._restorer_key = None  # rebuild on next restore
+        self.statusBar().showMessage(f"Движок восстановления: {self._cfg.restorer}")
+
     def _toggle_restored(self) -> None:
         if self._current is None or str(self._current) not in self._restored:
             return
@@ -664,6 +694,7 @@ class MainWindow(QMainWindow):
             self._show(Path(cur.data(Qt.UserRole)))
         elif self.file_list.count() == 0:
             self.view.set_image(None, [])
+        self._refresh_marks()  # rows shifted — remap marks to new indices
         self._update_nav()
 
     @staticmethod
@@ -679,14 +710,27 @@ class MainWindow(QMainWindow):
         return folder / f"{stem} ({i}){suffix}"
 
     # ----------------------------------------------------------------- detail
+    # Row tints in the file list: red = censorship found, green = checked & clean.
+    _TINT_HIT = QColor(200, 80, 80, 70)
+    _TINT_CLEAN = QColor(90, 160, 90, 50)
+
     def _tag_file(self, path: Path, count: int) -> None:
         for i in range(self.file_list.count()):
             item = self.file_list.item(i)
             if item.data(Qt.UserRole) == str(path):
                 base = item.data(Qt.UserRole + 1)
                 item.setText(f"{base}  ·  {count}" if count else f"{base}  ·  —")
+                item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN)
                 return
 
+    def _refresh_marks(self) -> None:
+        """Project frames-with-detections onto the scrubber as marks."""
+        marks = {
+            i for i in range(self.file_list.count())
+            if self._results.get(self.file_list.item(i).data(Qt.UserRole))
+        }
+        self.frame_slider.set_marks(marks)
+
     def _fill_detail_table(self, path: Path, img, dets: list[Detection] | None) -> None:
         h, w = (img.shape[0], img.shape[1]) if img is not None else (0, 0)
         if dets is None:

hvideotool/ui/marker_slider.py

@@ -0,0 +1,66 @@
+"""A horizontal QSlider that paints marks at frames where censorship was found.
+
+Used as the navigation scrubber under the image: the file list / `_results`
+cache is projected onto the groove as small vertical ticks, so you can see at a
+glance where the detected regions are along the whole sequence.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from PySide6.QtCore import Qt
+from PySide6.QtGui import QColor, QPainter
+from PySide6.QtWidgets import QSlider, QStyle, QStyleOptionSlider
+
+
+class MarkerSlider(QSlider):
+    def __init__(self, orientation=Qt.Horizontal, parent=None) -> None:
+        super().__init__(orientation, parent)
+        self._marks: set[int] = set()
+        self._mark_color = QColor(220, 70, 70)
+
+    def set_marks(self, marks: Iterable[int]) -> None:
+        marks = set(marks)
+        if marks != self._marks:
+            self._marks = marks
+            self.update()
+
+    def clear_marks(self) -> None:
+        if self._marks:
+            self._marks = set()
+            self.update()
+
+    def paintEvent(self, event) -> None:
+        super().paintEvent(event)
+        if not self._marks or self.maximum() <= self.minimum():
+            return
+
+        opt = QStyleOptionSlider()
+        self.initStyleOption(opt)
+        groove = self.style().subControlRect(
+            QStyle.CC_Slider, opt, QStyle.SC_SliderGroove, self
+        )
+        handle = self.style().subControlRect(
+            QStyle.CC_Slider, opt, QStyle.SC_SliderHandle, self
+        )
+        span = groove.width() - handle.width()
+        if span <= 0:
+            return
+        lo, hi = self.minimum(), self.maximum()
+        half = handle.width() // 2
+        top = groove.center().y() - 5
+        bottom = groove.center().y() + 5
+
+        painter = QPainter(self)
+        painter.setPen(self._mark_color)
+        # Many frames can collapse onto the same pixel column — dedupe to keep
+        # repaint cheap on big folders (tens of thousands of frames).
+        seen_x: set[int] = set()
+        for m in self._marks:
+            pos = QStyle.sliderPositionFromValue(lo, hi, m, span, opt.upsideDown)
+            x = groove.x() + half + pos
+            if x not in seen_x:
+                seen_x.add(x)
+                painter.drawLine(x, top, x, bottom)
+        painter.end()

hvideotool/ui/restore_dialog.py

@@ -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"