"""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"{path.name} · {w}×{h} · не рассчитано " "(двойной клик по файлу или «Рассчитать кадр»)" ) self.detail_table.setRowCount(0) self.view.set_highlight(None) return by_type: dict[str, int] = {} for d in dets: by_type[d.type.value] = by_type.get(d.type.value, 0) + 1 summary = ", ".join(f"{k}: {v}" for k, v in sorted(by_type.items())) or "ничего не найдено" self.detail_header.setText(f"{path.name} · {w}×{h} · всего {len(dets)} ({summary})") self.detail_table.blockSignals(True) self.detail_table.setRowCount(len(dets)) for row, d in enumerate(dets): x, y, bw, bh = d.bbox cells = [d.type.value, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))] for col, text in enumerate(cells): self.detail_table.setItem(row, col, QTableWidgetItem(text)) self.detail_table.blockSignals(False) self.detail_table.clearSelection() self.detail_table.resizeColumnsToContents() self.view.set_highlight(None) def _on_detail_selected(self) -> None: rows = self.detail_table.selectionModel().selectedRows() self.view.set_highlight(rows[0].row() if rows else None) def _on_threshold_changed(self, value: float) -> None: self._cfg.default_threshold = value self.view.set_threshold(value) settings_store.save(self._cfg)