Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.
This commit is contained in:
@@ -0,0 +1,62 @@
|
||||
"""Options dialog for "Создать из ролика…" — sampling mode, step, downscale.
|
||||
|
||||
Keeps the speed levers in one place: keyframe-only (fast) vs every-Nth-frame, the
|
||||
step, and an optional max-side downscale (smaller files → less disk/AV pressure).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from PySide6.QtWidgets import (
|
||||
QComboBox,
|
||||
QDialog,
|
||||
QDialogButtonBox,
|
||||
QFormLayout,
|
||||
QLabel,
|
||||
QSpinBox,
|
||||
QWidget,
|
||||
)
|
||||
|
||||
|
||||
class ExtractDialog(QDialog):
|
||||
def __init__(self, parent: QWidget | None = None) -> None:
|
||||
super().__init__(parent)
|
||||
self.setWindowTitle("Создать из ролика")
|
||||
|
||||
self.mode = QComboBox()
|
||||
self.mode.addItem("Только ключевые кадры (быстро)", userData=True)
|
||||
self.mode.addItem("Каждый N-й кадр", userData=False)
|
||||
self.mode.currentIndexChanged.connect(self._sync)
|
||||
|
||||
self.step = QSpinBox()
|
||||
self.step.setRange(1, 100000)
|
||||
self.step.setValue(15)
|
||||
|
||||
self.max_dim = QSpinBox()
|
||||
self.max_dim.setRange(0, 8192)
|
||||
self.max_dim.setSingleStep(120)
|
||||
self.max_dim.setValue(0)
|
||||
self.max_dim.setSpecialValueText("оригинал")
|
||||
|
||||
form = QFormLayout(self)
|
||||
form.addRow("Режим:", self.mode)
|
||||
form.addRow("Брать каждый N-й кадр:", self.step)
|
||||
form.addRow("Макс. сторона, px:", self.max_dim)
|
||||
hint = QLabel(
|
||||
"Ключевые кадры — в разы быстрее (декодируются только I-кадры),\n"
|
||||
"но реже по времени. Даунскейл уменьшает файлы и нагрузку на диск."
|
||||
)
|
||||
hint.setWordWrap(True)
|
||||
form.addRow(hint)
|
||||
|
||||
buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
|
||||
buttons.accepted.connect(self.accept)
|
||||
buttons.rejected.connect(self.reject)
|
||||
form.addRow(buttons)
|
||||
self._sync()
|
||||
|
||||
def _sync(self) -> None:
|
||||
self.step.setEnabled(not self.mode.currentData())
|
||||
|
||||
def options(self) -> tuple[bool, int, int]:
|
||||
"""Return (keyframes_only, step, max_dim)."""
|
||||
return bool(self.mode.currentData()), self.step.value(), self.max_dim.value()
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Widget that renders an image and draws detection overlays.
|
||||
|
||||
Overlay visibility and the confidence threshold are applied at paint time, so
|
||||
toggling them is instant. One detection can be *highlighted* (selected in the
|
||||
detail table) — it is drawn boldly even if below the threshold, while the others
|
||||
dim, so the user can inspect exactly what the detector found.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PySide6.QtCore import QPointF, QRectF, Qt
|
||||
from PySide6.QtGui import QBrush, QColor, QFont, QImage, QPainter, QPen, QPolygonF
|
||||
from PySide6.QtWidgets import QWidget
|
||||
|
||||
from ..config import OverlayConfig
|
||||
from ..core.detection.types import CensorType, Detection
|
||||
|
||||
|
||||
class ImageView(QWidget):
|
||||
def __init__(self, overlay_cfg: OverlayConfig, parent: QWidget | None = None) -> None:
|
||||
super().__init__(parent)
|
||||
self._cfg = overlay_cfg
|
||||
self._qimage: QImage | None = None
|
||||
self._dets: list[Detection] = []
|
||||
self._overlay_enabled = True
|
||||
self._threshold = 0.0
|
||||
self._highlight: int | None = None
|
||||
self.setMinimumSize(480, 360)
|
||||
|
||||
# ------------------------------------------------------------------ slots
|
||||
def set_image(self, image_bgr: np.ndarray | None, dets: list[Detection]) -> None:
|
||||
if image_bgr is None:
|
||||
self._qimage = None
|
||||
else:
|
||||
rgb = np.ascontiguousarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
|
||||
h, w, ch = rgb.shape
|
||||
# .copy() so the QImage owns its pixels (the numpy buffer can be freed).
|
||||
self._qimage = QImage(rgb.data, w, h, ch * w, QImage.Format_RGB888).copy()
|
||||
self._dets = dets
|
||||
self._highlight = None
|
||||
self.update()
|
||||
|
||||
def set_overlay_enabled(self, enabled: bool) -> None:
|
||||
self._overlay_enabled = enabled
|
||||
self.update()
|
||||
|
||||
def set_threshold(self, threshold: float) -> None:
|
||||
self._threshold = threshold
|
||||
self.update()
|
||||
|
||||
def set_highlight(self, index: int | None) -> None:
|
||||
self._highlight = index
|
||||
self.update()
|
||||
|
||||
# ------------------------------------------------------------------ paint
|
||||
def _color(self, ctype: CensorType) -> QColor:
|
||||
r, g, b = self._cfg.colors.get(ctype.value, (255, 0, 0))
|
||||
return QColor(r, g, b)
|
||||
|
||||
def paintEvent(self, event) -> None: # noqa: N802 - Qt signature
|
||||
painter = QPainter(self)
|
||||
painter.fillRect(self.rect(), QColor(18, 18, 18))
|
||||
|
||||
if self._qimage is None:
|
||||
painter.setPen(QColor(160, 160, 160))
|
||||
painter.drawText(self.rect(), Qt.AlignCenter, "Откройте папку с картинками (Файл → Открыть папку…)")
|
||||
painter.end()
|
||||
return
|
||||
|
||||
iw, ih = self._qimage.width(), self._qimage.height()
|
||||
scale = min(self.width() / iw, self.height() / ih)
|
||||
dw, dh = iw * scale, ih * scale
|
||||
ox, oy = (self.width() - dw) / 2, (self.height() - dh) / 2
|
||||
|
||||
painter.setRenderHint(QPainter.SmoothPixmapTransform, True)
|
||||
painter.drawImage(QRectF(ox, oy, dw, dh), self._qimage)
|
||||
|
||||
if self._overlay_enabled and self._dets:
|
||||
painter.setRenderHint(QPainter.Antialiasing, True)
|
||||
for i, d in enumerate(self._dets):
|
||||
highlighted = i == self._highlight
|
||||
# A highlighted detection is always drawn; others respect the threshold.
|
||||
if not highlighted and d.score < self._threshold:
|
||||
continue
|
||||
dim = self._highlight is not None and not highlighted
|
||||
self._draw_detection(painter, d, ox, oy, scale, highlighted, dim)
|
||||
painter.end()
|
||||
|
||||
def _draw_detection(
|
||||
self, painter: QPainter, d: Detection, ox: float, oy: float,
|
||||
scale: float, highlighted: bool, dim: bool,
|
||||
) -> None:
|
||||
color = self._color(d.type)
|
||||
width = self._cfg.line_width * (2 if highlighted else 1)
|
||||
pen_color = QColor(color)
|
||||
if dim:
|
||||
pen_color.setAlpha(70)
|
||||
painter.setPen(QPen(pen_color, width))
|
||||
fill = QColor(color)
|
||||
fill.setAlpha(0 if dim else (self._cfg.fill_alpha * 2 if highlighted else self._cfg.fill_alpha))
|
||||
painter.setBrush(QBrush(fill))
|
||||
|
||||
points = d.polygon if len(d.polygon) >= 3 else self._bbox_points(d.bbox)
|
||||
poly = QPolygonF([QPointF(ox + x * scale, oy + y * scale) for x, y in points])
|
||||
painter.drawPolygon(poly)
|
||||
|
||||
if self._cfg.show_labels and not dim:
|
||||
x, y, _w, _h = d.bbox
|
||||
self._draw_label(painter, f"{d.type.value} {d.score:.2f}", ox + x * scale, oy + y * scale, color)
|
||||
|
||||
@staticmethod
|
||||
def _bbox_points(bbox: tuple[int, int, int, int]) -> list[tuple[int, int]]:
|
||||
x, y, w, h = bbox
|
||||
return [(x, y), (x + w, y), (x + w, y + h), (x, y + h)]
|
||||
|
||||
def _draw_label(self, painter: QPainter, text: str, x: float, y: float, color: QColor) -> None:
|
||||
font = QFont()
|
||||
font.setPointSize(9)
|
||||
painter.setFont(font)
|
||||
metrics = painter.fontMetrics()
|
||||
tw = metrics.horizontalAdvance(text) + 8
|
||||
th = metrics.height() + 2
|
||||
bg = QRectF(x, max(0.0, y - th), tw, th)
|
||||
painter.fillRect(bg, color)
|
||||
painter.setPen(QColor(0, 0, 0))
|
||||
painter.drawText(bg, Qt.AlignCenter, text)
|
||||
@@ -0,0 +1,494 @@
|
||||
"""Main window: open a folder of images and inspect what the detector found.
|
||||
|
||||
Layout: a toolbar (open folder · detector · model · calc-frame · detect-all ·
|
||||
threshold), then a splitter with three panes — the file list (left), the image
|
||||
with overlays (center), and a detail table of every detection (right).
|
||||
|
||||
Viewing and detecting are decoupled, so browsing a big folder stays instant even
|
||||
with a slow (CPU) detector:
|
||||
- selecting a file just **shows** it (with its cached result, if any);
|
||||
- **double-clicking** a file, or "Рассчитать кадр", runs the detector on it;
|
||||
- "Детектировать все" runs the whole folder.
|
||||
Both folder loading and detect-all show a progress bar. Results are cached;
|
||||
switching detector/model clears the cache.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from PySide6.QtCore import Qt
|
||||
from PySide6.QtGui import QAction
|
||||
from PySide6.QtWidgets import (
|
||||
QAbstractItemView,
|
||||
QApplication,
|
||||
QComboBox,
|
||||
QDialog,
|
||||
QDoubleSpinBox,
|
||||
QFileDialog,
|
||||
QInputDialog,
|
||||
QLabel,
|
||||
QListWidget,
|
||||
QListWidgetItem,
|
||||
QMainWindow,
|
||||
QMessageBox,
|
||||
QProgressBar,
|
||||
QSplitter,
|
||||
QTableWidget,
|
||||
QTableWidgetItem,
|
||||
QVBoxLayout,
|
||||
QWidget,
|
||||
)
|
||||
|
||||
from .. import settings_store
|
||||
from ..config import AppConfig
|
||||
from ..core.detection.factory import build_detector
|
||||
from ..core.detection.types import Detection
|
||||
from ..core.imageio import imread_unicode
|
||||
from ..core.video.extract import extract_frames
|
||||
from ..core.video.frame import Frame
|
||||
from .extract_dialog import ExtractDialog
|
||||
from .image_view import ImageView
|
||||
|
||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||||
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
|
||||
_DETECTORS = ["classic", "yolo", "combined"]
|
||||
|
||||
|
||||
class MainWindow(QMainWindow):
|
||||
def __init__(self, config: AppConfig) -> None:
|
||||
super().__init__()
|
||||
self._cfg = config
|
||||
self._detector = None
|
||||
self._detector_key = None
|
||||
self._folder: Path | None = None
|
||||
self._files: list[Path] = []
|
||||
self._results: dict[str, list[Detection]] = {} # path -> detections (cache)
|
||||
self._current: Path | None = None
|
||||
self._collection: Path | None = None # active destination folder for moves
|
||||
|
||||
self.setWindowTitle("HVideoTool — инспектор детекции цензуры")
|
||||
self.resize(1180, 720)
|
||||
|
||||
self._build_toolbar()
|
||||
self._build_central()
|
||||
self._build_statusbar()
|
||||
self._build_menu()
|
||||
self.statusBar().showMessage("Откройте папку с картинками")
|
||||
|
||||
# ------------------------------------------------------------------ setup
|
||||
def _build_menu(self) -> None:
|
||||
file_menu = self.menuBar().addMenu("Файл")
|
||||
file_menu.addAction("Открыть папку…", self._choose_folder)
|
||||
file_menu.addAction("Создать из ролика…", self._create_from_video)
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Рассчитать кадр", self._recompute_current).setShortcut("Space")
|
||||
file_menu.addAction("Детектировать все", self._detect_all)
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Создать коллекцию…", self._create_collection)
|
||||
file_menu.addAction("В коллекцию", self._move_to_collection).setShortcut("Ctrl+M")
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Выход", self.close)
|
||||
|
||||
def _build_toolbar(self) -> None:
|
||||
tb = self.addToolBar("Главная")
|
||||
tb.setMovable(False)
|
||||
|
||||
tb.addAction(QAction("Открыть папку…", self, triggered=self._choose_folder))
|
||||
from_video = QAction("Создать из ролика…", self, triggered=self._create_from_video)
|
||||
from_video.setToolTip("Разложить видео на кадры в папку-коллекцию и открыть её")
|
||||
tb.addAction(from_video)
|
||||
tb.addSeparator()
|
||||
|
||||
tb.addWidget(QLabel(" Детектор: "))
|
||||
self.detector_combo = QComboBox()
|
||||
self.detector_combo.addItems(_DETECTORS)
|
||||
self.detector_combo.setCurrentText(self._cfg.detector)
|
||||
self.detector_combo.currentTextChanged.connect(self._on_detector_changed)
|
||||
tb.addWidget(self.detector_combo)
|
||||
|
||||
self.model_action = QAction("Модель…", self, triggered=self._choose_model)
|
||||
tb.addAction(self.model_action)
|
||||
|
||||
tb.addSeparator()
|
||||
calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current)
|
||||
calc.setToolTip("Запустить детектор на выбранном кадре (Space / двойной клик по файлу)")
|
||||
tb.addAction(calc)
|
||||
tb.addAction(QAction("Детектировать все", self, triggered=self._detect_all))
|
||||
|
||||
tb.addSeparator()
|
||||
tb.addWidget(QLabel(" Порог: "))
|
||||
self.threshold_spin = QDoubleSpinBox()
|
||||
self.threshold_spin.setRange(0.0, 1.0)
|
||||
self.threshold_spin.setSingleStep(0.05)
|
||||
self.threshold_spin.setValue(self._cfg.default_threshold)
|
||||
self.threshold_spin.valueChanged.connect(self._on_threshold_changed)
|
||||
tb.addWidget(self.threshold_spin)
|
||||
|
||||
tb.addSeparator()
|
||||
tb.addAction(QAction("Создать коллекцию…", self, triggered=self._create_collection))
|
||||
move = QAction("В коллекцию →", self, triggered=self._move_to_collection)
|
||||
move.setToolTip("Переместить выбранные кадры в активную коллекцию (Ctrl+M)")
|
||||
tb.addAction(move)
|
||||
self.collection_label = QLabel(" коллекция: —")
|
||||
tb.addWidget(self.collection_label)
|
||||
|
||||
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)
|
||||
|
||||
self.view = ImageView(self._cfg.overlay)
|
||||
self.view.set_threshold(self._cfg.default_threshold)
|
||||
|
||||
right = QWidget()
|
||||
rlayout = QVBoxLayout(right)
|
||||
rlayout.setContentsMargins(4, 4, 4, 4)
|
||||
self.detail_header = QLabel("Детекции")
|
||||
self.detail_header.setWordWrap(True)
|
||||
rlayout.addWidget(self.detail_header)
|
||||
self.detail_table = QTableWidget(0, 4)
|
||||
self.detail_table.setHorizontalHeaderLabels(["Тип", "Увер.", "BBox (x,y,w,h)", "Полигон"])
|
||||
self.detail_table.verticalHeader().setVisible(False)
|
||||
self.detail_table.setSelectionBehavior(QTableWidget.SelectRows)
|
||||
self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers)
|
||||
self.detail_table.itemSelectionChanged.connect(self._on_detail_selected)
|
||||
rlayout.addWidget(self.detail_table)
|
||||
|
||||
splitter = QSplitter(Qt.Horizontal)
|
||||
splitter.addWidget(self.file_list)
|
||||
splitter.addWidget(self.view)
|
||||
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_statusbar(self) -> None:
|
||||
self.progress = QProgressBar()
|
||||
self.progress.setMaximumWidth(260)
|
||||
self.progress.setVisible(False)
|
||||
self.statusBar().addPermanentWidget(self.progress)
|
||||
|
||||
# --------------------------------------------------------------- detector
|
||||
def _make_detector(self):
|
||||
d = self._cfg.detection
|
||||
key = (self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz)
|
||||
if key != self._detector_key:
|
||||
self._detector = build_detector(self._cfg) # may raise ValueError / import / file errors
|
||||
self._detector_key = key
|
||||
return self._detector
|
||||
|
||||
def _on_detector_changed(self, name: str) -> None:
|
||||
self._cfg.detector = name
|
||||
# YOLO/combined need a model — offer to pick one if missing.
|
||||
if name in ("yolo", "combined") and not self._cfg.model_path:
|
||||
self._choose_model()
|
||||
settings_store.save(self._cfg)
|
||||
self._invalidate_results()
|
||||
|
||||
def _choose_model(self) -> None:
|
||||
start = self._cfg.model_path or str(Path.cwd() / "models")
|
||||
path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)")
|
||||
if path:
|
||||
self._cfg.model_path = path
|
||||
settings_store.save(self._cfg)
|
||||
self.statusBar().showMessage(f"Модель: {path}")
|
||||
self._invalidate_results()
|
||||
|
||||
def _invalidate_results(self) -> None:
|
||||
"""Detector changed — drop the cache and refresh the current image."""
|
||||
self._detector_key = None
|
||||
self._results.clear()
|
||||
for i in range(self.file_list.count()):
|
||||
self.file_list.item(i).setText(self.file_list.item(i).data(Qt.UserRole + 1))
|
||||
if self._current is not None:
|
||||
self._show(self._current)
|
||||
|
||||
# --------------------------------------------------------------- handlers
|
||||
def open_path(self, folder: str) -> None:
|
||||
self._load_folder(Path(folder))
|
||||
|
||||
def _choose_folder(self) -> None:
|
||||
start = settings_store.last_dir() or ""
|
||||
folder = QFileDialog.getExistingDirectory(self, "Открыть папку с картинками", start)
|
||||
if folder:
|
||||
self._load_folder(Path(folder))
|
||||
|
||||
def _create_from_video(self) -> None:
|
||||
"""Decode a video into a folder of frames (a collection) and open it."""
|
||||
path, _ = QFileDialog.getOpenFileName(
|
||||
self, "Выберите ролик", settings_store.last_dir() or "", _VIDEO_FILTER
|
||||
)
|
||||
if not path:
|
||||
return
|
||||
dialog = ExtractDialog(self)
|
||||
if dialog.exec() != QDialog.Accepted:
|
||||
return
|
||||
keyframes_only, step, max_dim = dialog.options()
|
||||
video = Path(path)
|
||||
out = video.parent / f"{video.stem}_frames"
|
||||
|
||||
self.progress.setRange(0, 1000) # promille of duration
|
||||
self.progress.setValue(0)
|
||||
self.progress.setVisible(True)
|
||||
|
||||
def cb(done: float, total: float) -> bool:
|
||||
if total > 0:
|
||||
self.progress.setValue(int(1000 * min(done, total) / total))
|
||||
self.statusBar().showMessage(f"Извлечение кадров: {done:.0f}/{total:.0f} с…")
|
||||
QApplication.processEvents()
|
||||
return True
|
||||
|
||||
try:
|
||||
saved = extract_frames(
|
||||
str(video), str(out), step=step, keyframes_only=keyframes_only,
|
||||
max_dim=max_dim, progress=cb,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - surface decode errors to the user
|
||||
self.progress.setVisible(False)
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}")
|
||||
return
|
||||
finally:
|
||||
self.progress.setVisible(False)
|
||||
|
||||
if saved == 0:
|
||||
QMessageBox.warning(self, "Пусто", "Из ролика не удалось извлечь ни одного кадра.")
|
||||
return
|
||||
self.statusBar().showMessage(f"Извлечено {saved} кадров → {out}")
|
||||
self._load_folder(out)
|
||||
|
||||
def _load_folder(self, folder: Path) -> None:
|
||||
if not folder.is_dir():
|
||||
QMessageBox.warning(self, "Ошибка", f"Папка не найдена: {folder}")
|
||||
return
|
||||
self.statusBar().showMessage(f"Сканирую папку: {folder}…")
|
||||
QApplication.processEvents()
|
||||
files = sorted(p for p in folder.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
|
||||
self._folder = folder
|
||||
self._files = files
|
||||
self._results.clear()
|
||||
self._current = None
|
||||
settings_store.set_last_dir(str(folder))
|
||||
|
||||
self.file_list.blockSignals(True)
|
||||
self.file_list.setUpdatesEnabled(False)
|
||||
self.file_list.clear()
|
||||
self.progress.setRange(0, len(files))
|
||||
self.progress.setVisible(True)
|
||||
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()
|
||||
self.file_list.setUpdatesEnabled(True)
|
||||
self.file_list.blockSignals(False)
|
||||
self.progress.setVisible(False)
|
||||
|
||||
if not files:
|
||||
self.view.set_image(None, [])
|
||||
self.statusBar().showMessage(f"В папке нет картинок: {folder}")
|
||||
return
|
||||
self.statusBar().showMessage(f"{len(files)} картинок · {folder} · двойной клик / «Рассчитать кадр» для детекции")
|
||||
self.file_list.setCurrentRow(0)
|
||||
|
||||
def _on_file_selected(self, current: QListWidgetItem | None, _prev=None) -> None:
|
||||
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.
|
||||
path = Path(item.data(Qt.UserRole))
|
||||
if str(path) not in self._results and self._detect(path) is None:
|
||||
return
|
||||
self._show(path)
|
||||
|
||||
def _detect(self, path: Path) -> list[Detection] | None:
|
||||
"""Run (or fetch cached) detections for one image. None on failure."""
|
||||
key = str(path)
|
||||
if key in self._results:
|
||||
return self._results[key]
|
||||
img = imread_unicode(key)
|
||||
if img is None:
|
||||
self.statusBar().showMessage(f"Не удалось прочитать: {path.name}")
|
||||
return None
|
||||
try:
|
||||
detector = self._make_detector()
|
||||
except Exception as exc: # noqa: BLE001 - surface config/model errors to the user
|
||||
QMessageBox.warning(self, "Детектор недоступен", str(exc))
|
||||
return None
|
||||
self.statusBar().showMessage(f"Детекция: {path.name}…")
|
||||
QApplication.processEvents()
|
||||
dets = detector.detect(Frame(image=img, index=0, pts=0.0))
|
||||
dets.sort(key=lambda d: d.score, reverse=True)
|
||||
self._results[key] = dets
|
||||
self._tag_file(path, len(dets))
|
||||
return dets
|
||||
|
||||
def _show(self, path: Path) -> None:
|
||||
"""Display the image with its cached detections (does not run the detector)."""
|
||||
self._current = path
|
||||
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)
|
||||
|
||||
def _recompute_current(self) -> None:
|
||||
"""Toolbar/Space: (re)run the detector on the selected frame."""
|
||||
if self._current is None:
|
||||
return
|
||||
self._results.pop(str(self._current), None)
|
||||
self._detector_key = None # rebuild the detector so settings changes take effect
|
||||
if self._detect(self._current) is None:
|
||||
return
|
||||
self._show(self._current)
|
||||
|
||||
def _detect_all(self) -> None:
|
||||
if not self._files:
|
||||
return
|
||||
total = len(self._files)
|
||||
self.progress.setRange(0, total)
|
||||
self.progress.setVisible(True)
|
||||
try:
|
||||
for i, p in enumerate(self._files, 1):
|
||||
self.progress.setValue(i)
|
||||
self.statusBar().showMessage(f"Детекция {i}/{total}: {p.name}")
|
||||
QApplication.processEvents()
|
||||
if self._detect(p) is None:
|
||||
return # detector unavailable — message already shown
|
||||
finally:
|
||||
self.progress.setVisible(False)
|
||||
hits = sum(1 for p in self._files if self._results.get(str(p)))
|
||||
self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок")
|
||||
if self._current is not None:
|
||||
self._show(self._current)
|
||||
|
||||
# ------------------------------------------------------------ collections
|
||||
def _collections_base(self) -> Path:
|
||||
"""Where new collections are created: next to the opened folder, else home."""
|
||||
if self._folder is not None:
|
||||
return self._folder.parent
|
||||
return Path.home() / "HVideoTool" / "collections"
|
||||
|
||||
def _create_collection(self) -> None:
|
||||
name, ok = QInputDialog.getText(self, "Создать коллекцию", "Имя коллекции:")
|
||||
name = name.strip()
|
||||
if not ok or not name:
|
||||
return
|
||||
path = self._collections_base() / name
|
||||
try:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as exc:
|
||||
QMessageBox.warning(self, "Ошибка", f"Не удалось создать коллекцию:\n{exc}")
|
||||
return
|
||||
self._collection = path
|
||||
self._update_collection_label()
|
||||
self.statusBar().showMessage(f"Активная коллекция: {path}")
|
||||
|
||||
def _update_collection_label(self) -> None:
|
||||
self.collection_label.setText(
|
||||
f" коллекция: {self._collection.name}" if self._collection else " коллекция: —"
|
||||
)
|
||||
|
||||
def _move_to_collection(self) -> None:
|
||||
if self._collection is None:
|
||||
QMessageBox.information(
|
||||
self, "Нет коллекции",
|
||||
"Сначала создайте коллекцию (кнопка «Создать коллекцию…»).",
|
||||
)
|
||||
return
|
||||
items = self.file_list.selectedItems()
|
||||
if not items:
|
||||
QMessageBox.information(self, "Нет выбора", "Выберите кадры в списке слева.")
|
||||
return
|
||||
|
||||
moved = 0
|
||||
for item in items:
|
||||
src = Path(item.data(Qt.UserRole))
|
||||
if not src.exists():
|
||||
continue
|
||||
dst = self._unique_dest(self._collection, src.name)
|
||||
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
|
||||
|
||||
self.statusBar().showMessage(f"Перемещено {moved} → {self._collection.name}")
|
||||
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, [])
|
||||
|
||||
@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
|
||||
def _tag_file(self, path: Path, count: int) -> None:
|
||||
for i in range(self.file_list.count()):
|
||||
item = self.file_list.item(i)
|
||||
if item.data(Qt.UserRole) == str(path):
|
||||
base = item.data(Qt.UserRole + 1)
|
||||
item.setText(f"{base} · {count}" if count else f"{base} · —")
|
||||
return
|
||||
|
||||
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.type.value] = by_type.get(d.type.value, 0) + 1
|
||||
summary = ", ".join(f"{k}: {v}" for k, v in sorted(by_type.items())) or "ничего не найдено"
|
||||
self.detail_header.setText(f"<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.type.value, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
|
||||
for col, text in enumerate(cells):
|
||||
self.detail_table.setItem(row, col, QTableWidgetItem(text))
|
||||
self.detail_table.blockSignals(False)
|
||||
self.detail_table.clearSelection()
|
||||
self.detail_table.resizeColumnsToContents()
|
||||
self.view.set_highlight(None)
|
||||
|
||||
def _on_detail_selected(self) -> None:
|
||||
rows = self.detail_table.selectionModel().selectedRows()
|
||||
self.view.set_highlight(rows[0].row() if rows else None)
|
||||
|
||||
def _on_threshold_changed(self, value: float) -> None:
|
||||
self._cfg.default_threshold = value
|
||||
self.view.set_threshold(value)
|
||||
settings_store.save(self._cfg)
|
||||
Reference in New Issue
Block a user