Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.

2026-06-06 15:11:53 +03:00
parent 10bf87aa47
commit 33f20fe681
28 changed files with 2256 additions and 0 deletions
@@ -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)