"""Main window: open a **project** and inspect what the detector found.
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 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 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
import shutil
from pathlib import Path
from PySide6.QtCore import Qt
from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
from PySide6.QtWidgets import (
QAbstractItemView,
QApplication,
QComboBox,
QDialog,
QDoubleSpinBox,
QFileDialog,
QHBoxLayout,
QInputDialog,
QLabel,
QListWidget,
QListWidgetItem,
QMainWindow,
QMessageBox,
QProgressBar,
QPushButton,
QSplitter,
QTableWidget,
QTableWidgetItem,
QVBoxLayout,
QWidget,
)
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
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);;Все файлы (*.*)"
_DETECTORS = ["classic", "yolo", "combined"]
class MainWindow(QMainWindow):
def __init__(self, config: AppConfig) -> None:
super().__init__()
self._cfg = config
self._detector = None
self._detector_key = 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._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)
self._build_toolbar()
self._build_central()
self._build_statusbar()
self._build_menu()
self.statusBar().showMessage("Создайте или откройте проект (Файл)")
# ------------------------------------------------------------------ setup
def _build_menu(self) -> None:
file_menu = self.menuBar().addMenu("Файл")
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("Детектировать все (дозапуск)", 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._move_to_favorites).setShortcut("Ctrl+M")
file_menu.addSeparator()
file_menu.addAction("Выход", self.close)
def _build_toolbar(self) -> None:
tb = self.addToolBar("Главная")
tb.setMovable(False)
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("Разложить видео на кадры в новый проект и открыть его")
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)
self.model_action = QAction("Модель…", self, triggered=self._choose_model)
tb.addAction(self.model_action)
tb.addSeparator()
calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
calc.setToolTip("Запустить детектор на выбранном кадре (Space / двойной клик по файлу)")
tb.addAction(calc)
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)
restore.setToolTip("Восстановить найденные области на текущем кадре")
tb.addAction(restore)
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
self.toggle_restored_action.setEnabled(False)
tb.addAction(self.toggle_restored_action)
self.save_restored_action = QAction("Сохранить результат", self, triggered=self._save_restored)
self.save_restored_action.setEnabled(False)
tb.addAction(self.save_restored_action)
tb.addSeparator()
tb.addWidget(QLabel(" Порог: "))
self.threshold_spin = QDoubleSpinBox()
self.threshold_spin.setRange(0.0, 1.0)
self.threshold_spin.setSingleStep(0.05)
self.threshold_spin.setValue(self._cfg.default_threshold)
self.threshold_spin.valueChanged.connect(self._on_threshold_changed)
tb.addWidget(self.threshold_spin)
def _build_central(self) -> None:
self.file_list = QListWidget()
self.file_list.setSelectionMode(QAbstractItemView.ExtendedSelection) # multi-select for moves
self.file_list.currentItemChanged.connect(self._on_file_selected)
self.file_list.itemDoubleClicked.connect(self._on_file_activated)
# 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.addWidget(self.file_list, 1)
left_layout.addWidget(move_btn)
self.view = ImageView(self._cfg.overlay)
self.view.set_threshold(self._cfg.default_threshold)
center = QWidget()
clayout = QVBoxLayout(center)
clayout.setContentsMargins(0, 0, 0, 0)
clayout.setSpacing(2)
clayout.addWidget(self.view, 1)
clayout.addWidget(self._build_nav_bar())
right = QWidget()
rlayout = QVBoxLayout(right)
rlayout.setContentsMargins(4, 4, 4, 4)
self.detail_header = QLabel("Детекции")
self.detail_header.setWordWrap(True)
rlayout.addWidget(self.detail_header)
self.detail_table = QTableWidget(0, 4)
self.detail_table.setHorizontalHeaderLabels(["Тип", "Увер.", "BBox (x,y,w,h)", "Полигон"])
self.detail_table.verticalHeader().setVisible(False)
self.detail_table.setSelectionBehavior(QTableWidget.SelectRows)
self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers)
self.detail_table.itemSelectionChanged.connect(self._on_detail_selected)
rlayout.addWidget(self.detail_table)
splitter = QSplitter(Qt.Horizontal)
splitter.addWidget(left)
splitter.addWidget(center)
splitter.addWidget(right)
splitter.setStretchFactor(0, 0)
splitter.setStretchFactor(1, 1)
splitter.setStretchFactor(2, 0)
splitter.setSizes([240, 640, 300])
self.setCentralWidget(splitter)
def _build_nav_bar(self) -> QWidget:
bar = QWidget()
h = QHBoxLayout(bar)
h.setContentsMargins(4, 2, 4, 2)
self.prev_btn = QPushButton("◀")
self.prev_btn.setToolTip("Предыдущий кадр (,)")
self.prev_btn.clicked.connect(lambda: self._step(-1))
self.next_btn = QPushButton("▶")
self.next_btn.setToolTip("Следующий кадр (.)")
self.next_btn.clicked.connect(lambda: self._step(1))
self.frame_slider = MarkerSlider(Qt.Horizontal)
self.frame_slider.setMinimum(0)
self.frame_slider.setMaximum(0)
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
self.frame_slider.valueChanged.connect(self._on_slider)
self.pos_label = QLabel("0 / 0")
self.pos_label.setMinimumWidth(90)
self.pos_label.setAlignment(Qt.AlignCenter)
self.prev_hit_btn = QPushButton("◀ детекция")
self.prev_hit_btn.setToolTip("Предыдущий кадр с детекцией ([)")
self.prev_hit_btn.clicked.connect(lambda: self._step_hit(-1))
self.next_hit_btn = QPushButton("детекция ▶")
self.next_hit_btn.setToolTip("Следующий кадр с детекцией (])")
self.next_hit_btn.clicked.connect(lambda: self._step_hit(1))
for wdg in (self.prev_btn, self.next_btn, self.frame_slider, self.pos_label,
self.prev_hit_btn, self.next_hit_btn):
h.addWidget(wdg, 1 if wdg is self.frame_slider else 0)
# Keyboard shortcuts (window-wide), chosen to not clash with list/slider arrows.
QShortcut(QKeySequence(","), self, lambda: self._step(-1))
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
def _step(self, delta: int) -> None:
n = self.file_list.count()
if n == 0:
return
row = max(0, min(n - 1, self.file_list.currentRow() + delta))
self.file_list.setCurrentRow(row)
def _step_hit(self, direction: int) -> None:
"""Jump to the nearest frame (in `direction`) that has detections."""
n = self.file_list.count()
if n == 0:
return
row = self.file_list.currentRow()
i = row + direction
while 0 <= i < n:
path = self.file_list.item(i).data(Qt.UserRole)
if self._results.get(path):
self.file_list.setCurrentRow(i)
return
i += direction
self.statusBar().showMessage(
"Больше нет кадров с детекцией в эту сторону "
"(сначала «Детектировать все»)"
)
def _on_slider(self, value: int) -> None:
if self._nav_sync:
return
if value != self.file_list.currentRow():
self.file_list.setCurrentRow(value)
def _update_nav(self) -> None:
n = self.file_list.count()
row = self.file_list.currentRow()
self._nav_sync = True
self.frame_slider.setMaximum(max(0, n - 1))
self.frame_slider.setValue(max(0, row))
self._nav_sync = False
self.pos_label.setText(f"{row + 1 if row >= 0 else 0} / {n}")
def _build_statusbar(self) -> None:
self.progress = QProgressBar()
self.progress.setMaximumWidth(260)
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
key = (self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz)
if key != self._detector_key:
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
self._detector_key = key
return self._detector
def _on_detector_changed(self, name: str) -> None:
self._cfg.detector = name
# 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()
self._persist_settings()
self._invalidate_results()
def _choose_model(self) -> None:
start = self._cfg.model_path or str(Path.cwd() / "models")
path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)")
if path:
self._cfg.model_path = path
self._persist_settings()
self.statusBar().showMessage(f"Модель: {path}")
self._invalidate_results()
def _invalidate_results(self) -> None:
"""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._clear_results()
if self._current is not None:
self._show(self._current)
# ----------------------------------------------------------------- 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 _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 ""
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 new project's frames/ and open the project."""
if self._busy:
return
path, _ = QFileDialog.getOpenFileName(
self, "Выберите ролик", settings_store.last_dir() or "", _VIDEO_FILTER
)
if not path:
return
dialog = ExtractDialog(self)
if dialog.exec() != QDialog.Accepted:
return
keyframes_only, step, max_dim = dialog.options()
video = Path(path)
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._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 not self._cancel # returning False stops extraction
try:
saved = extract_frames(
str(video), str(out), step=step, keyframes_only=keyframes_only,
max_dim=max_dim, progress=cb,
)
except Exception as exc: # noqa: BLE001 - surface decode errors to the user
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
return
finally:
cancelled = self._cancel
self._end_busy()
if saved == 0:
msg = "Извлечение отменено — кадров нет." if cancelled \
else "Из ролика не удалось извлечь ни одного кадра."
QMessageBox.warning(self, "Пусто", msg)
return
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
self.statusBar().showMessage(f"Сканирую папку: {folder}…")
QApplication.processEvents()
files = sorted(p for p in folder.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
self._folder = folder
self._files = files
self._results.clear()
self._current = None
self.file_list.blockSignals(True)
self.file_list.setUpdatesEnabled(False)
self.file_list.clear()
self._begin_busy(len(files))
for i, p in enumerate(files, 1):
item = QListWidgetItem(p.name)
item.setData(Qt.UserRole, str(p))
item.setData(Qt.UserRole + 1, p.name) # base label, without the count suffix
self.file_list.addItem(item)
if i % 1000 == 0:
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._end_busy()
loaded = self._load_cached_results() # reuse a matching on-disk cache
self._refresh_marks()
if not files:
self.view.set_image(None, [])
self._update_nav()
self.statusBar().showMessage(
"В проекте пока нет кадров — импортируйте папку или создайте из ролика"
)
return
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:
self._update_nav()
if current is not None:
self._show(Path(current.data(Qt.UserRole))) # view only — no detection
def _on_file_activated(self, item: QListWidgetItem) -> None:
# Double-click: compute if not already cached, then show.
if self._busy:
return
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:
"""Run (or fetch cached) detections for one image. None on failure."""
key = str(path)
if key in self._results:
return self._results[key]
img = imread_unicode(key)
if img is None:
self.statusBar().showMessage(f"Не удалось прочитать: {path.name}")
return None
try:
detector = self._make_detector()
except Exception as exc: # noqa: BLE001 - surface config/model errors to the user
QMessageBox.warning(self, "Детектор недоступен", str(exc))
return None
self.statusBar().showMessage(f"Детекция: {path.name}…")
QApplication.processEvents()
dets = detector.detect(Frame(image=img, index=0, pts=0.0))
dets.sort(key=lambda d: d.score, reverse=True)
self._results[key] = dets
self._tag_file(path, len(dets))
return dets
def _show(self, path: Path) -> None:
"""Display the image with its cached detections (does not run the detector)."""
self._current = path
self._showing_restored = False
img = imread_unicode(str(path))
dets = self._results.get(str(path)) # None => not yet computed
self.view.set_image(img, dets or [])
self._fill_detail_table(path, img, dets)
self._update_restore_actions()
def _recompute_current(self) -> None:
"""Toolbar/Space: (re)run the detector on the selected frame."""
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, 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._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._end_busy()
hits = sum(1 for p in self._files if self._results.get(str(p)))
self._refresh_marks()
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 or self._busy:
return
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("Нет найденных областей — нечего расцензуривать")
return
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, 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, [])
self._update_restore_actions()
self.statusBar().showMessage(
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()
self._persist_settings()
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
self._showing_restored = not self._showing_restored
key = str(self._current)
if self._showing_restored:
self.view.set_image(self._restored[key], [])
else:
self.view.set_image(imread_unicode(key), self._results.get(key) or [])
self._update_restore_actions()
def _update_restore_actions(self) -> None:
has = self._current is not None and str(self._current) in self._restored
self.toggle_restored_action.setEnabled(has)
self.toggle_restored_action.setText(
"Показать оригинал" if self._showing_restored else "Показать результат"
)
self.save_restored_action.setEnabled(has)
def _save_restored(self) -> None:
if self._current is None or str(self._current) not in self._restored:
return
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, "Ошибка", "Не удалось сохранить файл.")
# -------------------------------------------------------------- favorites
def _move_to_favorites(self) -> None:
"""Move the selected frames into the project's single default collection."""
if self._busy:
return
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(dest, src.name)
try:
shutil.move(str(src), str(dst))
except OSError as exc:
QMessageBox.warning(self, "Ошибка", f"Не удалось переместить {src.name}:\n{exc}")
continue
moved += 1
self._results.pop(str(src), None)
self._files = [p for p in self._files if p != src]
self.file_list.takeItem(self.file_list.row(item))
if self._current == src:
self._current = None
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)))
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
def _unique_dest(folder: Path, name: str) -> Path:
"""Avoid clobbering: foo.jpg -> foo (1).jpg if it already exists."""
dst = folder / name
if not dst.exists():
return dst
stem, suffix = dst.stem, dst.suffix
i = 1
while (folder / f"{stem} ({i}){suffix}").exists():
i += 1
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 _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):
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 = {
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:
self.detail_header.setText(
f"{path.name} · {w}×{h} · не рассчитано "
"(двойной клик по файлу или «Рассчитать кадр»)"
)
self.detail_table.setRowCount(0)
self.view.set_highlight(None)
return
by_type: dict[str, int] = {}
for d in dets:
by_type[d.type.value] = by_type.get(d.type.value, 0) + 1
summary = ", ".join(f"{k}: {v}" for k, v in sorted(by_type.items())) or "ничего не найдено"
self.detail_header.setText(f"{path.name} · {w}×{h} · всего {len(dets)} ({summary})")
self.detail_table.blockSignals(True)
self.detail_table.setRowCount(len(dets))
for row, d in enumerate(dets):
x, y, bw, bh = d.bbox
cells = [d.type.value, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
for col, text in enumerate(cells):
self.detail_table.setItem(row, col, QTableWidgetItem(text))
self.detail_table.blockSignals(False)
self.detail_table.clearSelection()
self.detail_table.resizeColumnsToContents()
self.view.set_highlight(None)
def _on_detail_selected(self) -> None:
rows = self.detail_table.selectionModel().selectedRows()
self.view.set_highlight(rows[0].row() if rows else None)
def _on_threshold_changed(self, value: float) -> None:
self._cfg.default_threshold = value
self.view.set_threshold(value)
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)