1722 lines
79 KiB
Python
1722 lines
79 KiB
Python
"""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 · model · calc-frame · detect-all ·
|
||
restore · 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.
|
||
|
||
Detection is YOLO-only; restoration is DeepMosaics-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 the model clears the cache.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import os
|
||
import shutil
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from PySide6.QtCore import Qt, QThreadPool
|
||
from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcut
|
||
from PySide6.QtWidgets import (
|
||
QAbstractItemView,
|
||
QApplication,
|
||
QComboBox,
|
||
QDialog,
|
||
QDialogButtonBox,
|
||
QDoubleSpinBox,
|
||
QFileDialog,
|
||
QFormLayout,
|
||
QHBoxLayout,
|
||
QInputDialog,
|
||
QLabel,
|
||
QListWidget,
|
||
QListWidgetItem,
|
||
QMainWindow,
|
||
QMenu,
|
||
QMessageBox,
|
||
QPlainTextEdit,
|
||
QProgressBar,
|
||
QPushButton,
|
||
QSizePolicy,
|
||
QSplitter,
|
||
QTableWidget,
|
||
QTableWidgetItem,
|
||
QToolButton,
|
||
QVBoxLayout,
|
||
QWidget,
|
||
)
|
||
|
||
from .. import settings_store
|
||
from ..config import AppConfig
|
||
from ..core.detection import cache as detection_cache
|
||
from ..core.detection import registry as model_registry
|
||
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.factory import build_restorer, restorer_needs_detections
|
||
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
|
||
from .workers import Job
|
||
|
||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
|
||
|
||
|
||
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._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._restored_count = 0 # frames with output in restored/ (for the progress summary)
|
||
self._row_restored: set[str] = set() # frame paths that have a restored version (row ✓)
|
||
self._showing_restored = False
|
||
self._filter_mode = "all" # file-list filter (see filter_combo)
|
||
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._pool = QThreadPool.globalInstance()
|
||
self._job: Job | None = None # the running background job, if any
|
||
self._tick_count = 0 # throttles scrubber-mark refreshes during detect-all
|
||
self._device_info: dict | None = None # torch/CUDA probe result (for the badge)
|
||
self._probe_job: Job | None = None
|
||
|
||
self.setWindowTitle("HVideoTool — инспектор детекции цензуры")
|
||
self.resize(1180, 720)
|
||
|
||
self._build_toolbar()
|
||
self._build_central()
|
||
self._build_statusbar()
|
||
self._build_menu()
|
||
self._probe_device() # determine CUDA/CPU in the background and fill the badge
|
||
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.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
|
||
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||
file_menu.addAction("Расцензурить найденное (по детекции)", lambda: self._restore_all(only_detected=True))
|
||
file_menu.addAction("Открыть папку результатов", self._open_restored_dir)
|
||
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.setToolButtonStyle(Qt.ToolButtonTextOnly)
|
||
|
||
# --- Проект: rarely-touched session actions collapsed into one dropdown.
|
||
project_btn = self._dropdown_button("Проект ▾", "Действия с проектом")
|
||
m = project_btn.menu()
|
||
m.addAction("Создать проект…", self._create_project)
|
||
m.addAction("Открыть проект…", self._open_project_dialog)
|
||
m.addAction("Создать из ролика…", self._create_from_video)
|
||
m.addAction("Импортировать папку как проект…", self._import_folder_as_project)
|
||
tb.addWidget(project_btn)
|
||
tb.addSeparator()
|
||
|
||
# --- Детекторы: active YOLO models picker (unchanged).
|
||
tb.addWidget(QLabel(" Детекторы: "))
|
||
self._models_menu = QMenu(self)
|
||
self.models_button = QToolButton()
|
||
self.models_button.setPopupMode(QToolButton.InstantPopup)
|
||
self.models_button.setMenu(self._models_menu)
|
||
self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)")
|
||
tb.addWidget(self.models_button)
|
||
self._rebuild_models_menu()
|
||
tb.addSeparator()
|
||
|
||
# --- Детекция: primary "Рассчитать кадр" + the bulk runs in its dropdown.
|
||
detect_btn = self._split_button(
|
||
"Рассчитать кадр", self._recompute_current,
|
||
"Запустить детектор на выбранном кадре (Space / двойной клик по файлу)",
|
||
)
|
||
dm = detect_btn.menu()
|
||
dm.addAction("Детектировать все (дозапуск)", lambda: self._detect_all(False))
|
||
dm.addAction("Детектировать все заново", lambda: self._detect_all(True))
|
||
tb.addWidget(detect_btn)
|
||
|
||
# --- Расцензурить: primary "Расцензурить кадр" + the bulk/engine actions.
|
||
restore_btn = self._split_button(
|
||
"Расцензурить кадр", self._restore_current,
|
||
"Восстановить мозаику на текущем кадре (результат сохраняется в restored/)",
|
||
)
|
||
rm = restore_btn.menu()
|
||
rm.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
|
||
hits = rm.addAction("Расцензурить найденное (по детекции)", lambda: self._restore_all(only_detected=True))
|
||
hits.setToolTip(
|
||
"Расцензурить только кадры с детекцией (быстро, пропускает чистые).\n"
|
||
"ВНИМАНИЕ: YOLO ловит не всю мозаику — может пропустить."
|
||
)
|
||
rm.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||
rm.addSeparator()
|
||
rm.addAction("Движок восстановления…", self._open_restore_settings)
|
||
rm.addAction("Открыть папку результатов", self._open_restored_dir)
|
||
rm.addSeparator()
|
||
self.save_restored_action = QAction("Сохранить результат…", self, triggered=self._save_restored)
|
||
self.save_restored_action.setEnabled(False)
|
||
self.save_restored_action.setToolTip("Экспортировать <имя>_restored.jpg рядом с кадром")
|
||
rm.addAction(self.save_restored_action)
|
||
tb.addWidget(restore_btn)
|
||
|
||
# --- View toggle: kept visible (checkable) so the result/original state is obvious.
|
||
self.toggle_restored_action = QAction("Показать расцензуренное", self, triggered=self._toggle_restored)
|
||
self.toggle_restored_action.setCheckable(True)
|
||
self.toggle_restored_action.setEnabled(False)
|
||
self.toggle_restored_action.setShortcut("R")
|
||
self.toggle_restored_action.setToolTip(
|
||
"Переключить просмотр оригинал ⇄ расцензуренное для всего проекта (R)"
|
||
)
|
||
tb.addAction(self.toggle_restored_action)
|
||
tb.addSeparator()
|
||
|
||
# --- Stop: kept visible — must be reachable instantly during a long run.
|
||
self.stop_action = QAction("■ Стоп", self, triggered=self._request_cancel)
|
||
self.stop_action.setToolTip("Отменить текущую операцию (Esc)")
|
||
self.stop_action.setEnabled(False)
|
||
tb.addAction(self.stop_action)
|
||
|
||
# Push the threshold control to the right edge.
|
||
spacer = QWidget()
|
||
spacer.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Preferred)
|
||
tb.addWidget(spacer)
|
||
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 _dropdown_button(self, text: str, tooltip: str) -> QToolButton:
|
||
"""A toolbar button that just opens a menu (no default action)."""
|
||
btn = QToolButton()
|
||
btn.setText(text)
|
||
btn.setToolTip(tooltip)
|
||
btn.setPopupMode(QToolButton.InstantPopup)
|
||
btn.setMenu(QMenu(btn))
|
||
return btn
|
||
|
||
def _split_button(self, text: str, slot, tooltip: str) -> QToolButton:
|
||
"""A split button: click runs ``slot``; the arrow opens a menu of related actions."""
|
||
btn = QToolButton()
|
||
btn.setText(text)
|
||
btn.setToolTip(tooltip)
|
||
btn.setPopupMode(QToolButton.MenuButtonPopup)
|
||
action = QAction(text, btn, triggered=slot)
|
||
action.setToolTip(tooltip)
|
||
btn.setDefaultAction(action)
|
||
btn.setMenu(QMenu(btn))
|
||
return btn
|
||
|
||
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)
|
||
|
||
# Filter the (possibly huge) list down to the frames you care about.
|
||
self.filter_combo = QComboBox()
|
||
for label, mode in (
|
||
("Все кадры", "all"),
|
||
("С цензурой", "hits"),
|
||
("Чистые", "clean"),
|
||
("Не рассчитано", "uncomputed"),
|
||
("Расцензуренные", "restored"),
|
||
("Без расцензуривания", "unrestored"),
|
||
):
|
||
self.filter_combo.addItem(label, mode)
|
||
self.filter_combo.setToolTip("Показывать только кадры выбранной категории")
|
||
self.filter_combo.currentIndexChanged.connect(self._on_filter_changed)
|
||
|
||
# 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.filter_combo)
|
||
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())
|
||
# Processing summary under the scrubber: how much of the sequence is done.
|
||
self.stats_label = QLabel("")
|
||
self.stats_label.setAlignment(Qt.AlignCenter)
|
||
self.stats_label.setStyleSheet("QLabel{color:#888; padding:1px;}")
|
||
clayout.addWidget(self.stats_label)
|
||
|
||
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, 5)
|
||
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)
|
||
|
||
# Clickable position readout — opens "go to frame N" (handy on huge sequences
|
||
# where the scrubber is too coarse, ~50 frames/px on 29k).
|
||
self.pos_label = QPushButton("0 / 0")
|
||
self.pos_label.setFlat(True)
|
||
self.pos_label.setMinimumWidth(90)
|
||
self.pos_label.setCursor(Qt.PointingHandCursor)
|
||
self.pos_label.setToolTip("Перейти к кадру по номеру")
|
||
self.pos_label.clicked.connect(self._jump_to_frame)
|
||
|
||
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
|
||
# Skip rows hidden by the filter so prev/next walk only the visible frames.
|
||
i = self.file_list.currentRow() + delta
|
||
while 0 <= i < n:
|
||
if not self.file_list.item(i).isHidden():
|
||
self.file_list.setCurrentRow(i)
|
||
return
|
||
i += delta
|
||
|
||
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 _jump_to_frame(self) -> None:
|
||
"""Ask for a 1-based frame number and select it (clamped to range)."""
|
||
n = self.file_list.count()
|
||
if n == 0:
|
||
return
|
||
cur = self.file_list.currentRow() + 1
|
||
num, ok = QInputDialog.getInt(
|
||
self, "Перейти к кадру", f"Номер кадра (1–{n}):", cur, 1, n
|
||
)
|
||
if ok:
|
||
self.file_list.setCurrentRow(num - 1)
|
||
|
||
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.device_badge = QPushButton("⏳ устройство…")
|
||
self.device_badge.setFlat(True)
|
||
self.device_badge.setCursor(Qt.PointingHandCursor)
|
||
self.device_badge.setToolTip("Устройство вычислений (нажмите для подробностей)")
|
||
self.device_badge.clicked.connect(self._show_device_info)
|
||
self.statusBar().addPermanentWidget(self.device_badge)
|
||
|
||
self.progress = QProgressBar()
|
||
self.progress.setMaximumWidth(260)
|
||
self.progress.setVisible(False)
|
||
self.statusBar().addPermanentWidget(self.progress)
|
||
|
||
# ------------------------------------------------------------- device badge
|
||
def _probe_device(self) -> None:
|
||
"""Determine CUDA/CPU off the GUI thread (importing torch is slow)."""
|
||
from ..core import torch_info
|
||
|
||
job = Job(lambda _job: torch_info.gather())
|
||
self._probe_job = job # keep alive until `done`
|
||
job.signals.done.connect(self._set_device_badge)
|
||
job.signals.failed.connect(lambda _msg: self._set_device_badge(None))
|
||
self._pool.start(job)
|
||
|
||
def _set_device_badge(self, info: dict | None) -> None:
|
||
self._probe_job = None
|
||
self._device_info = info or {}
|
||
if self._device_info.get("cuda_available"):
|
||
name = self._device_info.get("device_name") or "GPU"
|
||
self.device_badge.setText("⚡ CUDA")
|
||
self.device_badge.setToolTip(f"Вычисления на GPU: {name} (нажмите для подробностей)")
|
||
self.device_badge.setStyleSheet("QPushButton{color:#16a085; font-weight:bold;}")
|
||
else:
|
||
self.device_badge.setText("🖥 CPU")
|
||
self.device_badge.setToolTip(
|
||
"Вычисления на CPU — нажмите, чтобы узнать почему и как включить GPU"
|
||
)
|
||
self.device_badge.setStyleSheet("QPushButton{color:#cc8400; font-weight:bold;}")
|
||
|
||
def _show_device_info(self) -> None:
|
||
from ..core import torch_info
|
||
|
||
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 = [
|
||
"ВЕРДИКТ:",
|
||
a["summary"],
|
||
"",
|
||
"Диагностика:",
|
||
*(f" • {d}" for d in a["details"]),
|
||
"",
|
||
"Что делать:",
|
||
a["steps"],
|
||
]
|
||
|
||
# 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:
|
||
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:
|
||
"""Enter a cancellable long operation. ``total=None`` => busy spinner."""
|
||
self._busy = True
|
||
self._cancel = False
|
||
self.stop_action.setEnabled(True)
|
||
# Disable inputs that would race a running job (they clear cache / rebuild engines).
|
||
self.models_button.setEnabled(False)
|
||
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.models_button.setEnabled(True)
|
||
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
|
||
if self._job is not None:
|
||
self._job.cancel() # stops the background loop at its next check
|
||
self.statusBar().showMessage("Отмена…")
|
||
|
||
# ------------------------------------------------------------- background jobs
|
||
def _start_job(self, fn, total: int | None, *, on_tick=None, on_done=None) -> None:
|
||
"""Run ``fn(job)`` on the thread pool; marshal results back to the GUI.
|
||
|
||
``on_tick(payload)`` handles incremental results (GUI thread); ``on_done(result,
|
||
cancelled)`` runs when the job finishes. Only one job runs at a time (callers
|
||
guard with ``self._busy``).
|
||
"""
|
||
self._begin_busy(total)
|
||
self._tick_count = 0
|
||
self._job_start = time.monotonic() # for ETA in the progress messages
|
||
job = Job(fn)
|
||
self._job = job
|
||
if on_tick is not None:
|
||
job.signals.tick.connect(on_tick)
|
||
job.signals.progress.connect(self._on_job_progress)
|
||
job.signals.done.connect(lambda result: self._finish_job(result, on_done))
|
||
job.signals.failed.connect(self._on_job_failed)
|
||
self._pool.start(job)
|
||
|
||
def _on_job_progress(self, done: int, total: int, message: str) -> None:
|
||
if total > 0:
|
||
self.progress.setRange(0, total)
|
||
self.progress.setValue(done)
|
||
if message:
|
||
eta = self._eta_suffix(done, total)
|
||
self.statusBar().showMessage(message + eta)
|
||
|
||
def _eta_suffix(self, done: int, total: int) -> str:
|
||
""" ' · осталось ~Xм Yс' estimated from the average rate so far (or '' if N/A)."""
|
||
start = getattr(self, "_job_start", None)
|
||
if not start or done <= 0 or total <= 0 or done >= total:
|
||
return ""
|
||
elapsed = time.monotonic() - start
|
||
if elapsed < 0.5:
|
||
return ""
|
||
remaining = elapsed / done * (total - done)
|
||
secs = int(remaining)
|
||
if secs >= 3600:
|
||
text = f"{secs // 3600}ч {secs % 3600 // 60}м"
|
||
elif secs >= 60:
|
||
text = f"{secs // 60}м {secs % 60}с"
|
||
else:
|
||
text = f"{secs}с"
|
||
return f" · осталось ~{text}"
|
||
|
||
def _finish_job(self, result, on_done) -> None:
|
||
cancelled = self._job.cancelled if self._job is not None else False
|
||
self._job = None
|
||
self._end_busy()
|
||
if on_done is not None:
|
||
on_done(result, cancelled)
|
||
|
||
def _on_job_failed(self, message: str) -> None:
|
||
self._job = None
|
||
self._end_busy()
|
||
QMessageBox.warning(self, "Ошибка", message)
|
||
|
||
# --------------------------------------------------------------- detector
|
||
def _make_detector(self):
|
||
d = self._cfg.detection
|
||
key = (
|
||
tuple(sorted(self._cfg.detector_models)), d.yolo_conf, d.yolo_imgsz,
|
||
self._cfg.cross_model_nms, self._cfg.nms_iou,
|
||
)
|
||
if key != self._detector_key:
|
||
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
|
||
self._detector_key = key
|
||
return self._detector
|
||
|
||
# --------------------------------------------------------- model selection
|
||
def _ensure_models(self) -> None:
|
||
"""Drop selected models that vanished; default to all discovered if none picked.
|
||
|
||
Called on project open. The factory raises a clear message if a detect runs with
|
||
nothing selected, so we don't prompt here.
|
||
"""
|
||
entries = model_registry.discover_models()
|
||
available = {e.path for e in entries}
|
||
kept = [m for m in self._cfg.detector_models if m in available]
|
||
if not kept and entries: # first run / fresh project — enable everything found
|
||
kept = [e.path for e in entries]
|
||
if kept != self._cfg.detector_models:
|
||
self._cfg.detector_models = kept
|
||
self._persist_settings()
|
||
|
||
def _rebuild_models_menu(self) -> None:
|
||
"""Repopulate the toolbar "Модели" menu with a checkable item per discovered model."""
|
||
self._models_menu.clear()
|
||
selected = set(self._cfg.detector_models)
|
||
entries = model_registry.discover_models()
|
||
if not entries:
|
||
self._models_menu.addAction("(нет моделей в models/yolo/<категория>)").setEnabled(False)
|
||
else:
|
||
last_cat = None
|
||
for e in entries:
|
||
if e.category != last_cat:
|
||
self._models_menu.addSection(e.category)
|
||
last_cat = e.category
|
||
act = self._models_menu.addAction(e.name)
|
||
act.setCheckable(True)
|
||
act.setChecked(e.path in selected)
|
||
act.toggled.connect(lambda on, p=e.path: self._on_model_toggled(p, on))
|
||
self._models_menu.addSeparator()
|
||
nms = self._models_menu.addAction("Объединять пересечения (NMS)")
|
||
nms.setCheckable(True)
|
||
nms.setChecked(self._cfg.cross_model_nms)
|
||
nms.setToolTip(
|
||
"Убирать дублирующие рамки от перекрывающихся моделей (по IoU; остаётся\n"
|
||
"рамка с большей уверенностью). Меняет результат — кэш пересчитывается."
|
||
)
|
||
nms.toggled.connect(self._on_nms_toggled)
|
||
self._models_menu.addAction("Пороги по моделям…", self._edit_model_thresholds)
|
||
self._models_menu.addSeparator()
|
||
self._models_menu.addAction("Добавить модель…", self._add_model)
|
||
self._models_menu.addAction("Открыть папку моделей", self._open_models_dir)
|
||
self._models_menu.addAction("Обновить список", self._rebuild_models_menu)
|
||
self._update_models_button()
|
||
|
||
def _on_nms_toggled(self, on: bool) -> None:
|
||
self._cfg.cross_model_nms = on
|
||
self._persist_settings()
|
||
self._invalidate_results() # merging changes detections => recompute
|
||
self.statusBar().showMessage(
|
||
"Объединение пересечений (NMS): " + ("вкл" if on else "выкл")
|
||
)
|
||
|
||
def _edit_model_thresholds(self) -> None:
|
||
"""Dialog: per-model overlay threshold overrides (display-only, not detection)."""
|
||
models = [m for m in self._cfg.detector_models if Path(m).is_file()]
|
||
if not models:
|
||
QMessageBox.information(
|
||
self, "Нет моделей", "Сначала отметьте хотя бы одну модель."
|
||
)
|
||
return
|
||
dlg = QDialog(self)
|
||
dlg.setWindowTitle("Пороги отображения по моделям")
|
||
layout = QVBoxLayout(dlg)
|
||
layout.addWidget(QLabel(
|
||
f"Порог отображения для каждой модели (по умолчанию {self._cfg.default_threshold:.2f}).\n"
|
||
"Влияет только на отрисовку/подсветку, не на саму детекцию."
|
||
))
|
||
form = QFormLayout()
|
||
spins: dict[str, QDoubleSpinBox] = {}
|
||
for m in models:
|
||
stem = Path(m).stem
|
||
spin = QDoubleSpinBox()
|
||
spin.setRange(0.0, 1.0)
|
||
spin.setSingleStep(0.05)
|
||
spin.setValue(self._cfg.model_thresholds.get(stem, self._cfg.default_threshold))
|
||
spins[stem] = spin
|
||
form.addRow(stem, spin)
|
||
layout.addLayout(form)
|
||
buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
|
||
buttons.accepted.connect(dlg.accept)
|
||
buttons.rejected.connect(dlg.reject)
|
||
layout.addWidget(buttons)
|
||
if dlg.exec() != QDialog.Accepted:
|
||
return
|
||
# Store only the overrides that differ from the global default — keeps it tidy.
|
||
thresholds = {
|
||
stem: round(spin.value(), 4)
|
||
for stem, spin in spins.items()
|
||
if abs(spin.value() - self._cfg.default_threshold) > 1e-9
|
||
}
|
||
self._cfg.model_thresholds = thresholds
|
||
self.view.set_model_thresholds(thresholds)
|
||
self._persist_settings()
|
||
self.statusBar().showMessage(
|
||
f"Пороги по моделям обновлены ({len(thresholds)} переопределений)"
|
||
)
|
||
|
||
def _update_models_button(self) -> None:
|
||
n = len([m for m in self._cfg.detector_models if Path(m).is_file()])
|
||
self.models_button.setText(f"Модели ({n}) ▾")
|
||
|
||
def _on_model_toggled(self, path: str, on: bool) -> None:
|
||
sel = [m for m in self._cfg.detector_models if m != path]
|
||
if on:
|
||
sel.append(path)
|
||
self._cfg.detector_models = sel
|
||
self._persist_settings()
|
||
self._invalidate_results() # different model set => recompute
|
||
self._update_models_button()
|
||
|
||
def _add_model(self) -> None:
|
||
"""Copy a chosen .pt into models/yolo/<category>/ and tick it."""
|
||
path, _ = QFileDialog.getOpenFileName(
|
||
self, "Выберите веса YOLO (.pt)", str(Path.cwd()), "Веса YOLO (*.pt)"
|
||
)
|
||
if not path:
|
||
return
|
||
category, ok = QInputDialog.getText(
|
||
self, "Категория модели",
|
||
"Категория (папка под models/yolo, напр. mosaic, face):", text="misc"
|
||
)
|
||
if not ok:
|
||
return
|
||
category = (category.strip() or "misc")
|
||
dest_dir = model_registry.yolo_root() / category
|
||
dest_dir.mkdir(parents=True, exist_ok=True)
|
||
dest = self._unique_dest(dest_dir, Path(path).name)
|
||
try:
|
||
shutil.copy2(path, dest)
|
||
except OSError as exc:
|
||
QMessageBox.warning(self, "Ошибка", f"Не удалось скопировать модель:\n{exc}")
|
||
return
|
||
self._cfg.detector_models = [*self._cfg.detector_models, str(dest)]
|
||
self._persist_settings()
|
||
self._rebuild_models_menu()
|
||
self._invalidate_results()
|
||
self.statusBar().showMessage(f"Модель добавлена: {dest.name} → {category}")
|
||
|
||
def _open_models_dir(self) -> None:
|
||
root = model_registry.yolo_root()
|
||
root.mkdir(parents=True, exist_ok=True)
|
||
with contextlib.suppress(OSError, AttributeError):
|
||
os.startfile(str(root)) # noqa: S606 - Windows: open in Explorer
|
||
|
||
def _open_restored_dir(self) -> None:
|
||
"""Open the project's restored/ folder (where restored images are saved)."""
|
||
if self._project is None:
|
||
QMessageBox.information(self, "Нет проекта", "Сначала откройте проект.")
|
||
return
|
||
d = self._project.restored_dir
|
||
d.mkdir(parents=True, exist_ok=True)
|
||
with contextlib.suppress(OSError, AttributeError):
|
||
os.startfile(str(d)) # noqa: S606 - Windows: open in Explorer
|
||
|
||
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:
|
||
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):
|
||
with contextlib.suppress(OSError, ValueError):
|
||
self._open_project(Project.load(last))
|
||
|
||
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():
|
||
with contextlib.suppress(OSError):
|
||
shutil.copy2(str(old_sidecar), str(project.cache_path))
|
||
|
||
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._ensure_models() # prune missing / default-select discovered models
|
||
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._rebuild_models_menu() # reflect this project's model selection
|
||
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)
|
||
self.view.set_model_thresholds(self._cfg.model_thresholds)
|
||
|
||
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, jpg_quality = 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, jpg_quality=jpg_quality, progress=cb,
|
||
)
|
||
except Exception as exc:
|
||
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._files = files
|
||
self._results.clear()
|
||
self._row_restored.clear()
|
||
self._showing_restored = False # start a project in original-view mode
|
||
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()
|
||
self._refresh_restored_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) in self._results:
|
||
self._show(path)
|
||
return
|
||
self._detect_one(path, then_show=True)
|
||
|
||
@staticmethod
|
||
def _compute(detector, path: Path) -> list[Detection]:
|
||
"""Pure read + detect for one image (runs on a worker thread; no Qt)."""
|
||
img = imread_unicode(str(path))
|
||
if img is None:
|
||
raise RuntimeError(f"Не удалось прочитать: {Path(path).name}")
|
||
dets = detector.detect(Frame(image=img))
|
||
dets.sort(key=lambda d: d.score, reverse=True)
|
||
return dets
|
||
|
||
def _detect_one(self, path: Path, *, then_show: bool) -> None:
|
||
"""Detect one image on a background thread, then cache/tag/show it."""
|
||
if self._busy:
|
||
return
|
||
|
||
def fn(job):
|
||
return (str(path), self._compute(self._make_detector(), path))
|
||
|
||
def done(result, cancelled):
|
||
if result is None:
|
||
return
|
||
key, dets = result
|
||
self._results[key] = dets
|
||
self._tag_file(Path(key))
|
||
self._refresh_marks()
|
||
self._save_results()
|
||
if then_show or self._current == Path(key):
|
||
self._show(Path(key))
|
||
self.statusBar().showMessage(f"Детекция: {Path(key).name} — {len(dets)} обл.")
|
||
|
||
self.statusBar().showMessage(f"Детекция: {path.name}…")
|
||
self._start_job(fn, None, on_done=done)
|
||
|
||
def _show(self, path: Path) -> None:
|
||
"""Display the image (does not run the detector).
|
||
|
||
Honours the global "show restored" view mode (``_showing_restored``): when on and
|
||
a restored version exists (memory or ``restored/``), the restored image is shown
|
||
(no overlays); otherwise the original frame with its cached detections.
|
||
"""
|
||
self._current = path
|
||
dets = self._results.get(str(path)) # None => not yet computed
|
||
restored = self._restored_image_for(path) if self._showing_restored else None
|
||
if restored is not None:
|
||
self.view.set_image(restored, []) # restored: no overlays
|
||
self._fill_detail_table(path, restored, dets)
|
||
else:
|
||
img = imread_unicode(str(path))
|
||
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 (background)."""
|
||
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
|
||
self._detect_one(self._current, then_show=True)
|
||
|
||
def _detect_all(self, force: bool = False) -> None:
|
||
"""Detect the whole folder on a background thread. ``force`` clears the cache
|
||
first (full regen); otherwise already-computed frames are skipped (resume/top-up).
|
||
The GUI stays responsive — results stream in via per-frame ticks."""
|
||
if not self._files or self._busy:
|
||
return
|
||
if force:
|
||
self._clear_results()
|
||
pending = [p for p in self._files if str(p) not in self._results]
|
||
if not pending:
|
||
self.statusBar().showMessage("Все кадры уже посчитаны (см. «Все заново»)")
|
||
return
|
||
total = len(pending)
|
||
|
||
def fn(job):
|
||
detector = self._make_detector() # built on the worker thread (may raise)
|
||
for i, p in enumerate(pending, 1):
|
||
if job.cancelled:
|
||
break
|
||
try:
|
||
dets = self._compute(detector, p)
|
||
except RuntimeError:
|
||
continue # unreadable image — skip, keep going
|
||
job.tick((str(p), dets))
|
||
job.progress(i, total, f"Детекция {i}/{total}: {p.name}")
|
||
return None
|
||
|
||
def done(_result, cancelled):
|
||
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
||
self._refresh_marks()
|
||
self._save_results() # persist progress (completed or cancelled)
|
||
if cancelled:
|
||
self.statusBar().showMessage(f"Отменено · детекции на {hits} картинках")
|
||
else:
|
||
self.statusBar().showMessage(
|
||
f"Готово: детекции на {hits} из {len(self._files)} картинок"
|
||
)
|
||
if self._current is not None:
|
||
self._show(self._current)
|
||
|
||
self._start_job(fn, total, on_tick=self._apply_detection, on_done=done)
|
||
|
||
def _apply_detection(self, payload) -> None:
|
||
"""GUI-thread handler for one streamed detect-all result."""
|
||
key, dets = payload
|
||
self._results[key] = dets
|
||
self._tag_file(Path(key))
|
||
# 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 on a background thread, then show it.
|
||
|
||
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)
|
||
# The diffusion engine builds its mask from detections — capture them on the GUI
|
||
# thread (DeepMosaics ignores them). Warn if it needs them but none are computed.
|
||
needs_dets = restorer_needs_detections(self._cfg.restorer)
|
||
dets_for_restore = list(self._results.get(key, []))
|
||
if needs_dets and not dets_for_restore:
|
||
self.statusBar().showMessage(
|
||
"Diffusion перерисовывает по детекции — на этом кадре цензура не найдена "
|
||
"(сначала «Рассчитать кадр»)"
|
||
)
|
||
return
|
||
|
||
def fn(job):
|
||
img = imread_unicode(key)
|
||
if img is None:
|
||
raise RuntimeError(f"Не удалось прочитать: {path.name}")
|
||
restorer = self._make_restorer()
|
||
dets = dets_for_restore if restorer.needs_detections else []
|
||
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
|
||
return ("restored", key, restored, restorer.name)
|
||
|
||
def done(result, cancelled):
|
||
if cancelled:
|
||
self.statusBar().showMessage("Восстановление отменено")
|
||
return
|
||
if result is None:
|
||
return
|
||
_, k, restored, engine = result
|
||
self._restored[k] = restored
|
||
# Persist to the project's restored/ folder (like the batch run), so a single
|
||
# restore is saved on disk and survives reopening — not just held in memory.
|
||
saved_to = ""
|
||
if self._project is not None:
|
||
try:
|
||
self._project.restored_dir.mkdir(parents=True, exist_ok=True)
|
||
dst = self._project.restored_dir / f"{Path(k).stem}.jpg"
|
||
if imwrite_unicode(str(dst), restored):
|
||
saved_to = f" → {self._project.restored_dir.name}/"
|
||
except OSError:
|
||
pass
|
||
self._refresh_restored_marks()
|
||
if self._current is not None and str(self._current) == k:
|
||
self._showing_restored = True
|
||
self._show(self._current)
|
||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}{saved_to}")
|
||
|
||
self.statusBar().showMessage(f"Восстановление: {path.name}…")
|
||
self._start_job(fn, None, on_done=done)
|
||
|
||
def _restore_all(self, force: bool = False, only_detected: bool = False) -> None:
|
||
"""Restore frames 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 a contiguous sequence in order via
|
||
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
|
||
|
||
``only_detected`` uses the YOLO detection cache to skip frames known clean:
|
||
per-frame → restore just the frames with detections; temporal → restrict the run
|
||
to the contiguous span [first hit … last hit] (clean frames inside it still run,
|
||
for recurrence). NOTE: LADA misses some mosaic, so this can miss censorship YOLO
|
||
didn't flag — "Расцензурить все" stays the thorough option.
|
||
"""
|
||
if not self._files or self._project is None or self._busy:
|
||
return
|
||
files = list(self._files) # snapshot — favorites/move mutate self._files
|
||
is_temporal = self._cfg.restorer == "deepmosaics_video"
|
||
# The diffusion engine masks the detections, so it inherently runs only on hits and
|
||
# requires detection to be computed — same gating as the explicit "найденное" mode.
|
||
needs_dets = restorer_needs_detections(self._cfg.restorer)
|
||
hits = [i for i, p in enumerate(files) if self._results.get(str(p))]
|
||
# Per-index detection snapshot (GUI-thread read) — fed to engines that need a mask.
|
||
dets_by_index = (
|
||
{i: list(self._results.get(str(files[i]), [])) for i in hits} if needs_dets else {}
|
||
)
|
||
need_hits = only_detected or needs_dets
|
||
|
||
if need_hits:
|
||
if not self._results:
|
||
self.statusBar().showMessage(
|
||
"Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)"
|
||
)
|
||
return
|
||
if not hits:
|
||
self.statusBar().showMessage("Цензура не найдена ни на одном кадре — нечего расцензуривать")
|
||
return
|
||
span = (hits[0], hits[-1]) if is_temporal else None
|
||
total = (span[1] - span[0] + 1) if span else len(hits)
|
||
else:
|
||
span = None
|
||
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
|
||
state = {"done": 0} # frames processed (for the progress bar)
|
||
|
||
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, *, verb="Расцензуривание"):
|
||
if restored is not None:
|
||
imwrite_unicode(str(out_path(files[i])), restored)
|
||
state["done"] += 1
|
||
job.progress(state["done"], total, f"{verb} {state['done']}/{total}: {files[i].name}")
|
||
|
||
if restorer.temporal:
|
||
start = span[0] if span else 0
|
||
end = span[1] if span else len(files) - 1
|
||
restorer.restore_sequence(
|
||
end - start + 1,
|
||
lambda li: get_frame(start + li),
|
||
lambda _li: [],
|
||
lambda li, res: emit(start + li, res),
|
||
should_cancel=lambda: job.cancelled,
|
||
)
|
||
else:
|
||
indices = hits if need_hits else range(len(files))
|
||
use_dets = restorer.needs_detections
|
||
for i in indices:
|
||
if job.cancelled:
|
||
break
|
||
p = files[i]
|
||
if not force and out_path(p).is_file():
|
||
emit(i, None, verb="Пропуск") # already restored — count, don't rewrite
|
||
continue
|
||
dets = dets_by_index.get(i, []) if use_dets else []
|
||
emit(i, restorer.restore(get_frame(i), dets, should_cancel=lambda: job.cancelled))
|
||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||
return None
|
||
|
||
def done(_result, cancelled):
|
||
self._refresh_restored_marks()
|
||
if not cancelled and self._row_restored:
|
||
self._showing_restored = True # auto-switch to viewing the results
|
||
if self._current is not None: # re-show current frame in the (new) mode
|
||
self._show(self._current)
|
||
self.statusBar().showMessage(
|
||
"Расцензуривание отменено" if cancelled
|
||
else f"Готово: результаты в {out_dir.name}/ ({total} кадров) — показываю расцензуренное"
|
||
)
|
||
|
||
scope = " (по детекции)" if only_detected else ""
|
||
self.statusBar().showMessage(f"Пакетное расцензуривание{scope}…")
|
||
self._start_job(fn, total, on_done=done)
|
||
|
||
def _make_restorer(self):
|
||
c = self._cfg
|
||
key = (
|
||
c.restorer, c.dm_dir, c.dm_model, c.dm_gpu,
|
||
# diffusion identity — changing any of these must rebuild the engine
|
||
c.diff_backend, c.diff_url, c.diff_model, c.diff_prompt, c.diff_negative,
|
||
c.diff_steps, c.diff_cfg, c.diff_denoise, c.diff_seed,
|
||
c.diff_mask_dilate, c.diff_mask_blur,
|
||
)
|
||
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}")
|
||
|
||
# ---------------------------------------------------------- restored access
|
||
def _restored_disk_path(self, path: Path) -> Path | None:
|
||
"""Path of the restored output for ``path`` in ``restored/``, if it exists."""
|
||
if self._project is None:
|
||
return None
|
||
rp = self._project.restored_dir / f"{Path(path).stem}.jpg"
|
||
return rp if rp.is_file() else None
|
||
|
||
def _has_restored(self, path: Path) -> bool:
|
||
"""Cheap check (no decode): is there a restored version of ``path``?"""
|
||
return str(path) in self._restored or self._restored_disk_path(path) is not None
|
||
|
||
def _restored_image_for(self, path: Path):
|
||
"""Return the restored image for ``path`` (from memory or ``restored/``), or None."""
|
||
img = self._restored.get(str(path))
|
||
if img is not None:
|
||
return img
|
||
rp = self._restored_disk_path(path)
|
||
return imread_unicode(str(rp)) if rp is not None else None
|
||
|
||
def _toggle_restored(self) -> None:
|
||
"""Flip the global view mode between original and restored, then re-show."""
|
||
self._showing_restored = not self._showing_restored
|
||
if self._current is not None:
|
||
self._show(self._current)
|
||
else:
|
||
self._update_restore_actions()
|
||
self.statusBar().showMessage(
|
||
"Показ: расцензуренное (где есть)" if self._showing_restored else "Показ: оригинал"
|
||
)
|
||
|
||
def _update_restore_actions(self) -> None:
|
||
any_restored = bool(self._row_restored) or bool(self._restored)
|
||
self.toggle_restored_action.setEnabled(any_restored)
|
||
# Keep the checkable state in sync with the mode (setChecked emits `toggled`,
|
||
# not `triggered`, so this never re-enters `_toggle_restored`).
|
||
self.toggle_restored_action.setChecked(self._showing_restored and any_restored)
|
||
self.toggle_restored_action.setText(
|
||
"Показать оригинал" if self._showing_restored else "Показать расцензуренное"
|
||
)
|
||
self.save_restored_action.setEnabled(
|
||
self._current is not None and self._has_restored(self._current)
|
||
)
|
||
|
||
def _save_restored(self) -> None:
|
||
if self._current is None:
|
||
return
|
||
restored = self._restored_image_for(self._current)
|
||
if restored is None:
|
||
return
|
||
out = self._unique_dest(self._current.parent, f"{self._current.stem}_restored.jpg")
|
||
if imwrite_unicode(str(out), restored):
|
||
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._refresh_restored_marks()
|
||
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 _relabel_row(self, item: QListWidgetItem) -> None:
|
||
"""Set a row's text/tint/tooltip from its detection + restoration state.
|
||
|
||
Text: ``name · <count|—> ✓`` — the count/— suffix appears once detected (red tint
|
||
= censorship, green = clean), and a trailing ✓ marks frames that have a restored
|
||
version in ``restored/``.
|
||
"""
|
||
base = item.data(Qt.UserRole + 1)
|
||
path = item.data(Qt.UserRole)
|
||
restored = path in self._row_restored
|
||
if path in self._results: # `in`, not truthy: empty list = clean
|
||
dets = self._results[path]
|
||
suffix = f" · {len(dets)}" if dets else " · —"
|
||
item.setBackground(self._TINT_HIT if dets else self._TINT_CLEAN)
|
||
else:
|
||
suffix = ""
|
||
item.setBackground(QBrush())
|
||
item.setText(f"{base}{suffix}{' ✓' if restored else ''}")
|
||
item.setToolTip("Есть расцензуренная версия (restored/)" if restored else "")
|
||
item.setHidden(not self._row_matches_filter(path))
|
||
|
||
def _row_matches_filter(self, path: str) -> bool:
|
||
"""Whether a row should be visible under the current filter mode."""
|
||
mode = self._filter_mode
|
||
if mode == "all":
|
||
return True
|
||
if mode == "hits":
|
||
return bool(self._results.get(path))
|
||
if mode == "clean":
|
||
return path in self._results and not self._results[path]
|
||
if mode == "uncomputed":
|
||
return path not in self._results
|
||
if mode == "restored":
|
||
return path in self._row_restored
|
||
if mode == "unrestored":
|
||
return path not in self._row_restored
|
||
return True
|
||
|
||
def _on_filter_changed(self) -> None:
|
||
self._filter_mode = self.filter_combo.currentData() or "all"
|
||
self._relabel_all() # re-applies hidden state per row
|
||
# If the current row got hidden, jump to the first visible one so the view isn't stale.
|
||
cur = self.file_list.currentItem()
|
||
if cur is not None and cur.isHidden():
|
||
for i in range(self.file_list.count()):
|
||
if not self.file_list.item(i).isHidden():
|
||
self.file_list.setCurrentRow(i)
|
||
break
|
||
n_vis = sum(1 for i in range(self.file_list.count()) if not self.file_list.item(i).isHidden())
|
||
self.statusBar().showMessage(
|
||
f"Фильтр: {self.filter_combo.currentText()} — показано {n_vis} из {len(self._files)}"
|
||
)
|
||
|
||
def _relabel_all(self) -> None:
|
||
for i in range(self.file_list.count()):
|
||
self._relabel_row(self.file_list.item(i))
|
||
|
||
def _tag_file(self, path: Path) -> None:
|
||
for i in range(self.file_list.count()):
|
||
item = self.file_list.item(i)
|
||
if item.data(Qt.UserRole) == str(path):
|
||
self._relabel_row(item)
|
||
return
|
||
|
||
# ------------------------------------------------------------- result cache
|
||
def _results_key(self) -> dict:
|
||
"""Detector identity used to tag/validate the on-disk detection cache."""
|
||
d = self._cfg.detection
|
||
nms_iou = self._cfg.nms_iou if self._cfg.cross_model_nms else None
|
||
return detection_cache.make_key(
|
||
self._cfg.detector_models, d.yolo_conf, d.yolo_imgsz, nms_iou=nms_iou
|
||
)
|
||
|
||
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
|
||
self._relabel_all()
|
||
return len(cached)
|
||
|
||
def _clear_results(self) -> None:
|
||
"""Drop all cached detections and reset row labels/tints (keeps the detector).
|
||
|
||
Restored ✓ markers stay — restoration is independent of detection.
|
||
"""
|
||
self._results.clear()
|
||
self._relabel_all()
|
||
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)
|
||
self._update_counts_label()
|
||
|
||
def _update_counts_label(self) -> None:
|
||
"""Refresh the processing summary under the scrubber (cheap; counts only)."""
|
||
n = len(self._files)
|
||
detected = sum(1 for p in self._files if str(p) in self._results)
|
||
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
||
if n == 0:
|
||
self.stats_label.setText("")
|
||
return
|
||
self.stats_label.setText(
|
||
f"Кадров: {n} · детектировано: {detected}/{n} (с цензурой: {hits})"
|
||
f" · расцензурено: {self._restored_count}/{n}"
|
||
)
|
||
|
||
def _refresh_restored_marks(self) -> None:
|
||
"""Scan ``restored/`` and project restored frames onto the scrubber (green).
|
||
|
||
Done only on load / after a restore op (not per-frame) — it touches the disk.
|
||
"""
|
||
stems: set[str] = set()
|
||
if self._project is not None and self._project.restored_dir.is_dir():
|
||
stems = {p.stem for p in self._project.restored_dir.glob("*.jpg")}
|
||
mem = set(self._restored) # single-frame restores held in memory (not on disk yet)
|
||
rows: set[int] = set()
|
||
paths: set[str] = set()
|
||
if stems or mem:
|
||
for i in range(self.file_list.count()):
|
||
sp = self.file_list.item(i).data(Qt.UserRole)
|
||
if Path(sp).stem in stems or sp in mem:
|
||
rows.add(i)
|
||
paths.add(sp)
|
||
self._row_restored = paths
|
||
self._restored_count = len(rows)
|
||
self.frame_slider.set_restored_marks(rows)
|
||
self._relabel_all() # show/refresh the ✓ markers in the file list
|
||
self._update_counts_label()
|
||
|
||
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"<b>{path.name}</b> · {w}×{h} · <i>не рассчитано</i> "
|
||
"(двойной клик по файлу или «Рассчитать кадр»)"
|
||
)
|
||
self.detail_table.setRowCount(0)
|
||
self.view.set_highlight(None)
|
||
return
|
||
|
||
by_type: dict[str, int] = {}
|
||
for d in dets:
|
||
by_type[d.display] = by_type.get(d.display, 0) + 1
|
||
summary = ", ".join(f"{k}: {v}" for k, v in sorted(by_type.items())) or "ничего не найдено"
|
||
self.detail_header.setText(f"<b>{path.name}</b> · {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.display, d.model or "—", 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:
|
||
if self._job is not None: # stop a running background job before tearing down
|
||
self._job.cancel()
|
||
self._pool.waitForDone(3000)
|
||
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)
|