mrleo1nid/HVideoTool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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.

Browse Source
This commit is contained in:
Leonid Pershin
2026-06-07 04:02:25 +03:00
parent ddc8543647
commit 7f0121b7df
9 changed files with 413 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hvideotool/ui/main_window.py
+52 -8
View File
@@ -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
+66
View File
@@ -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
+105
View File
@@ -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"