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: