Refactor HVideoTool to support project-based workflow: introduced project management features, updated UI for project handling, and enhanced documentation in README and CLAUDE.md. The tool now organizes images and settings into projects, improving usability and detection caching.
This commit is contained in:
+416
-162
@@ -1,18 +1,23 @@
|
||||
"""Main window: open a folder of images and inspect what the detector found.
|
||||
"""Main window: open a **project** and inspect what the detector found.
|
||||
|
||||
Layout: a toolbar (open folder · from-video · detector · model · calc-frame ·
|
||||
A project is a folder (``project.json`` + ``frames/`` + ``detections.json`` +
|
||||
``collections/``) — see ``core/project.py``. The per-project settings (detector,
|
||||
model, threshold, restore engine) live in ``project.json``; the global
|
||||
``settings.json`` only seeds defaults for new projects.
|
||||
|
||||
Layout: a toolbar (new/open project · from-video · detector · model · calc-frame ·
|
||||
detect-all · threshold), then a splitter with three panes — left: collection
|
||||
controls + the file list; center: the image with overlays; right: a detail table
|
||||
of every detection. Collection controls sit by the file list (they act on its
|
||||
selection), keeping the toolbar to detection/entry actions only.
|
||||
|
||||
Viewing and detecting are decoupled, so browsing a big folder stays instant even
|
||||
Viewing and detecting are decoupled, so browsing a big project stays instant even
|
||||
with a slow (CPU) detector:
|
||||
- selecting a file just **shows** it (with its cached result, if any);
|
||||
- **double-clicking** a file, or "Рассчитать кадр", runs the detector on it;
|
||||
- "Детектировать все" runs the whole folder.
|
||||
Both folder loading and detect-all show a progress bar. Results are cached;
|
||||
switching detector/model clears the cache.
|
||||
- "Детектировать все" runs the whole project.
|
||||
Both project loading and detect-all show a progress bar. Results are cached in the
|
||||
project; switching detector/model clears the cache.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -47,9 +52,12 @@ from PySide6.QtWidgets import (
|
||||
|
||||
from .. import settings_store
|
||||
from ..config import AppConfig
|
||||
from ..core.detection import cache as detection_cache
|
||||
from ..core.detection.factory import build_detector
|
||||
from ..core.detection.types import Detection
|
||||
from ..core.imageio import imread_unicode, imwrite_unicode
|
||||
from ..core.project import PROJECT_FILE, Project
|
||||
from ..core.restore.base import Cancelled
|
||||
from ..core.restore.factory import build_restorer
|
||||
from ..core.video.extract import extract_frames
|
||||
from ..core.video.frame import Frame
|
||||
@@ -69,16 +77,18 @@ class MainWindow(QMainWindow):
|
||||
self._cfg = config
|
||||
self._detector = None
|
||||
self._detector_key = None
|
||||
self._folder: Path | None = None
|
||||
self._project: Project | None = None # the open project (None until one is opened)
|
||||
self._folder: Path | None = None # == project.frames_dir while a project is open
|
||||
self._files: list[Path] = []
|
||||
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 = 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
|
||||
self._busy = False # a long operation is running
|
||||
self._cancel = False # the user asked to stop it
|
||||
|
||||
self.setWindowTitle("HVideoTool — инспектор детекции цензуры")
|
||||
self.resize(1180, 720)
|
||||
@@ -87,22 +97,25 @@ class MainWindow(QMainWindow):
|
||||
self._build_central()
|
||||
self._build_statusbar()
|
||||
self._build_menu()
|
||||
self._refresh_collections()
|
||||
self.statusBar().showMessage("Откройте папку с картинками")
|
||||
self.statusBar().showMessage("Создайте или откройте проект (Файл)")
|
||||
|
||||
# ------------------------------------------------------------------ setup
|
||||
def _build_menu(self) -> None:
|
||||
file_menu = self.menuBar().addMenu("Файл")
|
||||
file_menu.addAction("Открыть папку…", self._choose_folder)
|
||||
file_menu.addAction("Создать проект…", self._create_project)
|
||||
file_menu.addAction("Открыть проект…", self._open_project_dialog)
|
||||
file_menu.addAction("Импортировать папку как проект…", self._import_folder_as_project)
|
||||
file_menu.addAction("Создать из ролика…", self._create_from_video)
|
||||
self._recent_menu = file_menu.addMenu("Недавние проекты")
|
||||
self._refresh_recent_menu()
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Рассчитать кадр", self._recompute_current).setShortcut("Space")
|
||||
file_menu.addAction("Детектировать все", self._detect_all)
|
||||
file_menu.addAction("Детектировать все (дозапуск)", lambda: self._detect_all(False))
|
||||
file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True))
|
||||
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.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M")
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Выход", self.close)
|
||||
|
||||
@@ -110,9 +123,10 @@ class MainWindow(QMainWindow):
|
||||
tb = self.addToolBar("Главная")
|
||||
tb.setMovable(False)
|
||||
|
||||
tb.addAction(QAction("Открыть папку…", self, triggered=self._choose_folder))
|
||||
tb.addAction(QAction("Создать проект…", self, triggered=self._create_project))
|
||||
tb.addAction(QAction("Открыть проект…", self, triggered=self._open_project_dialog))
|
||||
from_video = QAction("Создать из ролика…", self, triggered=self._create_from_video)
|
||||
from_video.setToolTip("Разложить видео на кадры в папку-коллекцию и открыть её")
|
||||
from_video.setToolTip("Разложить видео на кадры в новый проект и открыть его")
|
||||
tb.addAction(from_video)
|
||||
tb.addSeparator()
|
||||
|
||||
@@ -130,7 +144,17 @@ class MainWindow(QMainWindow):
|
||||
calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
|
||||
calc.setToolTip("Запустить детектор на выбранном кадре (Space / двойной клик по файлу)")
|
||||
tb.addAction(calc)
|
||||
tb.addAction(QAction("Детектировать все", self, triggered=self._detect_all))
|
||||
detect_all = QAction("Детектировать все", self, triggered=lambda: self._detect_all(False))
|
||||
detect_all.setToolTip("Рассчитать все ещё не посчитанные кадры (дозапуск; кэш сохраняется)")
|
||||
tb.addAction(detect_all)
|
||||
regen = QAction("Все заново", self, triggered=lambda: self._detect_all(True))
|
||||
regen.setToolTip("Очистить кэш детекций и пересчитать всю папку заново")
|
||||
tb.addAction(regen)
|
||||
|
||||
self.stop_action = QAction("■ Стоп", self, triggered=self._request_cancel)
|
||||
self.stop_action.setToolTip("Отменить текущую операцию (Esc)")
|
||||
self.stop_action.setEnabled(False)
|
||||
tb.addAction(self.stop_action)
|
||||
|
||||
tb.addSeparator()
|
||||
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
|
||||
@@ -158,27 +182,15 @@ class MainWindow(QMainWindow):
|
||||
self.file_list.currentItemChanged.connect(self._on_file_selected)
|
||||
self.file_list.itemDoubleClicked.connect(self._on_file_activated)
|
||||
|
||||
# Collection controls live next to the file list — they act on its selection.
|
||||
self.collection_combo = QComboBox()
|
||||
self.collection_combo.setToolTip("Активная коллекция, куда перемещаются кадры")
|
||||
self.collection_combo.activated.connect(self._on_collection_selected)
|
||||
new_coll = QPushButton("Создать")
|
||||
new_coll.clicked.connect(self._create_collection)
|
||||
move_btn = QPushButton("В коллекцию →")
|
||||
move_btn.setToolTip("Переместить выбранные кадры в активную коллекцию (Ctrl+M)")
|
||||
move_btn.clicked.connect(self._move_to_collection)
|
||||
|
||||
coll_row = QHBoxLayout()
|
||||
coll_row.setContentsMargins(0, 0, 0, 0)
|
||||
coll_row.addWidget(QLabel("Коллекция:"))
|
||||
coll_row.addWidget(self.collection_combo, 1)
|
||||
coll_row.addWidget(new_coll)
|
||||
# One default collection ("Избранное"); the button acts on the list selection.
|
||||
move_btn = QPushButton("★ В избранное")
|
||||
move_btn.setToolTip("Переместить выбранные кадры в избранное проекта (Ctrl+M)")
|
||||
move_btn.clicked.connect(self._move_to_favorites)
|
||||
|
||||
left = QWidget()
|
||||
left_layout = QVBoxLayout(left)
|
||||
left_layout.setContentsMargins(4, 4, 4, 4)
|
||||
left_layout.setSpacing(4)
|
||||
left_layout.addLayout(coll_row)
|
||||
left_layout.addWidget(self.file_list, 1)
|
||||
left_layout.addWidget(move_btn)
|
||||
|
||||
@@ -253,6 +265,7 @@ class MainWindow(QMainWindow):
|
||||
QShortcut(QKeySequence("."), self, lambda: self._step(1))
|
||||
QShortcut(QKeySequence("["), self, lambda: self._step_hit(-1))
|
||||
QShortcut(QKeySequence("]"), self, lambda: self._step_hit(1))
|
||||
QShortcut(QKeySequence(Qt.Key_Escape), self, self._request_cancel)
|
||||
return bar
|
||||
|
||||
# -------------------------------------------------------------- navigation
|
||||
@@ -302,6 +315,35 @@ class MainWindow(QMainWindow):
|
||||
self.progress.setVisible(False)
|
||||
self.statusBar().addPermanentWidget(self.progress)
|
||||
|
||||
# ------------------------------------------------------------- cancellation
|
||||
def _begin_busy(self, total: int | None = None) -> None:
|
||||
"""Enter a cancellable long operation. ``total=None`` => busy spinner."""
|
||||
self._busy = True
|
||||
self._cancel = False
|
||||
self.stop_action.setEnabled(True)
|
||||
if total is None:
|
||||
self.progress.setRange(0, 0) # indeterminate
|
||||
else:
|
||||
self.progress.setRange(0, total)
|
||||
self.progress.setValue(0)
|
||||
self.progress.setVisible(True)
|
||||
|
||||
def _end_busy(self) -> None:
|
||||
self._busy = False
|
||||
self.stop_action.setEnabled(False)
|
||||
self.progress.setVisible(False)
|
||||
self.progress.setRange(0, 100) # leave it determinate for the next user
|
||||
|
||||
def _request_cancel(self) -> None:
|
||||
if self._busy:
|
||||
self._cancel = True
|
||||
self.statusBar().showMessage("Отмена…")
|
||||
|
||||
def _poll_cancel(self) -> bool:
|
||||
"""Cancel hook for core engines: pump the UI so Стоп registers, then report."""
|
||||
QApplication.processEvents()
|
||||
return self._cancel
|
||||
|
||||
# --------------------------------------------------------------- detector
|
||||
def _make_detector(self):
|
||||
d = self._cfg.detection
|
||||
@@ -316,7 +358,7 @@ class MainWindow(QMainWindow):
|
||||
# YOLO/combined need a model — offer to pick one if missing.
|
||||
if name in ("yolo", "combined") and not self._cfg.model_path:
|
||||
self._choose_model()
|
||||
settings_store.save(self._cfg)
|
||||
self._persist_settings()
|
||||
self._invalidate_results()
|
||||
|
||||
def _choose_model(self) -> None:
|
||||
@@ -324,34 +366,210 @@ class MainWindow(QMainWindow):
|
||||
path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)")
|
||||
if path:
|
||||
self._cfg.model_path = path
|
||||
settings_store.save(self._cfg)
|
||||
self._persist_settings()
|
||||
self.statusBar().showMessage(f"Модель: {path}")
|
||||
self._invalidate_results()
|
||||
|
||||
def _invalidate_results(self) -> None:
|
||||
"""Detector changed — drop the cache and refresh the current image."""
|
||||
"""Detector changed — drop the in-memory cache and refresh the current image.
|
||||
|
||||
The on-disk cache is left as-is; it won't be reloaded for the new detector
|
||||
(key mismatch) and gets overwritten once results for the new detector exist.
|
||||
"""
|
||||
self._detector_key = None
|
||||
self._results.clear()
|
||||
for i in range(self.file_list.count()):
|
||||
item = self.file_list.item(i)
|
||||
item.setText(item.data(Qt.UserRole + 1))
|
||||
item.setBackground(QBrush())
|
||||
self._refresh_marks()
|
||||
self._clear_results()
|
||||
if self._current is not None:
|
||||
self._show(self._current)
|
||||
|
||||
# --------------------------------------------------------------- handlers
|
||||
def open_path(self, folder: str) -> None:
|
||||
self._load_folder(Path(folder))
|
||||
# ----------------------------------------------------------------- projects
|
||||
def open_path(self, path: str) -> None:
|
||||
"""Open a project at ``path`` (a project folder or its project.json)."""
|
||||
p = Path(path)
|
||||
if Project.is_project(p):
|
||||
try:
|
||||
self._open_project(Project.load(p))
|
||||
except (OSError, ValueError) as exc: # noqa: BLE001 - surface to the user
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}")
|
||||
elif p.is_dir():
|
||||
QMessageBox.information(
|
||||
self, "Не проект",
|
||||
"Это обычная папка, а не проект. Используйте "
|
||||
"«Импортировать папку как проект…».",
|
||||
)
|
||||
else:
|
||||
QMessageBox.warning(self, "Ошибка", f"Путь не найден: {p}")
|
||||
|
||||
def _choose_folder(self) -> None:
|
||||
def _auto_open_last(self) -> None:
|
||||
"""On startup, reopen the last project if it still exists (best-effort)."""
|
||||
last = settings_store.last_project()
|
||||
if last and Project.is_project(last):
|
||||
try:
|
||||
self._open_project(Project.load(last))
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
|
||||
def _new_project_root(self, default_name: str = "") -> Path | None:
|
||||
"""Prompt for a parent dir + name; return a fresh (empty) project root or None."""
|
||||
start = settings_store.last_dir() or str(Path.home())
|
||||
parent = QFileDialog.getExistingDirectory(self, "Где создать проект", start)
|
||||
if not parent:
|
||||
return None
|
||||
name, ok = QInputDialog.getText(self, "Новый проект", "Имя проекта:", text=default_name)
|
||||
name = name.strip()
|
||||
if not ok or not name:
|
||||
return None
|
||||
root = Path(parent) / name
|
||||
if root.exists() and any(root.iterdir()):
|
||||
QMessageBox.warning(self, "Папка занята", f"Папка уже существует и не пуста:\n{root}")
|
||||
return None
|
||||
settings_store.set_last_dir(parent)
|
||||
return root
|
||||
|
||||
def _create_project(self) -> None:
|
||||
if self._busy:
|
||||
return
|
||||
root = self._new_project_root()
|
||||
if root is None:
|
||||
return
|
||||
try:
|
||||
project = Project.create(root, name=root.name)
|
||||
except OSError as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось создать проект:\n{exc}")
|
||||
return
|
||||
project.update_from_config(self._cfg) # seed from current global defaults
|
||||
project.save()
|
||||
self._open_project(project)
|
||||
|
||||
def _open_project_dialog(self) -> None:
|
||||
if self._busy:
|
||||
return
|
||||
start = settings_store.last_dir() or str(Path.home())
|
||||
folder = QFileDialog.getExistingDirectory(self, "Открыть проект (папка проекта)", start)
|
||||
if not folder:
|
||||
return
|
||||
if not Project.is_project(folder):
|
||||
QMessageBox.warning(self, "Не проект", f"В папке нет {PROJECT_FILE}:\n{folder}")
|
||||
return
|
||||
try:
|
||||
project = Project.load(folder)
|
||||
except (OSError, ValueError) as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}")
|
||||
return
|
||||
settings_store.set_last_dir(str(Path(folder).parent))
|
||||
self._open_project(project)
|
||||
|
||||
def _import_folder_as_project(self) -> None:
|
||||
"""Create a project and copy a folder of images into its frames/."""
|
||||
if self._busy:
|
||||
return
|
||||
start = settings_store.last_dir() or ""
|
||||
folder = QFileDialog.getExistingDirectory(self, "Открыть папку с картинками", start)
|
||||
if folder:
|
||||
self._load_folder(Path(folder))
|
||||
src = QFileDialog.getExistingDirectory(self, "Папка с картинками для импорта", start)
|
||||
if not src:
|
||||
return
|
||||
src = Path(src)
|
||||
images = sorted(p for p in src.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
|
||||
if not images:
|
||||
QMessageBox.warning(self, "Пусто", f"В папке нет картинок:\n{src}")
|
||||
return
|
||||
root = self._new_project_root(default_name=src.name)
|
||||
if root is None:
|
||||
return
|
||||
try:
|
||||
project = Project.create(root, name=root.name, source=str(src))
|
||||
except OSError as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось создать проект:\n{exc}")
|
||||
return
|
||||
project.update_from_config(self._cfg)
|
||||
project.save()
|
||||
|
||||
self._begin_busy(len(images))
|
||||
copied = 0
|
||||
try:
|
||||
for i, p in enumerate(images, 1):
|
||||
self.progress.setValue(i)
|
||||
self.statusBar().showMessage(f"Импорт {i}/{len(images)}: {p.name}")
|
||||
QApplication.processEvents()
|
||||
if self._cancel:
|
||||
break
|
||||
dst = self._unique_dest(project.frames_dir, p.name)
|
||||
try:
|
||||
shutil.copy2(str(p), str(dst))
|
||||
copied += 1
|
||||
except OSError:
|
||||
continue
|
||||
finally:
|
||||
self._end_busy()
|
||||
|
||||
# Carry over an old sidecar detection cache (basename-keyed) if present.
|
||||
old_sidecar = src / ".hvideotool_detections.json"
|
||||
if old_sidecar.is_file():
|
||||
try:
|
||||
shutil.copy2(str(old_sidecar), str(project.cache_path))
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
self.statusBar().showMessage(f"Импортировано {copied} картинок → {project.name}")
|
||||
self._open_project(project)
|
||||
|
||||
def _refresh_recent_menu(self) -> None:
|
||||
self._recent_menu.clear()
|
||||
recents = settings_store.recent_projects()
|
||||
if not recents:
|
||||
empty = self._recent_menu.addAction("(пусто)")
|
||||
empty.setEnabled(False)
|
||||
return
|
||||
for path in recents:
|
||||
self._recent_menu.addAction(Path(path).name, lambda checked=False, p=path: self._open_recent(p))
|
||||
|
||||
def _open_recent(self, path: str) -> None:
|
||||
if self._busy:
|
||||
return
|
||||
if not Project.is_project(path):
|
||||
QMessageBox.warning(self, "Нет проекта", f"Проект не найден:\n{path}")
|
||||
return
|
||||
try:
|
||||
self._open_project(Project.load(path))
|
||||
except (OSError, ValueError) as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}")
|
||||
|
||||
def _open_project(self, project: Project) -> None:
|
||||
"""Core open: set project state, apply its settings, list its frames."""
|
||||
if self._busy:
|
||||
return
|
||||
self._project = project
|
||||
project.frames_dir.mkdir(parents=True, exist_ok=True)
|
||||
project.apply_to_config(self._cfg) # per-project settings -> live config
|
||||
self._sync_settings_ui()
|
||||
self._detector_key = None
|
||||
self._restorer_key = None
|
||||
self._restored.clear()
|
||||
settings_store.set_last_project(str(project.root))
|
||||
settings_store.add_recent_project(str(project.root))
|
||||
self._refresh_recent_menu()
|
||||
self.setWindowTitle(f"HVideoTool — {project.name}")
|
||||
self._load_folder(project.frames_dir)
|
||||
|
||||
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)
|
||||
self.view.set_threshold(self._cfg.default_threshold)
|
||||
|
||||
def _persist_settings(self) -> None:
|
||||
"""Save settings to the global defaults and (if open) into the project."""
|
||||
settings_store.save(self._cfg) # global defaults for new projects
|
||||
if self._project is not None:
|
||||
self._project.update_from_config(self._cfg)
|
||||
self._project.save()
|
||||
|
||||
def _create_from_video(self) -> None:
|
||||
"""Decode a video into a folder of frames (a collection) and open it."""
|
||||
"""Decode a video into a new project's frames/ and open the project."""
|
||||
if self._busy:
|
||||
return
|
||||
path, _ = QFileDialog.getOpenFileName(
|
||||
self, "Выберите ролик", settings_store.last_dir() or "", _VIDEO_FILTER
|
||||
)
|
||||
@@ -362,18 +580,26 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
keyframes_only, step, max_dim = dialog.options()
|
||||
video = Path(path)
|
||||
out = video.parent / f"{video.stem}_frames"
|
||||
root = video.parent / f"{video.stem}_frames"
|
||||
if Project.is_project(root):
|
||||
project = Project.load(root) # re-extract into the existing project
|
||||
elif root.exists() and any(root.iterdir()):
|
||||
QMessageBox.warning(self, "Папка занята", f"Папка уже существует и не пуста:\n{root}")
|
||||
return
|
||||
else:
|
||||
project = Project.create(root, name=root.name, source=str(video))
|
||||
project.update_from_config(self._cfg)
|
||||
project.save()
|
||||
out = project.frames_dir
|
||||
|
||||
self.progress.setRange(0, 1000) # promille of duration
|
||||
self.progress.setValue(0)
|
||||
self.progress.setVisible(True)
|
||||
self._begin_busy(1000) # promille of duration
|
||||
|
||||
def cb(done: float, total: float) -> bool:
|
||||
if total > 0:
|
||||
self.progress.setValue(int(1000 * min(done, total) / total))
|
||||
self.statusBar().showMessage(f"Извлечение кадров: {done:.0f}/{total:.0f} с…")
|
||||
QApplication.processEvents()
|
||||
return True
|
||||
return not self._cancel # returning False stops extraction
|
||||
|
||||
try:
|
||||
saved = extract_frames(
|
||||
@@ -381,19 +607,23 @@ class MainWindow(QMainWindow):
|
||||
max_dim=max_dim, progress=cb,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - surface decode errors to the user
|
||||
self.progress.setVisible(False)
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
|
||||
return
|
||||
finally:
|
||||
self.progress.setVisible(False)
|
||||
cancelled = self._cancel
|
||||
self._end_busy()
|
||||
|
||||
if saved == 0:
|
||||
QMessageBox.warning(self, "Пусто", "Из ролика не удалось извлечь ни одного кадра.")
|
||||
msg = "Извлечение отменено — кадров нет." if cancelled \
|
||||
else "Из ролика не удалось извлечь ни одного кадра."
|
||||
QMessageBox.warning(self, "Пусто", msg)
|
||||
return
|
||||
self.statusBar().showMessage(f"Извлечено {saved} кадров → {out}")
|
||||
self._load_folder(out)
|
||||
verb = "Отменено, извлечено" if cancelled else "Извлечено"
|
||||
self.statusBar().showMessage(f"{verb} {saved} кадров → {out}")
|
||||
self._open_project(project)
|
||||
|
||||
def _load_folder(self, folder: Path) -> None:
|
||||
"""List images from ``folder`` (a project's frames/) into the file list."""
|
||||
if not folder.is_dir():
|
||||
QMessageBox.warning(self, "Ошибка", f"Папка не найдена: {folder}")
|
||||
return
|
||||
@@ -404,13 +634,11 @@ class MainWindow(QMainWindow):
|
||||
self._files = files
|
||||
self._results.clear()
|
||||
self._current = None
|
||||
settings_store.set_last_dir(str(folder))
|
||||
|
||||
self.file_list.blockSignals(True)
|
||||
self.file_list.setUpdatesEnabled(False)
|
||||
self.file_list.clear()
|
||||
self.progress.setRange(0, len(files))
|
||||
self.progress.setVisible(True)
|
||||
self._begin_busy(len(files))
|
||||
for i, p in enumerate(files, 1):
|
||||
item = QListWidgetItem(p.name)
|
||||
item.setData(Qt.UserRole, str(p))
|
||||
@@ -420,18 +648,27 @@ class MainWindow(QMainWindow):
|
||||
self.progress.setValue(i)
|
||||
self.statusBar().showMessage(f"Загрузка списка: {i}/{len(files)}…")
|
||||
QApplication.processEvents()
|
||||
if self._cancel:
|
||||
self._files = files[:i] # keep only what we listed
|
||||
break
|
||||
self.file_list.setUpdatesEnabled(True)
|
||||
self.file_list.blockSignals(False)
|
||||
self.progress.setVisible(False)
|
||||
self._end_busy()
|
||||
loaded = self._load_cached_results() # reuse a matching on-disk cache
|
||||
self._refresh_marks()
|
||||
|
||||
self._refresh_collections()
|
||||
if not files:
|
||||
self.view.set_image(None, [])
|
||||
self._update_nav()
|
||||
self.statusBar().showMessage(f"В папке нет картинок: {folder}")
|
||||
self.statusBar().showMessage(
|
||||
"В проекте пока нет кадров — импортируйте папку или создайте из ролика"
|
||||
)
|
||||
return
|
||||
self.statusBar().showMessage(f"{len(files)} картинок · {folder} · двойной клик / «Рассчитать кадр» для детекции")
|
||||
cache_note = f" · загружен кэш детекций ({loaded})" if loaded else ""
|
||||
self.statusBar().showMessage(
|
||||
f"{len(files)} картинок · {folder} · двойной клик / «Рассчитать кадр» для детекции"
|
||||
+ cache_note
|
||||
)
|
||||
self.file_list.setCurrentRow(0)
|
||||
|
||||
def _on_file_selected(self, current: QListWidgetItem | None, _prev=None) -> None:
|
||||
@@ -441,10 +678,14 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def _on_file_activated(self, item: QListWidgetItem) -> None:
|
||||
# Double-click: compute if not already cached, then show.
|
||||
path = Path(item.data(Qt.UserRole))
|
||||
if str(path) not in self._results and self._detect(path) is None:
|
||||
if self._busy:
|
||||
return
|
||||
self._refresh_marks()
|
||||
path = Path(item.data(Qt.UserRole))
|
||||
if str(path) not in self._results:
|
||||
if self._detect(path) is None:
|
||||
return
|
||||
self._refresh_marks()
|
||||
self._save_results() # only when a detection actually ran
|
||||
self._show(path)
|
||||
|
||||
def _detect(self, path: Path) -> list[Detection] | None:
|
||||
@@ -481,40 +722,56 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def _recompute_current(self) -> None:
|
||||
"""Toolbar/Space: (re)run the detector on the selected frame."""
|
||||
if self._current is None:
|
||||
if self._current is None or self._busy:
|
||||
return
|
||||
self._results.pop(str(self._current), None)
|
||||
self._detector_key = None # rebuild the detector so settings changes take effect
|
||||
if self._detect(self._current) is None:
|
||||
return
|
||||
self._refresh_marks()
|
||||
self._save_results()
|
||||
self._show(self._current)
|
||||
|
||||
def _detect_all(self) -> None:
|
||||
if not self._files:
|
||||
def _detect_all(self, force: bool = False) -> None:
|
||||
"""Detect the whole folder. ``force`` clears the cache first (full regen);
|
||||
otherwise already-computed frames are skipped, so it resumes/tops-up."""
|
||||
if not self._files or self._busy:
|
||||
return
|
||||
if force:
|
||||
self._clear_results()
|
||||
total = len(self._files)
|
||||
self.progress.setRange(0, total)
|
||||
self.progress.setVisible(True)
|
||||
self._begin_busy(total)
|
||||
done = 0
|
||||
try:
|
||||
for i, p in enumerate(self._files, 1):
|
||||
self.progress.setValue(i)
|
||||
self.statusBar().showMessage(f"Детекция {i}/{total}: {p.name}")
|
||||
QApplication.processEvents()
|
||||
if self._cancel:
|
||||
break
|
||||
if self._detect(p) is None:
|
||||
return # detector unavailable — message already shown
|
||||
done = i
|
||||
if i % 50 == 0:
|
||||
self._refresh_marks() # let marks appear progressively
|
||||
finally:
|
||||
self.progress.setVisible(False)
|
||||
self._end_busy()
|
||||
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
||||
self._refresh_marks()
|
||||
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
||||
self._save_results() # persist progress (works for completed and cancelled runs)
|
||||
if self._cancel:
|
||||
self.statusBar().showMessage(
|
||||
f"Отменено на {done}/{total} · детекции на {hits} картинках"
|
||||
)
|
||||
else:
|
||||
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
||||
if self._current is not None:
|
||||
self._show(self._current)
|
||||
|
||||
# ------------------------------------------------------------- restoration
|
||||
def _restore_current(self) -> None:
|
||||
"""Run the restorer on the current frame's detected regions and show it."""
|
||||
if self._current is None:
|
||||
if self._current is None or self._busy:
|
||||
return
|
||||
key = str(self._current)
|
||||
if key not in self._results and self._detect(self._current) is None:
|
||||
@@ -527,14 +784,20 @@ class MainWindow(QMainWindow):
|
||||
img = imread_unicode(key)
|
||||
if img is None:
|
||||
return
|
||||
self._begin_busy() # indeterminate — engine drives the duration
|
||||
self.statusBar().showMessage(f"Восстановление: {self._current.name}…")
|
||||
QApplication.processEvents()
|
||||
try:
|
||||
restorer = self._make_restorer()
|
||||
restored = restorer.restore(img, dets)
|
||||
restored = restorer.restore(img, dets, should_cancel=self._poll_cancel)
|
||||
except Cancelled:
|
||||
self.statusBar().showMessage("Восстановление отменено")
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001 - surface model/engine errors
|
||||
QMessageBox.warning(self, "Ошибка восстановления", str(exc))
|
||||
return
|
||||
finally:
|
||||
self._end_busy()
|
||||
self._restored[key] = restored
|
||||
self._showing_restored = True
|
||||
self.view.set_image(restored, [])
|
||||
@@ -556,7 +819,7 @@ class MainWindow(QMainWindow):
|
||||
if dlg.exec() != QDialog.Accepted:
|
||||
return
|
||||
dlg.apply_to_config()
|
||||
settings_store.save(self._cfg)
|
||||
self._persist_settings()
|
||||
self._restorer_key = None # rebuild on next restore
|
||||
self.statusBar().showMessage(f"Движок восстановления: {self._cfg.restorer}")
|
||||
|
||||
@@ -582,100 +845,37 @@ class MainWindow(QMainWindow):
|
||||
def _save_restored(self) -> None:
|
||||
if self._current is None or str(self._current) not in self._restored:
|
||||
return
|
||||
dest_dir = self._collection or self._current.parent
|
||||
out = self._unique_dest(dest_dir, f"{self._current.stem}_restored.jpg")
|
||||
out = self._unique_dest(self._current.parent, f"{self._current.stem}_restored.jpg")
|
||||
if imwrite_unicode(str(out), self._restored[str(self._current)]):
|
||||
self.statusBar().showMessage(f"Сохранено: {out}")
|
||||
else:
|
||||
QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.")
|
||||
|
||||
# ------------------------------------------------------------ collections
|
||||
def _collections_base(self) -> Path:
|
||||
"""Where new collections are created: next to the opened folder, else home."""
|
||||
if self._folder is not None:
|
||||
return self._folder.parent
|
||||
return Path.home() / "HVideoTool" / "collections"
|
||||
|
||||
def _refresh_collections(self) -> None:
|
||||
"""Repopulate the collection combo from sibling folders of the opened folder.
|
||||
|
||||
Keeps the active collection selected (and present even if it lives elsewhere).
|
||||
"""
|
||||
base = self._collections_base()
|
||||
subdirs = []
|
||||
if base.exists():
|
||||
subdirs = sorted(
|
||||
(p for p in base.iterdir() if p.is_dir() and p != self._folder),
|
||||
key=lambda p: p.name.lower(),
|
||||
)
|
||||
self.collection_combo.blockSignals(True)
|
||||
self.collection_combo.clear()
|
||||
self.collection_combo.addItem("— не выбрана —", None)
|
||||
for p in subdirs:
|
||||
self.collection_combo.addItem(p.name, str(p))
|
||||
# Make sure the active collection is listed even if it's outside base.
|
||||
if self._collection is not None and self.collection_combo.findData(str(self._collection)) < 0:
|
||||
self.collection_combo.addItem(self._collection.name, str(self._collection))
|
||||
self.collection_combo.addItem("Выбрать папку…", "__browse__")
|
||||
self._select_active_in_combo()
|
||||
self.collection_combo.blockSignals(False)
|
||||
|
||||
def _select_active_in_combo(self) -> None:
|
||||
idx = self.collection_combo.findData(str(self._collection)) if self._collection else 0
|
||||
self.collection_combo.setCurrentIndex(max(0, idx))
|
||||
|
||||
def _on_collection_selected(self, _index: int) -> None:
|
||||
data = self.collection_combo.currentData()
|
||||
if data == "__browse__":
|
||||
self._browse_collection()
|
||||
# -------------------------------------------------------------- favorites
|
||||
def _move_to_favorites(self) -> None:
|
||||
"""Move the selected frames into the project's single default collection."""
|
||||
if self._busy:
|
||||
return
|
||||
self._collection = Path(data) if data else None
|
||||
if self._collection is not None:
|
||||
self.statusBar().showMessage(f"Активная коллекция: {self._collection}")
|
||||
|
||||
def _browse_collection(self) -> None:
|
||||
start = str(self._collections_base())
|
||||
folder = QFileDialog.getExistingDirectory(self, "Выбрать коллекцию", start)
|
||||
if folder:
|
||||
self._collection = Path(folder)
|
||||
self._refresh_collections()
|
||||
self.statusBar().showMessage(f"Активная коллекция: {folder}")
|
||||
else:
|
||||
self._select_active_in_combo() # revert the combo to the current collection
|
||||
|
||||
def _create_collection(self) -> None:
|
||||
name, ok = QInputDialog.getText(self, "Создать коллекцию", "Имя коллекции:")
|
||||
name = name.strip()
|
||||
if not ok or not name:
|
||||
return
|
||||
path = self._collections_base() / name
|
||||
try:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось создать коллекцию:\n{exc}")
|
||||
return
|
||||
self._collection = path
|
||||
self._refresh_collections()
|
||||
self.statusBar().showMessage(f"Активная коллекция: {path}")
|
||||
|
||||
def _move_to_collection(self) -> None:
|
||||
if self._collection is None:
|
||||
QMessageBox.information(
|
||||
self, "Нет коллекции",
|
||||
"Сначала выберите коллекцию в списке или создайте новую («Создать…»).",
|
||||
)
|
||||
if self._project is None:
|
||||
QMessageBox.information(self, "Нет проекта", "Сначала откройте или создайте проект.")
|
||||
return
|
||||
items = self.file_list.selectedItems()
|
||||
if not items:
|
||||
QMessageBox.information(self, "Нет выбора", "Выберите кадры в списке слева.")
|
||||
return
|
||||
dest = self._project.favorites_dir
|
||||
try:
|
||||
dest.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось создать избранное:\n{exc}")
|
||||
return
|
||||
|
||||
moved = 0
|
||||
for item in items:
|
||||
src = Path(item.data(Qt.UserRole))
|
||||
if not src.exists():
|
||||
continue
|
||||
dst = self._unique_dest(self._collection, src.name)
|
||||
dst = self._unique_dest(dest, src.name)
|
||||
try:
|
||||
shutil.move(str(src), str(dst))
|
||||
except OSError as exc:
|
||||
@@ -688,7 +888,9 @@ class MainWindow(QMainWindow):
|
||||
if self._current == src:
|
||||
self._current = None
|
||||
|
||||
self.statusBar().showMessage(f"Перемещено {moved} → {self._collection.name}")
|
||||
if moved:
|
||||
self._save_results() # cache file should forget the moved frames
|
||||
self.statusBar().showMessage(f"В избранное перемещено {moved}")
|
||||
cur = self.file_list.currentItem()
|
||||
if cur is not None:
|
||||
self._show(Path(cur.data(Qt.UserRole)))
|
||||
@@ -714,15 +916,60 @@ class MainWindow(QMainWindow):
|
||||
_TINT_HIT = QColor(200, 80, 80, 70)
|
||||
_TINT_CLEAN = QColor(90, 160, 90, 50)
|
||||
|
||||
def _set_row_tag(self, item: QListWidgetItem, count: int) -> None:
|
||||
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)
|
||||
|
||||
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)
|
||||
self._set_row_tag(item, count)
|
||||
return
|
||||
|
||||
# ------------------------------------------------------------- result cache
|
||||
def _results_key(self) -> dict:
|
||||
"""Detector identity used to tag/validate the on-disk detection cache."""
|
||||
d = self._cfg.detection
|
||||
return detection_cache.make_key(
|
||||
self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz
|
||||
)
|
||||
|
||||
def _save_results(self) -> None:
|
||||
"""Persist the detection cache in the project (skip if nothing to save)."""
|
||||
if self._project is None or not self._results:
|
||||
return
|
||||
detection_cache.save_results(
|
||||
self._project.cache_path, self._results_key(), self._results
|
||||
)
|
||||
|
||||
def _load_cached_results(self) -> int:
|
||||
"""Load a matching on-disk cache into `_results` and tag rows. Returns count."""
|
||||
if self._project is None:
|
||||
return 0
|
||||
cached = detection_cache.load_results(
|
||||
self._project.cache_path, self._results_key(), self._project.frames_dir
|
||||
)
|
||||
if not cached:
|
||||
return 0
|
||||
self._results = cached
|
||||
for i in range(self.file_list.count()):
|
||||
item = self.file_list.item(i)
|
||||
path = item.data(Qt.UserRole)
|
||||
if path in self._results: # `in`, not truthy: empty list = checked-clean
|
||||
self._set_row_tag(item, len(self._results[path]))
|
||||
return len(cached)
|
||||
|
||||
def _clear_results(self) -> None:
|
||||
"""Drop all cached detections and reset row labels/tints (keeps the detector)."""
|
||||
self._results.clear()
|
||||
for i in range(self.file_list.count()):
|
||||
item = self.file_list.item(i)
|
||||
item.setText(item.data(Qt.UserRole + 1))
|
||||
item.setBackground(QBrush())
|
||||
self._refresh_marks()
|
||||
|
||||
def _refresh_marks(self) -> None:
|
||||
"""Project frames-with-detections onto the scrubber as marks."""
|
||||
marks = {
|
||||
@@ -767,4 +1014,11 @@ class MainWindow(QMainWindow):
|
||||
def _on_threshold_changed(self, value: float) -> None:
|
||||
self._cfg.default_threshold = value
|
||||
self.view.set_threshold(value)
|
||||
settings_store.save(self._cfg)
|
||||
self._persist_settings()
|
||||
|
||||
def closeEvent(self, event) -> None: # noqa: N802 - Qt override
|
||||
self._save_results() # persist the detection cache on exit
|
||||
if self._project is not None:
|
||||
self._project.update_from_config(self._cfg)
|
||||
self._project.save()
|
||||
super().closeEvent(event)
|
||||
|
||||
Reference in New Issue
Block a user