Enhance HVideoTool's detection and restoration capabilities: introduced per-model overlay threshold settings and cross-model non-maximum suppression (NMS) to improve detection accuracy. Updated configuration management to support these features, and refined the UI for better user experience. Documentation in CLAUDE.md has been updated to reflect these changes.
This commit is contained in:
+427
-105
@@ -25,6 +25,7 @@ from __future__ import annotations
|
||||
import contextlib
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from PySide6.QtCore import Qt, QThreadPool
|
||||
@@ -32,9 +33,12 @@ from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcu
|
||||
from PySide6.QtWidgets import (
|
||||
QAbstractItemView,
|
||||
QApplication,
|
||||
QComboBox,
|
||||
QDialog,
|
||||
QDialogButtonBox,
|
||||
QDoubleSpinBox,
|
||||
QFileDialog,
|
||||
QFormLayout,
|
||||
QHBoxLayout,
|
||||
QInputDialog,
|
||||
QLabel,
|
||||
@@ -46,6 +50,7 @@ from PySide6.QtWidgets import (
|
||||
QPlainTextEdit,
|
||||
QProgressBar,
|
||||
QPushButton,
|
||||
QSizePolicy,
|
||||
QSplitter,
|
||||
QTableWidget,
|
||||
QTableWidgetItem,
|
||||
@@ -89,7 +94,9 @@ class MainWindow(QMainWindow):
|
||||
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
|
||||
@@ -126,6 +133,8 @@ class MainWindow(QMainWindow):
|
||||
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()
|
||||
@@ -134,14 +143,19 @@ class MainWindow(QMainWindow):
|
||||
def _build_toolbar(self) -> None:
|
||||
tb = self.addToolBar("Главная")
|
||||
tb.setMovable(False)
|
||||
tb.setToolButtonStyle(Qt.ToolButtonTextOnly)
|
||||
|
||||
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)
|
||||
# --- Проект: 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()
|
||||
@@ -150,41 +164,62 @@ class MainWindow(QMainWindow):
|
||||
self.models_button.setToolTip("Выбрать активные YOLO-модели (models/yolo/<категория>)")
|
||||
tb.addWidget(self.models_button)
|
||||
self._rebuild_models_menu()
|
||||
|
||||
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)
|
||||
|
||||
# --- Детекция: 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)
|
||||
|
||||
tb.addSeparator()
|
||||
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
|
||||
restore.setToolTip("Восстановить найденные области на текущем кадре")
|
||||
tb.addAction(restore)
|
||||
restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False))
|
||||
restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)")
|
||||
tb.addAction(restore_all)
|
||||
restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True))
|
||||
restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)")
|
||||
tb.addAction(restore_regen)
|
||||
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
|
||||
self.toggle_restored_action.setEnabled(False)
|
||||
tb.addAction(self.toggle_restored_action)
|
||||
self.save_restored_action = QAction("Сохранить результат", self, triggered=self._save_restored)
|
||||
self.save_restored_action.setEnabled(False)
|
||||
tb.addAction(self.save_restored_action)
|
||||
|
||||
tb.addSeparator()
|
||||
# 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)
|
||||
@@ -193,12 +228,47 @@ class MainWindow(QMainWindow):
|
||||
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)")
|
||||
@@ -208,6 +278,7 @@ class MainWindow(QMainWindow):
|
||||
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)
|
||||
|
||||
@@ -269,9 +340,14 @@ class MainWindow(QMainWindow):
|
||||
self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией")
|
||||
self.frame_slider.valueChanged.connect(self._on_slider)
|
||||
|
||||
self.pos_label = QLabel("0 / 0")
|
||||
# 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.setAlignment(Qt.AlignCenter)
|
||||
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("Предыдущий кадр с детекцией ([)")
|
||||
@@ -297,8 +373,13 @@ class MainWindow(QMainWindow):
|
||||
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)
|
||||
# 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."""
|
||||
@@ -318,6 +399,18 @@ class MainWindow(QMainWindow):
|
||||
"(сначала «Детектировать все»)"
|
||||
)
|
||||
|
||||
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
|
||||
@@ -469,6 +562,7 @@ class MainWindow(QMainWindow):
|
||||
"""
|
||||
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:
|
||||
@@ -483,7 +577,26 @@ class MainWindow(QMainWindow):
|
||||
self.progress.setRange(0, total)
|
||||
self.progress.setValue(done)
|
||||
if message:
|
||||
self.statusBar().showMessage(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
|
||||
@@ -500,7 +613,10 @@ class MainWindow(QMainWindow):
|
||||
# --------------------------------------------------------------- detector
|
||||
def _make_detector(self):
|
||||
d = self._cfg.detection
|
||||
key = (tuple(sorted(self._cfg.detector_models)), d.yolo_conf, d.yolo_imgsz)
|
||||
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
|
||||
@@ -540,11 +656,74 @@ class MainWindow(QMainWindow):
|
||||
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}) ▾")
|
||||
@@ -592,6 +771,16 @@ class MainWindow(QMainWindow):
|
||||
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.
|
||||
|
||||
@@ -775,6 +964,7 @@ class MainWindow(QMainWindow):
|
||||
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."""
|
||||
@@ -795,7 +985,7 @@ class MainWindow(QMainWindow):
|
||||
dialog = ExtractDialog(self)
|
||||
if dialog.exec() != QDialog.Accepted:
|
||||
return
|
||||
keyframes_only, step, max_dim = dialog.options()
|
||||
keyframes_only, step, max_dim, jpg_quality = dialog.options()
|
||||
video = Path(path)
|
||||
root = video.parent / f"{video.stem}_frames"
|
||||
if Project.is_project(root):
|
||||
@@ -821,7 +1011,7 @@ class MainWindow(QMainWindow):
|
||||
try:
|
||||
saved = extract_frames(
|
||||
str(video), str(out), step=step, keyframes_only=keyframes_only,
|
||||
max_dim=max_dim, progress=cb,
|
||||
max_dim=max_dim, jpg_quality=jpg_quality, progress=cb,
|
||||
)
|
||||
except Exception as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
|
||||
@@ -849,6 +1039,8 @@ class MainWindow(QMainWindow):
|
||||
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)
|
||||
@@ -926,7 +1118,7 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
key, dets = result
|
||||
self._results[key] = dets
|
||||
self._tag_file(Path(key), len(dets))
|
||||
self._tag_file(Path(key))
|
||||
self._refresh_marks()
|
||||
self._save_results()
|
||||
if then_show or self._current == Path(key):
|
||||
@@ -937,13 +1129,22 @@ class MainWindow(QMainWindow):
|
||||
self._start_job(fn, None, on_done=done)
|
||||
|
||||
def _show(self, path: Path) -> None:
|
||||
"""Display the image with its cached detections (does not run the detector)."""
|
||||
"""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
|
||||
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)
|
||||
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:
|
||||
@@ -1000,7 +1201,7 @@ class MainWindow(QMainWindow):
|
||||
"""GUI-thread handler for one streamed detect-all result."""
|
||||
key, dets = payload
|
||||
self._results[key] = dets
|
||||
self._tag_file(Path(key), len(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)
|
||||
@@ -1036,28 +1237,61 @@ class MainWindow(QMainWindow):
|
||||
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.view.set_image(restored, [])
|
||||
self._update_restore_actions()
|
||||
self._refresh_restored_marks()
|
||||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}")
|
||||
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) -> None:
|
||||
"""Restore every frame on a background thread, writing results to ``restored/``.
|
||||
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 the whole contiguous sequence in order via
|
||||
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
|
||||
total = len(files)
|
||||
is_temporal = self._cfg.restorer == "deepmosaics_video"
|
||||
hits = [i for i, p in enumerate(files) if self._results.get(str(p))]
|
||||
|
||||
if only_detected:
|
||||
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)
|
||||
|
||||
@@ -1067,6 +1301,7 @@ class MainWindow(QMainWindow):
|
||||
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)
|
||||
@@ -1079,42 +1314,48 @@ class MainWindow(QMainWindow):
|
||||
frame_cache[i] = img
|
||||
return img
|
||||
|
||||
def emit(i, restored):
|
||||
imwrite_unicode(str(out_path(files[i])), restored)
|
||||
job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}")
|
||||
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(
|
||||
total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled
|
||||
end - start + 1,
|
||||
lambda li: get_frame(start + li),
|
||||
lambda _li: [],
|
||||
lambda li, res: emit(start + li, res),
|
||||
should_cancel=lambda: job.cancelled,
|
||||
)
|
||||
else:
|
||||
for i, p in enumerate(files):
|
||||
indices = hits if only_detected else range(len(files))
|
||||
for i in indices:
|
||||
if job.cancelled:
|
||||
break
|
||||
p = files[i]
|
||||
if not force and out_path(p).is_file():
|
||||
job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}")
|
||||
emit(i, None, verb="Пропуск") # already restored — count, don't rewrite
|
||||
continue
|
||||
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
||||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||||
return None
|
||||
|
||||
def done(_result, cancelled):
|
||||
if self._current is not None: # live-preview the current frame's result, if any
|
||||
rp = out_path(self._current)
|
||||
if rp.is_file():
|
||||
img = imread_unicode(str(rp))
|
||||
if img is not None:
|
||||
self._restored[str(self._current)] = img
|
||||
self._showing_restored = True
|
||||
self.view.set_image(img, [])
|
||||
self._update_restore_actions()
|
||||
self._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} кадров)"
|
||||
else f"Готово: результаты в {out_dir.name}/ ({total} кадров) — показываю расцензуренное"
|
||||
)
|
||||
|
||||
self.statusBar().showMessage("Пакетное расцензуривание…")
|
||||
scope = " (по детекции)" if only_detected else ""
|
||||
self.statusBar().showMessage(f"Пакетное расцензуривание{scope}…")
|
||||
self._start_job(fn, total, on_done=done)
|
||||
|
||||
def _make_restorer(self):
|
||||
@@ -1133,30 +1374,58 @@ class MainWindow(QMainWindow):
|
||||
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:
|
||||
if self._current is None or str(self._current) not in self._restored:
|
||||
return
|
||||
"""Flip the global view mode between original and restored, then re-show."""
|
||||
self._showing_restored = not self._showing_restored
|
||||
key = str(self._current)
|
||||
if self._showing_restored:
|
||||
self.view.set_image(self._restored[key], [])
|
||||
if self._current is not None:
|
||||
self._show(self._current)
|
||||
else:
|
||||
self.view.set_image(imread_unicode(key), self._results.get(key) or [])
|
||||
self._update_restore_actions()
|
||||
self._update_restore_actions()
|
||||
self.statusBar().showMessage(
|
||||
"Показ: расцензуренное (где есть)" if self._showing_restored else "Показ: оригинал"
|
||||
)
|
||||
|
||||
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)
|
||||
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 "Показать результат"
|
||||
"Показать оригинал" if self._showing_restored else "Показать расцензуренное"
|
||||
)
|
||||
self.save_restored_action.setEnabled(
|
||||
self._current is not None and self._has_restored(self._current)
|
||||
)
|
||||
self.save_restored_action.setEnabled(has)
|
||||
|
||||
def _save_restored(self) -> None:
|
||||
if self._current is None or str(self._current) not in self._restored:
|
||||
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), self._restored[str(self._current)]):
|
||||
if imwrite_unicode(str(out), restored):
|
||||
self.statusBar().showMessage(f"Сохранено: {out}")
|
||||
else:
|
||||
QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.")
|
||||
@@ -1227,24 +1496,77 @@ class MainWindow(QMainWindow):
|
||||
_TINT_HIT = QColor(200, 80, 80, 70)
|
||||
_TINT_CLEAN = QColor(90, 160, 90, 50)
|
||||
|
||||
def _set_row_tag(self, item: QListWidgetItem, count: int) -> None:
|
||||
base = item.data(Qt.UserRole + 1)
|
||||
item.setText(f"{base} · {count}" if count else f"{base} · —")
|
||||
item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN)
|
||||
def _relabel_row(self, item: QListWidgetItem) -> None:
|
||||
"""Set a row's text/tint/tooltip from its detection + restoration state.
|
||||
|
||||
def _tag_file(self, path: Path, count: int) -> None:
|
||||
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._set_row_tag(item, count)
|
||||
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
|
||||
self._cfg.detector_models, d.yolo_conf, d.yolo_imgsz, nms_iou=nms_iou
|
||||
)
|
||||
|
||||
def _save_results(self) -> None:
|
||||
@@ -1265,20 +1587,16 @@ class MainWindow(QMainWindow):
|
||||
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]))
|
||||
self._relabel_all()
|
||||
return len(cached)
|
||||
|
||||
def _clear_results(self) -> None:
|
||||
"""Drop all cached detections and reset row labels/tints (keeps the detector)."""
|
||||
"""Drop all cached detections and reset row labels/tints (keeps the detector).
|
||||
|
||||
Restored ✓ markers stay — restoration is independent of detection.
|
||||
"""
|
||||
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._relabel_all()
|
||||
self._refresh_marks()
|
||||
|
||||
def _refresh_marks(self) -> None:
|
||||
@@ -1313,13 +1631,17 @@ class MainWindow(QMainWindow):
|
||||
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()):
|
||||
fp = Path(self.file_list.item(i).data(Qt.UserRole))
|
||||
if fp.stem in stems or str(fp) in mem:
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user