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

Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.

Browse Source
This commit is contained in:
Leonid Pershin
2026-06-07 06:16:05 +03:00
parent cc518cc3e6
commit 9c471ca701
17 changed files with 798 additions and 676 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hvideotool/ui/main_window.py
+156 -68
View File
@@ -26,11 +26,10 @@ import shutil
from pathlib import Path
from PySide6.QtCore import Qt, QThreadPool
from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcut
from PySide6.QtWidgets import (
QAbstractItemView,
QApplication,
QComboBox,
QDialog,
QDoubleSpinBox,
QFileDialog,
@@ -41,6 +40,7 @@ from PySide6.QtWidgets import (
QListWidgetItem,
QMainWindow,
QMessageBox,
QPlainTextEdit,
QProgressBar,
QPushButton,
QSplitter,
@@ -51,7 +51,7 @@ from PySide6.QtWidgets import (
)
from .. import settings_store
from ..config import AppConfig
from ..config import AppConfig, normalize_config
from ..core.detection import cache as detection_cache
from ..core.detection.factory import build_detector
from ..core.detection.types import Detection
@@ -68,7 +68,6 @@ from .workers import Job
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
_DETECTORS = ["classic", "yolo", "combined"]
class MainWindow(QMainWindow):
@@ -120,6 +119,8 @@ class MainWindow(QMainWindow):
file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True))
file_menu.addSeparator()
file_menu.addAction("Движок восстановления…", self._open_restore_settings)
file_menu.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
file_menu.addSeparator()
file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M")
file_menu.addSeparator()
@@ -136,14 +137,9 @@ class MainWindow(QMainWindow):
tb.addAction(from_video)
tb.addSeparator()
tb.addWidget(QLabel(" Детектор: "))
self.detector_combo = QComboBox()
self.detector_combo.addItems(_DETECTORS)
self.detector_combo.setCurrentText(self._cfg.detector)
self.detector_combo.currentTextChanged.connect(self._on_detector_changed)
tb.addWidget(self.detector_combo)
tb.addWidget(QLabel(" Детектор: YOLO "))
self.model_action = QAction("Модель…", self, triggered=self._choose_model)
self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики")
tb.addAction(self.model_action)
tb.addSeparator()
@@ -166,6 +162,12 @@ class MainWindow(QMainWindow):
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
restore.setToolTip("Восстановить найденные области на текущем кадре")
tb.addAction(restore)
restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False))
restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)")
tb.addAction(restore_all)
restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True))
restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)")
tb.addAction(restore_regen)
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
self.toggle_restored_action.setEnabled(False)
tb.addAction(self.toggle_restored_action)
@@ -359,27 +361,58 @@ class MainWindow(QMainWindow):
info = self._device_info if self._device_info else torch_info.gather()
cuda = bool(info.get("cuda_available"))
a = torch_info.analyze(info)
self._install_command = a["command"] # what the Copy button will copy
lines = [
torch_info.reason(info),
"ВЕРДИКТ:",
a["summary"],
"",
"Диагностика:",
f" • PyTorch: {info.get('version') or 'не установлен'}",
f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}",
f" • CUDA доступна: {'да' if cuda else 'нет'}",
*(f" • {d}" for d in a["details"]),
"",
"Что делать:",
a["steps"],
]
if info.get("device_name"):
lines.append(f" • GPU: {info['device_name']}")
if info.get("import_error"):
lines.append(f" • Ошибка импорта torch: {info['import_error']}")
box = QMessageBox(self)
box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning)
box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU"))
box.setText("\n".join(lines))
# A real dialog (not QMessageBox) so the text — incl. the install command — is
# selectable, and a Copy button drops the pip command straight onto the clipboard.
dlg = QDialog(self)
dlg.setWindowTitle("Почему " + ("GPU" if cuda else "CPU") + " — диагностика PyTorch/CUDA")
dlg.resize(620, 480)
layout = QVBoxLayout(dlg)
text = QPlainTextEdit()
text.setReadOnly(True)
text.setPlainText("\n".join(lines))
mono = QFont("Consolas")
mono.setStyleHint(QFont.Monospace)
text.setFont(mono)
layout.addWidget(text, 1)
buttons = QHBoxLayout()
if not cuda:
box.setInformativeText(torch_info.install_hint())
box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands
box.exec()
copy_btn = QPushButton("Скопировать команду установки")
copy_btn.clicked.connect(self._copy_install_command)
buttons.addWidget(copy_btn)
recheck = QPushButton("Проверить заново")
recheck.setToolTip("Перепроверить torch/CUDA (например, после переустановки)")
recheck.clicked.connect(lambda: (self._probe_device(), dlg.accept()))
buttons.addWidget(recheck)
buttons.addStretch(1)
close_btn = QPushButton("Закрыть")
close_btn.clicked.connect(dlg.accept)
buttons.addWidget(close_btn)
layout.addLayout(buttons)
dlg.exec()
def _copy_install_command(self) -> None:
command = getattr(self, "_install_command", None)
if not command:
from ..core import torch_info
command = torch_info.install_command()
QApplication.clipboard().setText(command)
self.statusBar().showMessage("Команда установки скопирована в буфер обмена")
# ------------------------------------------------------------- cancellation
def _begin_busy(self, total: int | None = None) -> None:
@@ -388,7 +421,6 @@ class MainWindow(QMainWindow):
self._cancel = False
self.stop_action.setEnabled(True)
# Disable inputs that would race a running job (they clear cache / rebuild engines).
self.detector_combo.setEnabled(False)
self.model_action.setEnabled(False)
if total is None:
self.progress.setRange(0, 0) # indeterminate
@@ -400,7 +432,6 @@ class MainWindow(QMainWindow):
def _end_busy(self) -> None:
self._busy = False
self.stop_action.setEnabled(False)
self.detector_combo.setEnabled(True)
self.model_action.setEnabled(True)
self.progress.setVisible(False)
self.progress.setRange(0, 100) # leave it determinate for the next user
@@ -459,19 +490,19 @@ class MainWindow(QMainWindow):
self._detector_key = key
return self._detector
def _on_detector_changed(self, name: str) -> None:
self._cfg.detector = name
# YOLO/combined need a model. Auto-pick a known one from models/ if we have it;
# only prompt when nothing suitable is found (don't nag when the path is obvious).
if name in ("yolo", "combined") and not self._cfg.model_path:
found = self._auto_find_model()
if found:
self._cfg.model_path = found
self.statusBar().showMessage(f"Модель найдена автоматически: {found}")
else:
self._choose_model()
self._persist_settings()
self._invalidate_results()
def _ensure_model(self) -> None:
"""Make sure the YOLO detector has weights — auto-pick from ./models silently.
Called on project open. Doesn't prompt (the user can pick via "Модель…"); the
detector factory raises a clear message if a detect is attempted without one.
"""
if self._cfg.model_path and Path(self._cfg.model_path).is_file():
return
found = self._auto_find_model()
if found:
self._cfg.model_path = found
self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}")
self._persist_settings()
@staticmethod
def _auto_find_model() -> str | None:
@@ -667,6 +698,8 @@ class MainWindow(QMainWindow):
self._project = project
project.frames_dir.mkdir(parents=True, exist_ok=True)
project.apply_to_config(self._cfg) # per-project settings -> live config
normalize_config(self._cfg) # coerce any legacy classic/inpaint values
self._ensure_model() # YOLO needs weights — auto-pick if missing
self._sync_settings_ui()
self._detector_key = None
self._restorer_key = None
@@ -679,9 +712,6 @@ class MainWindow(QMainWindow):
def _sync_settings_ui(self) -> None:
"""Reflect the (project's) config onto the toolbar widgets without signal loops."""
self.detector_combo.blockSignals(True)
self.detector_combo.setCurrentText(self._cfg.detector)
self.detector_combo.blockSignals(False)
self.threshold_spin.blockSignals(True)
self.threshold_spin.setValue(self._cfg.default_threshold)
self.threshold_spin.blockSignals(False)
@@ -912,41 +942,32 @@ class MainWindow(QMainWindow):
key, dets = payload
self._results[key] = dets
self._tag_file(Path(key), len(dets))
# If the frame being viewed was just computed, show its overlay live.
if not self._showing_restored and self._current is not None and str(self._current) == key:
self._show(self._current)
self._tick_count += 1
if self._tick_count % 25 == 0:
self._refresh_marks() # let marks appear progressively (throttled)
# ------------------------------------------------------------- restoration
def _restore_current(self) -> None:
"""Restore the current frame's regions on a background thread, then show it.
"""Restore the current frame on a background thread, then show it.
Detections are computed first (in the same job) if not cached. The DeepMosaics
engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
DeepMosaics locates the mosaic itself, so no detection step is needed — we just
run the engine on the frame (if there's no mosaic the frame comes back unchanged).
The engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
if self._current is None or self._busy:
return
path = self._current
key = str(path)
def fn(job):
dets = self._results.get(key)
if dets is None:
dets = self._compute(self._make_detector(), path)
job.tick(("dets", key, dets)) # cache them on the GUI thread
if not dets:
return ("empty", key)
img = imread_unicode(key)
if img is None:
raise RuntimeError(f"Не удалось прочитать: {path.name}")
restorer = self._make_restorer()
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
return ("restored", key, restored, len(dets), restorer.name)
def tick(payload):
if payload[0] == "dets":
_, k, dets = payload
self._results[k] = dets
self._tag_file(Path(k), len(dets))
self._refresh_marks()
restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled)
return ("restored", key, restored, restorer.name)
def done(result, cancelled):
if cancelled:
@@ -954,19 +975,86 @@ class MainWindow(QMainWindow):
return
if result is None:
return
if result[0] == "empty":
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
return
_, k, restored, n, engine = result
_, k, restored, engine = result
self._restored[k] = restored
if self._current is not None and str(self._current) == k:
self._showing_restored = True
self.view.set_image(restored, [])
self._update_restore_actions()
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name} — {n} обл.")
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}")
self.statusBar().showMessage(f"Восстановление: {path.name}…")
self._start_job(fn, None, on_tick=tick, on_done=done)
self._start_job(fn, None, on_done=done)
def _restore_all(self, force: bool = False) -> None:
"""Restore every frame on a background thread, writing results to ``restored/``.
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
engine skips frames already restored (resume) unless ``force``. The temporal
engine (DeepMosaics-video) runs the whole contiguous sequence in order via
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
"""
if not self._files or self._project is None or self._busy:
return
files = list(self._files) # snapshot — favorites/move mutate self._files
total = len(files)
out_dir = self._project.restored_dir
out_dir.mkdir(parents=True, exist_ok=True)
def out_path(p: Path) -> Path:
return out_dir / f"{p.stem}.jpg"
def fn(job):
restorer = self._make_restorer() # built on the worker (may raise)
frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads
def get_frame(i):
img = frame_cache.get(i)
if img is None:
img = imread_unicode(str(files[i]))
if img is None:
raise RuntimeError(f"Не удалось прочитать: {files[i].name}")
if len(frame_cache) > 24:
frame_cache.clear()
frame_cache[i] = img
return img
def emit(i, restored):
imwrite_unicode(str(out_path(files[i])), restored)
job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}")
if restorer.temporal:
restorer.restore_sequence(
total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled
)
else:
for i, p in enumerate(files):
if job.cancelled:
break
if not force and out_path(p).is_file():
job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}")
continue
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
frame_cache.pop(i, None) # per-frame: don't accumulate
return None
def done(_result, cancelled):
if self._current is not None: # live-preview the current frame's result, if any
rp = out_path(self._current)
if rp.is_file():
img = imread_unicode(str(rp))
if img is not None:
self._restored[str(self._current)] = img
self._showing_restored = True
self.view.set_image(img, [])
self._update_restore_actions()
self.statusBar().showMessage(
"Расцензуривание отменено" if cancelled
else f"Готово: результаты в {out_dir.name}/ ({total} кадров)"
)
self.statusBar().showMessage("Пакетное расцензуривание…")
self._start_job(fn, total, on_done=done)
def _make_restorer(self):
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,