diff --git a/CLAUDE.md b/CLAUDE.md index 0de80f3..97ac7e3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -98,8 +98,9 @@ hvideotool/ │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type ├── restore/ # "un-censor" detected regions (per-frame) │ ├── base.py # Restorer ABC: restore(image, detections) -> image - │ ├── factory.py # build_restorer(name) -> inpaint (deepmosaics/lada = not wired yet) + │ ├── factory.py # build_restorer(name, config) -> inpaint | deepmosaics (lada = TODO) │ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs + │ ├── deepmosaics.py # DeepMosaicsRestorer — shells out to user-installed DeepMosaics │ └── mask.py # detections_to_mask(shape, dets, dilate) └── detection/ ├── base.py # Detector ABC: detect(frame) -> list[Detection] @@ -142,9 +143,15 @@ hvideotool/ caches the result in `_restored[path]`, and shows it overlay-free. "Показать оригинал/ результат" toggles (`_showing_restored`); "Сохранить результат" writes `_restored.jpg` into the active collection (or beside the frame). The baseline - is cv2 inpaint — honest: it fills, doesn't reconstruct. To add a real engine, - implement `core/restore/base.Restorer`, register it in `restore/factory.build_restorer`, - and swap `self._restorer`. `_show` resets `_showing_restored` + `_update_restore_actions`. + is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), which + shells out to a user-installed DeepMosaics repo (`deepmosaic.py --mode clean`, stdin + closed so its error `input()` can't hang; reads the newest image from a temp + result_dir). The engine + paths are set in `RestoreDialog` (Файл → Движок + восстановления…), persisted, and built lazily/cached in `_make_restorer` (like + `_make_detector`). NOTE: DeepMosaics locates mosaics itself (its `mosaic_position.pth`) + — our detections aren't passed to it. `_show` resets `_showing_restored` + + `_update_restore_actions`. To add another engine (e.g. LADA), implement + `core/restore/base.Restorer` and register it in `restore/factory.build_restorer`. - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, diff --git a/README.md b/README.md index 37ffa35..3c72887 100644 --- a/README.md +++ b/README.md @@ -35,14 +35,42 @@ Кнопка **«Расцензурить кадр»** восстанавливает найденные области на текущем кадре, **«Показать оригинал/результат»** переключает вид, **«Сохранить результат»** пишет -`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром). +`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром). Движок выбирается в +меню **Файл → Движок восстановления…**. -> ⚠️ Сейчас движок восстановления — **инпейнт (cv2)**: он *заполняет* область по -> окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). -> Настоящее восстановление мозаики требует генеративной модели (DeepMosaics или -> [LADA](https://github.com/ladaapp/lada) на BasicVSR++) и **GPU NVIDIA/CUDA**; на -> аниме качество таких моделей ограничено. Движок подключается через интерфейс -> `core/restore/base.Restorer` (`build_restorer`) — инпейнт сейчас стоит заглушкой. +Два движка: + +- **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по + окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). +- **DeepMosaics** — реальное генеративное удаление мозаики. Требует **GPU NVIDIA/CUDA** + (на CPU очень медленно) и отдельной установки. На аниме качество ограничено + (модели обучены на реальном видео). + +### Настройка DeepMosaics + +```powershell +git clone https://github.com/HypoX64/DeepMosaics +# установите зависимости DeepMosaics (см. его README; нужен torch с CUDA) +# скачайте веса (clean_youknow_resnet_9blocks.pth + mosaic_position.pth) в pretrained_models/mosaic +``` + +Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, укажите папку +DeepMosaics (с `deepmosaic.py`), файл весов и GPU id (`-1` = CPU). Приложение вызывает +DeepMosaics на текущем кадре и показывает результат. + +> **Важно:** для покадрового режима берите **картиночную** модель +> `clean_youknow_resnet_9blocks.pth`. Видеомодель `clean_youknow_video.pth` (BVDNet) +> покадрово **не работает** — ей нужен соседний кадр. Если кадр без мозаики, движок +> вернёт его без изменений. +> +> DeepMosaics 2021 года рассчитан на старые версии (`torch 1.7`, `numpy 1.19`). +> Проверено: на современных `torch 2.x`/`numpy 2.x` картиночная модель запускается +> (CPU ~7 с/кадр), но если столкнётесь с несовместимостью — заведите для DeepMosaics +> отдельное окружение по его `requirements.txt` и укажите его `python.exe` в диалоге. + +> Движок подключается через интерфейс `core/restore/base.Restorer` (`build_restorer`). +> [LADA](https://github.com/ladaapp/lada) (BasicVSR++, лучшее качество на реальном +> видео, но видеомодель) пока не подключён. ## Что НЕ делает (осознанно вне области задачи) diff --git a/hvideotool/config.py b/hvideotool/config.py index 64d81cd..66c0075 100644 --- a/hvideotool/config.py +++ b/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 diff --git a/hvideotool/core/restore/deepmosaics.py b/hvideotool/core/restore/deepmosaics.py new file mode 100644 index 0000000..1442b1d --- /dev/null +++ b/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) + ) diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py index 1fc6c4f..464509e 100644 --- a/hvideotool/core/restore/factory.py +++ b/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}") diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py index d4daa59..0274bed 100644 --- a/hvideotool/settings_store.py +++ b/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) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index fec0550..e0ed7ca 100644 --- a/hvideotool/ui/main_window.py +++ b/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: diff --git a/hvideotool/ui/marker_slider.py b/hvideotool/ui/marker_slider.py new file mode 100644 index 0000000..8690af7 --- /dev/null +++ b/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() diff --git a/hvideotool/ui/restore_dialog.py b/hvideotool/ui/restore_dialog.py new file mode 100644 index 0000000..a1a8870 --- /dev/null +++ b/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"