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.
This commit is contained in:
@@ -98,8 +98,9 @@ hvideotool/
|
|||||||
│ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
|
│ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
|
||||||
├── restore/ # "un-censor" detected regions (per-frame)
|
├── restore/ # "un-censor" detected regions (per-frame)
|
||||||
│ ├── base.py # Restorer ABC: restore(image, detections) -> image
|
│ ├── 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
|
│ ├── 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)
|
│ └── mask.py # detections_to_mask(shape, dets, dilate)
|
||||||
└── detection/
|
└── detection/
|
||||||
├── base.py # Detector ABC: detect(frame) -> list[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. "Показать оригинал/
|
caches the result in `_restored[path]`, and shows it overlay-free. "Показать оригинал/
|
||||||
результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
||||||
`<stem>_restored.jpg` into the active collection (or beside the frame). The baseline
|
`<stem>_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,
|
is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), which
|
||||||
implement `core/restore/base.Restorer`, register it in `restore/factory.build_restorer`,
|
shells out to a user-installed DeepMosaics repo (`deepmosaic.py --mode clean`, stdin
|
||||||
and swap `self._restorer`. `_show` resets `_showing_restored` + `_update_restore_actions`.
|
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
|
- **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys
|
||||||
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label`
|
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label`
|
||||||
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
|
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
|
||||||
|
|||||||
@@ -35,14 +35,42 @@
|
|||||||
|
|
||||||
Кнопка **«Расцензурить кадр»** восстанавливает найденные области на текущем кадре,
|
Кнопка **«Расцензурить кадр»** восстанавливает найденные области на текущем кадре,
|
||||||
**«Показать оригинал/результат»** переключает вид, **«Сохранить результат»** пишет
|
**«Показать оригинал/результат»** переключает вид, **«Сохранить результат»** пишет
|
||||||
`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром).
|
`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром). Движок выбирается в
|
||||||
|
меню **Файл → Движок восстановления…**.
|
||||||
|
|
||||||
> ⚠️ Сейчас движок восстановления — **инпейнт (cv2)**: он *заполняет* область по
|
Два движка:
|
||||||
> окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает).
|
|
||||||
> Настоящее восстановление мозаики требует генеративной модели (DeepMosaics или
|
- **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по
|
||||||
> [LADA](https://github.com/ladaapp/lada) на BasicVSR++) и **GPU NVIDIA/CUDA**; на
|
окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает).
|
||||||
> аниме качество таких моделей ограничено. Движок подключается через интерфейс
|
- **DeepMosaics** — реальное генеративное удаление мозаики. Требует **GPU NVIDIA/CUDA**
|
||||||
> `core/restore/base.Restorer` (`build_restorer`) — инпейнт сейчас стоит заглушкой.
|
(на 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++, лучшее качество на реальном
|
||||||
|
> видео, но видеомодель) пока не подключён.
|
||||||
|
|
||||||
## Что НЕ делает (осознанно вне области задачи)
|
## Что НЕ делает (осознанно вне области задачи)
|
||||||
|
|
||||||
|
|||||||
@@ -65,3 +65,10 @@ class AppConfig:
|
|||||||
detector: str = "classic" # "classic" | "yolo" | "combined"
|
detector: str = "classic" # "classic" | "yolo" | "combined"
|
||||||
model_path: str | None = None # weights path, used by the YOLO detector
|
model_path: str | None = None # weights path, used by the YOLO detector
|
||||||
default_threshold: float = 0.20 # initial overlay confidence threshold
|
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
|
||||||
|
|||||||
@@ -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
|
- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct).
|
||||||
(DeepMosaics / LADA BasicVSR++) are placeholders — they need model weights and a
|
- ``deepmosaics``: real generative mosaic removal via a user-installed DeepMosaics
|
||||||
CUDA GPU, and raise a clear, actionable error until integrated. See README.
|
(subprocess). Needs the DeepMosaics folder + clean weights + (ideally) a CUDA GPU.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from .base import Restorer
|
from .base import Restorer
|
||||||
from .inpaint import InpaintRestorer
|
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":
|
if name == "inpaint":
|
||||||
return InpaintRestorer()
|
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(
|
raise ValueError(
|
||||||
"Генеративное восстановление пока не подключено.\n"
|
"Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README."
|
||||||
"Нужна модель (DeepMosaics / LADA) и GPU (CUDA). См. README → Восстановление."
|
|
||||||
)
|
)
|
||||||
raise ValueError(f"Неизвестный режим восстановления: {name!r}")
|
raise ValueError(f"Неизвестный режим восстановления: {name!r}")
|
||||||
|
|||||||
@@ -35,6 +35,11 @@ def apply(config: AppConfig) -> None:
|
|||||||
config.model_path = data["model_path"]
|
config.model_path = data["model_path"]
|
||||||
if "threshold" in data:
|
if "threshold" in data:
|
||||||
config.default_threshold = float(data["threshold"])
|
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:
|
def save(config: AppConfig) -> None:
|
||||||
@@ -44,6 +49,11 @@ def save(config: AppConfig) -> None:
|
|||||||
detector=config.detector,
|
detector=config.detector,
|
||||||
model_path=config.model_path,
|
model_path=config.model_path,
|
||||||
threshold=config.default_threshold,
|
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)
|
_write(data)
|
||||||
|
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ import shutil
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from PySide6.QtCore import Qt
|
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 (
|
from PySide6.QtWidgets import (
|
||||||
QAbstractItemView,
|
QAbstractItemView,
|
||||||
QApplication,
|
QApplication,
|
||||||
@@ -38,7 +38,6 @@ from PySide6.QtWidgets import (
|
|||||||
QMessageBox,
|
QMessageBox,
|
||||||
QProgressBar,
|
QProgressBar,
|
||||||
QPushButton,
|
QPushButton,
|
||||||
QSlider,
|
|
||||||
QSplitter,
|
QSplitter,
|
||||||
QTableWidget,
|
QTableWidget,
|
||||||
QTableWidgetItem,
|
QTableWidgetItem,
|
||||||
@@ -56,6 +55,8 @@ from ..core.video.extract import extract_frames
|
|||||||
from ..core.video.frame import Frame
|
from ..core.video.frame import Frame
|
||||||
from .extract_dialog import ExtractDialog
|
from .extract_dialog import ExtractDialog
|
||||||
from .image_view import ImageView
|
from .image_view import ImageView
|
||||||
|
from .marker_slider import MarkerSlider
|
||||||
|
from .restore_dialog import RestoreDialog
|
||||||
|
|
||||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||||||
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
|
_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._results: dict[str, list[Detection]] = {} # path -> detections (cache)
|
||||||
self._current: Path | None = None
|
self._current: Path | None = None
|
||||||
self._collection: Path | None = None # active destination folder for moves
|
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._restored: dict[str, "object"] = {} # path -> restored image (BGR ndarray)
|
||||||
self._showing_restored = False
|
self._showing_restored = False
|
||||||
self._nav_sync = False # guard against slider<->list signal loops
|
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._recompute_current).setShortcut("Space")
|
||||||
file_menu.addAction("Детектировать все", self._detect_all)
|
file_menu.addAction("Детектировать все", self._detect_all)
|
||||||
file_menu.addSeparator()
|
file_menu.addSeparator()
|
||||||
|
file_menu.addAction("Движок восстановления…", self._open_restore_settings)
|
||||||
|
file_menu.addSeparator()
|
||||||
file_menu.addAction("Создать коллекцию…", self._create_collection)
|
file_menu.addAction("Создать коллекцию…", self._create_collection)
|
||||||
file_menu.addAction("В коллекцию", self._move_to_collection).setShortcut("Ctrl+M")
|
file_menu.addAction("В коллекцию", self._move_to_collection).setShortcut("Ctrl+M")
|
||||||
file_menu.addSeparator()
|
file_menu.addSeparator()
|
||||||
@@ -223,10 +227,10 @@ class MainWindow(QMainWindow):
|
|||||||
self.next_btn.setToolTip("Следующий кадр (.)")
|
self.next_btn.setToolTip("Следующий кадр (.)")
|
||||||
self.next_btn.clicked.connect(lambda: self._step(1))
|
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.setMinimum(0)
|
||||||
self.frame_slider.setMaximum(0)
|
self.frame_slider.setMaximum(0)
|
||||||
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →)")
|
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
|
||||||
self.frame_slider.valueChanged.connect(self._on_slider)
|
self.frame_slider.valueChanged.connect(self._on_slider)
|
||||||
|
|
||||||
self.pos_label = QLabel("0 / 0")
|
self.pos_label = QLabel("0 / 0")
|
||||||
@@ -329,7 +333,10 @@ class MainWindow(QMainWindow):
|
|||||||
self._detector_key = None
|
self._detector_key = None
|
||||||
self._results.clear()
|
self._results.clear()
|
||||||
for i in range(self.file_list.count()):
|
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:
|
if self._current is not None:
|
||||||
self._show(self._current)
|
self._show(self._current)
|
||||||
|
|
||||||
@@ -416,6 +423,7 @@ class MainWindow(QMainWindow):
|
|||||||
self.file_list.setUpdatesEnabled(True)
|
self.file_list.setUpdatesEnabled(True)
|
||||||
self.file_list.blockSignals(False)
|
self.file_list.blockSignals(False)
|
||||||
self.progress.setVisible(False)
|
self.progress.setVisible(False)
|
||||||
|
self._refresh_marks()
|
||||||
|
|
||||||
self._refresh_collections()
|
self._refresh_collections()
|
||||||
if not files:
|
if not files:
|
||||||
@@ -436,6 +444,7 @@ class MainWindow(QMainWindow):
|
|||||||
path = Path(item.data(Qt.UserRole))
|
path = Path(item.data(Qt.UserRole))
|
||||||
if str(path) not in self._results and self._detect(path) is None:
|
if str(path) not in self._results and self._detect(path) is None:
|
||||||
return
|
return
|
||||||
|
self._refresh_marks()
|
||||||
self._show(path)
|
self._show(path)
|
||||||
|
|
||||||
def _detect(self, path: Path) -> list[Detection] | None:
|
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
|
self._detector_key = None # rebuild the detector so settings changes take effect
|
||||||
if self._detect(self._current) is None:
|
if self._detect(self._current) is None:
|
||||||
return
|
return
|
||||||
|
self._refresh_marks()
|
||||||
self._show(self._current)
|
self._show(self._current)
|
||||||
|
|
||||||
def _detect_all(self) -> None:
|
def _detect_all(self) -> None:
|
||||||
@@ -496,6 +506,7 @@ class MainWindow(QMainWindow):
|
|||||||
finally:
|
finally:
|
||||||
self.progress.setVisible(False)
|
self.progress.setVisible(False)
|
||||||
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
||||||
|
self._refresh_marks()
|
||||||
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
||||||
if self._current is not None:
|
if self._current is not None:
|
||||||
self._show(self._current)
|
self._show(self._current)
|
||||||
@@ -508,6 +519,7 @@ class MainWindow(QMainWindow):
|
|||||||
key = str(self._current)
|
key = str(self._current)
|
||||||
if key not in self._results and self._detect(self._current) is None:
|
if key not in self._results and self._detect(self._current) is None:
|
||||||
return
|
return
|
||||||
|
self._refresh_marks()
|
||||||
dets = self._results.get(key) or []
|
dets = self._results.get(key) or []
|
||||||
if not dets:
|
if not dets:
|
||||||
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
|
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
|
||||||
@@ -518,7 +530,8 @@ class MainWindow(QMainWindow):
|
|||||||
self.statusBar().showMessage(f"Восстановление: {self._current.name}…")
|
self.statusBar().showMessage(f"Восстановление: {self._current.name}…")
|
||||||
QApplication.processEvents()
|
QApplication.processEvents()
|
||||||
try:
|
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
|
except Exception as exc: # noqa: BLE001 - surface model/engine errors
|
||||||
QMessageBox.warning(self, "Ошибка восстановления", str(exc))
|
QMessageBox.warning(self, "Ошибка восстановления", str(exc))
|
||||||
return
|
return
|
||||||
@@ -527,9 +540,26 @@ class MainWindow(QMainWindow):
|
|||||||
self.view.set_image(restored, [])
|
self.view.set_image(restored, [])
|
||||||
self._update_restore_actions()
|
self._update_restore_actions()
|
||||||
self.statusBar().showMessage(
|
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:
|
def _toggle_restored(self) -> None:
|
||||||
if self._current is None or str(self._current) not in self._restored:
|
if self._current is None or str(self._current) not in self._restored:
|
||||||
return
|
return
|
||||||
@@ -664,6 +694,7 @@ class MainWindow(QMainWindow):
|
|||||||
self._show(Path(cur.data(Qt.UserRole)))
|
self._show(Path(cur.data(Qt.UserRole)))
|
||||||
elif self.file_list.count() == 0:
|
elif self.file_list.count() == 0:
|
||||||
self.view.set_image(None, [])
|
self.view.set_image(None, [])
|
||||||
|
self._refresh_marks() # rows shifted — remap marks to new indices
|
||||||
self._update_nav()
|
self._update_nav()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@@ -679,14 +710,27 @@ class MainWindow(QMainWindow):
|
|||||||
return folder / f"{stem} ({i}){suffix}"
|
return folder / f"{stem} ({i}){suffix}"
|
||||||
|
|
||||||
# ----------------------------------------------------------------- detail
|
# ----------------------------------------------------------------- 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:
|
def _tag_file(self, path: Path, count: int) -> None:
|
||||||
for i in range(self.file_list.count()):
|
for i in range(self.file_list.count()):
|
||||||
item = self.file_list.item(i)
|
item = self.file_list.item(i)
|
||||||
if item.data(Qt.UserRole) == str(path):
|
if item.data(Qt.UserRole) == str(path):
|
||||||
base = item.data(Qt.UserRole + 1)
|
base = item.data(Qt.UserRole + 1)
|
||||||
item.setText(f"{base} · {count}" if count else f"{base} · —")
|
item.setText(f"{base} · {count}" if count else f"{base} · —")
|
||||||
|
item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN)
|
||||||
return
|
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:
|
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)
|
h, w = (img.shape[0], img.shape[1]) if img is not None else (0, 0)
|
||||||
if dets is None:
|
if dets is None:
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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"
|
||||||
Reference in New Issue
Block a user