diff --git a/CLAUDE.md b/CLAUDE.md index 7685e1a..0de80f3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -38,14 +38,19 @@ a detector, draws the regions, and shows a detailed per-image list of what it fo Keep this scope sharp: -- It is a **detection + overlay/inspection** tool. It does **not** remove, restore, or - reconstruct censored content. -- It does **not** generate images. There is **no** ControlNet / SDXL / diffusion - pipeline. (`xinsir/controlnet-union-sdxl-1.0` was considered early but rejected — a - generative model, not a detector. Do not reintroduce it.) +- Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить") + step was added later (user-requested): per-frame, on-demand, behind a `Restorer` + interface. The shipped engine is a cv2 **inpaint baseline** (fills, does NOT truly + reconstruct); a generative engine (DeepMosaics / LADA BasicVSR++) is the intended + real engine but needs weights + a CUDA GPU and is not wired yet. +- Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was + rejected early — a generative *conditioning* model, not a censorship restorer. Don't + reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/ + LADA), not a general text-to-image diffusion pipeline. - It detects **already-censored** regions, not "content that should be censored" (i.e. not an NSFW classifier). -- It does **not** decode video. No PyAV. Input is image files only. +- Video is only a one-shot frame-extraction convenience (see below); detection and + restoration operate on image folders. ## Target environment @@ -91,6 +96,11 @@ hvideotool/ ├── video/ │ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type + ├── restore/ # "un-censor" detected regions (per-frame) + │ ├── base.py # Restorer ABC: restore(image, detections) -> image + │ ├── factory.py # build_restorer(name) -> inpaint (deepmosaics/lada = not wired yet) + │ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs + │ └── mask.py # detections_to_mask(shape, dets, dilate) └── detection/ ├── base.py # Detector ABC: detect(frame) -> list[Detection] ├── factory.py # build_detector(config) -> classic | yolo | combined @@ -115,12 +125,32 @@ hvideotool/ "Детектировать все" (whole folder, progress bar). Do NOT re-add auto-detect-on-select. Results cache in `_results`; the file-list row gets a count suffix when computed. Switching detector/model clears the cache (`_invalidate_results`). -- **Collections (curation).** "Создать коллекцию…" makes a destination folder - (`_collections_base()` = the opened folder's parent, else `~/HVideoTool/collections`) - and marks it active. The file list is `ExtendedSelection`; "В коллекцию" / Ctrl+M - **moves** (`shutil.move`, not copy) the selected frames there, removing them from the +- **Collections (curation).** A combo in the left pane (next to the file list, since it + acts on the list selection — not on the toolbar, to keep that uncluttered) + (`collection_combo`) picks the active destination; `_refresh_collections()` repopulates it from sibling folders of the + opened folder (`_collections_base()` = the opened folder's parent, else + `~/HVideoTool/collections`) on load/create, so previously-made collections are + reselectable. Items carry the path in itemData; "— не выбрана —" = None, + "Выбрать папку…" = `"__browse__"` sentinel → `_browse_collection()` for an arbitrary + folder (kept in the combo even if outside base). "Создать…" makes a new one and + selects it. The file list is `ExtendedSelection`; "В коллекцию" / Ctrl+M **moves** + (`shutil.move`, not copy) the selected frames there, removing them from list/`_files`/cache. `_unique_dest` avoids clobbering (`foo.jpg` → `foo (1).jpg`). Use case: sort frames into a training/example set while inspecting detections. +- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via + `build_restorer`) on the current frame's detections (computing them first if needed), + caches the result in `_restored[path]`, and shows it overlay-free. "Показать оригинал/ + результат" toggles (`_showing_restored`); "Сохранить результат" writes + `_restored.jpg` into the active collection (or beside the frame). The baseline + is cv2 inpaint — honest: it fills, doesn't reconstruct. To add a real engine, + implement `core/restore/base.Restorer`, register it in `restore/factory.build_restorer`, + and swap `self._restorer`. `_show` resets `_showing_restored` + `_update_restore_actions`. +- **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys + `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` + ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, + `_step_hit` scans `_results` for the next non-empty frame). The slider and file list + are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all + navigation ultimately drives `file_list.setCurrentRow`. - `image_view.ImageView` draws the image scaled-to-fit plus overlays. Overlay visibility/threshold are applied at paint time. Selecting a row in the detail table calls `set_highlight(i)` — that detection is drawn boldly (even below threshold) and diff --git a/README.md b/README.md index bf44c2f..37ffa35 100644 --- a/README.md +++ b/README.md @@ -31,13 +31,25 @@ уверенности прямо в тулбаре — удобно сравнивать. - Выбор файла весов модели кнопкой **«Модель…»**. +## Восстановление (расцензуривание) + +Кнопка **«Расцензурить кадр»** восстанавливает найденные области на текущем кадре, +**«Показать оригинал/результат»** переключает вид, **«Сохранить результат»** пишет +`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром). + +> ⚠️ Сейчас движок восстановления — **инпейнт (cv2)**: он *заполняет* область по +> окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). +> Настоящее восстановление мозаики требует генеративной модели (DeepMosaics или +> [LADA](https://github.com/ladaapp/lada) на BasicVSR++) и **GPU NVIDIA/CUDA**; на +> аниме качество таких моделей ограничено. Движок подключается через интерфейс +> `core/restore/base.Restorer` (`build_restorer`) — инпейнт сейчас стоит заглушкой. + ## Что НЕ делает (осознанно вне области задачи) -- Не **удаляет** и не **восстанавливает** зацензуренный контент. -- Не **генерирует** изображения (никакого ControlNet/SDXL/диффузии). +- Не **генерирует** изображения через диффузию (никакого ControlNet/SDXL). - Не детектирует «контент, который следовало бы зацензурить» (NSFW) — ищем именно **уже наложенную** цензуру. -- Не декодирует видео — работает с готовыми картинками. +- Полноценное генеративное восстановление пока не подключено (см. выше). --- diff --git a/hvideotool/core/restore/__init__.py b/hvideotool/core/restore/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/hvideotool/core/restore/__init__.py @@ -0,0 +1 @@ + diff --git a/hvideotool/core/restore/base.py b/hvideotool/core/restore/base.py new file mode 100644 index 0000000..6526d5a --- /dev/null +++ b/hvideotool/core/restore/base.py @@ -0,0 +1,26 @@ +"""Restorer interface — "un-censor" detected regions of an image. + +A Restorer takes an image plus the detected censored regions and returns a new +image with those regions reconstructed/filled. This mirrors the ``Detector`` +abstraction so different engines (classic inpaint now; a generative model like +DeepMosaics / LADA later) plug in behind the same interface. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import numpy as np + +from ..detection.types import Detection + + +class Restorer(ABC): + @property + def name(self) -> str: + return type(self).__name__ + + @abstractmethod + def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray: + """Return a copy of ``image`` with the detected regions reconstructed.""" + raise NotImplementedError diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py new file mode 100644 index 0000000..1fc6c4f --- /dev/null +++ b/hvideotool/core/restore/factory.py @@ -0,0 +1,22 @@ +"""Restorer factory: build a Restorer by name. + +Currently only the cv2 inpaint baseline is wired. Generative engines +(DeepMosaics / LADA BasicVSR++) are placeholders — they need model weights and a +CUDA GPU, and raise a clear, actionable error until integrated. See README. +""" + +from __future__ import annotations + +from .base import Restorer +from .inpaint import InpaintRestorer + + +def build_restorer(name: str = "inpaint", model_path: str | None = None) -> Restorer: + if name == "inpaint": + return InpaintRestorer() + if name in ("deepmosaics", "lada"): + raise ValueError( + "Генеративное восстановление пока не подключено.\n" + "Нужна модель (DeepMosaics / LADA) и GPU (CUDA). См. README → Восстановление." + ) + raise ValueError(f"Неизвестный режим восстановления: {name!r}") diff --git a/hvideotool/core/restore/inpaint.py b/hvideotool/core/restore/inpaint.py new file mode 100644 index 0000000..25e992f --- /dev/null +++ b/hvideotool/core/restore/inpaint.py @@ -0,0 +1,35 @@ +"""Classic inpainting restorer (cv2) — the always-available baseline. + +HONEST LIMITATION: cv2 inpainting fills the masked region by propagating +surrounding pixels. It removes the mosaic/bar but does NOT reconstruct the hidden +detail — it smooths/guesses. For real reconstruction a generative model +(DeepMosaics / LADA) is needed; this is the no-weights, no-GPU fallback so the +"Расцензурить кадр" flow works end-to-end today. +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from ..detection.types import Detection +from .base import Restorer +from .mask import detections_to_mask + + +class InpaintRestorer(Restorer): + def __init__(self, radius: int = 3, dilate: int = 2, method: str = "telea") -> None: + self.radius = radius + self.dilate = dilate + self.method = method + + @property + def name(self) -> str: + return f"InpaintRestorer({self.method})" + + def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray: + if not detections: + return image.copy() + mask = detections_to_mask(image.shape, detections, dilate=self.dilate) + flags = cv2.INPAINT_TELEA if self.method == "telea" else cv2.INPAINT_NS + return cv2.inpaint(image, mask, self.radius, flags) diff --git a/hvideotool/core/restore/mask.py b/hvideotool/core/restore/mask.py new file mode 100644 index 0000000..237549b --- /dev/null +++ b/hvideotool/core/restore/mask.py @@ -0,0 +1,26 @@ +"""Build a binary mask of the censored regions from detections.""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from ..detection.types import Detection + + +def detections_to_mask( + shape: tuple[int, int], detections: list[Detection], dilate: int = 0 +) -> np.ndarray: + """White (255) over every detected region (polygon if present, else bbox).""" + h, w = shape[:2] + mask = np.zeros((h, w), np.uint8) + for d in detections: + if len(d.polygon) >= 3: + cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255) + else: + x, y, bw, bh = d.bbox + cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1) + if dilate > 0: + k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8) + mask = cv2.dilate(mask, k) + return mask diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index 7e837c5..fec0550 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -1,8 +1,10 @@ """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). +Layout: a toolbar (open folder · from-video · detector · model · calc-frame · +detect-all · threshold), then a splitter with three panes — left: collection +controls + the file list; center: the image with overlays; right: a detail table +of every detection. Collection controls sit by the file list (they act on its +selection), keeping the toolbar to detection/entry actions only. Viewing and detecting are decoupled, so browsing a big folder stays instant even with a slow (CPU) detector: @@ -19,7 +21,7 @@ import shutil from pathlib import Path from PySide6.QtCore import Qt -from PySide6.QtGui import QAction +from PySide6.QtGui import QAction, QKeySequence, QShortcut from PySide6.QtWidgets import ( QAbstractItemView, QApplication, @@ -27,6 +29,7 @@ from PySide6.QtWidgets import ( QDialog, QDoubleSpinBox, QFileDialog, + QHBoxLayout, QInputDialog, QLabel, QListWidget, @@ -34,6 +37,8 @@ from PySide6.QtWidgets import ( QMainWindow, QMessageBox, QProgressBar, + QPushButton, + QSlider, QSplitter, QTableWidget, QTableWidgetItem, @@ -45,7 +50,8 @@ 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.imageio import imread_unicode, imwrite_unicode +from ..core.restore.factory import build_restorer from ..core.video.extract import extract_frames from ..core.video.frame import Frame from .extract_dialog import ExtractDialog @@ -67,6 +73,10 @@ class MainWindow(QMainWindow): 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._restorer = build_restorer("inpaint") # un-censor engine (baseline) + self._restored: dict[str, "object"] = {} # path -> restored image (BGR ndarray) + self._showing_restored = False + self._nav_sync = False # guard against slider<->list signal loops self.setWindowTitle("HVideoTool — инспектор детекции цензуры") self.resize(1180, 720) @@ -75,6 +85,7 @@ class MainWindow(QMainWindow): self._build_central() self._build_statusbar() self._build_menu() + self._refresh_collections() self.statusBar().showMessage("Откройте папку с картинками") # ------------------------------------------------------------------ setup @@ -117,6 +128,17 @@ class MainWindow(QMainWindow): tb.addAction(calc) tb.addAction(QAction("Детектировать все", self, triggered=self._detect_all)) + tb.addSeparator() + restore = QAction("Расцензурить кадр", self, triggered=self._restore_current) + restore.setToolTip("Восстановить найденные области на текущем кадре") + tb.addAction(restore) + 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() tb.addWidget(QLabel(" Порог: ")) self.threshold_spin = QDoubleSpinBox() @@ -126,22 +148,44 @@ class MainWindow(QMainWindow): 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) + # Collection controls live next to the file list — they act on its selection. + self.collection_combo = QComboBox() + self.collection_combo.setToolTip("Активная коллекция, куда перемещаются кадры") + self.collection_combo.activated.connect(self._on_collection_selected) + new_coll = QPushButton("Создать") + new_coll.clicked.connect(self._create_collection) + move_btn = QPushButton("В коллекцию →") + move_btn.setToolTip("Переместить выбранные кадры в активную коллекцию (Ctrl+M)") + move_btn.clicked.connect(self._move_to_collection) + + coll_row = QHBoxLayout() + coll_row.setContentsMargins(0, 0, 0, 0) + coll_row.addWidget(QLabel("Коллекция:")) + coll_row.addWidget(self.collection_combo, 1) + coll_row.addWidget(new_coll) + + left = QWidget() + left_layout = QVBoxLayout(left) + left_layout.setContentsMargins(4, 4, 4, 4) + left_layout.setSpacing(4) + left_layout.addLayout(coll_row) + left_layout.addWidget(self.file_list, 1) + left_layout.addWidget(move_btn) + self.view = ImageView(self._cfg.overlay) self.view.set_threshold(self._cfg.default_threshold) + center = QWidget() + clayout = QVBoxLayout(center) + clayout.setContentsMargins(0, 0, 0, 0) + clayout.setSpacing(2) + clayout.addWidget(self.view, 1) + clayout.addWidget(self._build_nav_bar()) right = QWidget() rlayout = QVBoxLayout(right) @@ -158,8 +202,8 @@ class MainWindow(QMainWindow): rlayout.addWidget(self.detail_table) splitter = QSplitter(Qt.Horizontal) - splitter.addWidget(self.file_list) - splitter.addWidget(self.view) + splitter.addWidget(left) + splitter.addWidget(center) splitter.addWidget(right) splitter.setStretchFactor(0, 0) splitter.setStretchFactor(1, 1) @@ -167,6 +211,87 @@ class MainWindow(QMainWindow): splitter.setSizes([240, 640, 300]) self.setCentralWidget(splitter) + def _build_nav_bar(self) -> QWidget: + bar = QWidget() + h = QHBoxLayout(bar) + h.setContentsMargins(4, 2, 4, 2) + + self.prev_btn = QPushButton("◀") + self.prev_btn.setToolTip("Предыдущий кадр (,)") + self.prev_btn.clicked.connect(lambda: self._step(-1)) + self.next_btn = QPushButton("▶") + self.next_btn.setToolTip("Следующий кадр (.)") + self.next_btn.clicked.connect(lambda: self._step(1)) + + self.frame_slider = QSlider(Qt.Horizontal) + self.frame_slider.setMinimum(0) + self.frame_slider.setMaximum(0) + self.frame_slider.setToolTip("Перемотка по кадрам (стрелки ← →)") + self.frame_slider.valueChanged.connect(self._on_slider) + + self.pos_label = QLabel("0 / 0") + self.pos_label.setMinimumWidth(90) + self.pos_label.setAlignment(Qt.AlignCenter) + + self.prev_hit_btn = QPushButton("◀ детекция") + self.prev_hit_btn.setToolTip("Предыдущий кадр с детекцией ([)") + self.prev_hit_btn.clicked.connect(lambda: self._step_hit(-1)) + self.next_hit_btn = QPushButton("детекция ▶") + self.next_hit_btn.setToolTip("Следующий кадр с детекцией (])") + self.next_hit_btn.clicked.connect(lambda: self._step_hit(1)) + + for wdg in (self.prev_btn, self.next_btn, self.frame_slider, self.pos_label, + self.prev_hit_btn, self.next_hit_btn): + h.addWidget(wdg, 1 if wdg is self.frame_slider else 0) + + # Keyboard shortcuts (window-wide), chosen to not clash with list/slider arrows. + QShortcut(QKeySequence(","), self, lambda: self._step(-1)) + QShortcut(QKeySequence("."), self, lambda: self._step(1)) + QShortcut(QKeySequence("["), self, lambda: self._step_hit(-1)) + QShortcut(QKeySequence("]"), self, lambda: self._step_hit(1)) + return bar + + # -------------------------------------------------------------- navigation + def _step(self, delta: int) -> None: + 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) + + def _step_hit(self, direction: int) -> None: + """Jump to the nearest frame (in `direction`) that has detections.""" + n = self.file_list.count() + if n == 0: + return + row = self.file_list.currentRow() + i = row + direction + while 0 <= i < n: + path = self.file_list.item(i).data(Qt.UserRole) + if self._results.get(path): + self.file_list.setCurrentRow(i) + return + i += direction + self.statusBar().showMessage( + "Больше нет кадров с детекцией в эту сторону " + "(сначала «Детектировать все»)" + ) + + def _on_slider(self, value: int) -> None: + if self._nav_sync: + return + if value != self.file_list.currentRow(): + self.file_list.setCurrentRow(value) + + def _update_nav(self) -> None: + n = self.file_list.count() + row = self.file_list.currentRow() + self._nav_sync = True + self.frame_slider.setMaximum(max(0, n - 1)) + self.frame_slider.setValue(max(0, row)) + self._nav_sync = False + self.pos_label.setText(f"{row + 1 if row >= 0 else 0} / {n}") + def _build_statusbar(self) -> None: self.progress = QProgressBar() self.progress.setMaximumWidth(260) @@ -292,14 +417,17 @@ class MainWindow(QMainWindow): self.file_list.blockSignals(False) self.progress.setVisible(False) + self._refresh_collections() if not files: self.view.set_image(None, []) + self._update_nav() 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: + self._update_nav() if current is not None: self._show(Path(current.data(Qt.UserRole))) # view only — no detection @@ -335,10 +463,12 @@ class MainWindow(QMainWindow): def _show(self, path: Path) -> None: """Display the image with its cached detections (does not run the detector).""" 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) + self._update_restore_actions() def _recompute_current(self) -> None: """Toolbar/Space: (re)run the detector on the selected frame.""" @@ -370,6 +500,65 @@ class MainWindow(QMainWindow): if self._current is not None: self._show(self._current) + # ------------------------------------------------------------- restoration + def _restore_current(self) -> None: + """Run the restorer on the current frame's detected regions and show it.""" + if self._current is None: + return + key = str(self._current) + if key not in self._results and self._detect(self._current) is None: + return + dets = self._results.get(key) or [] + if not dets: + self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать") + return + img = imread_unicode(key) + if img is None: + return + self.statusBar().showMessage(f"Восстановление: {self._current.name}…") + QApplication.processEvents() + try: + restored = self._restorer.restore(img, dets) + except Exception as exc: # noqa: BLE001 - surface model/engine errors + QMessageBox.warning(self, "Ошибка восстановления", str(exc)) + return + self._restored[key] = restored + self._showing_restored = True + self.view.set_image(restored, []) + self._update_restore_actions() + self.statusBar().showMessage( + f"Расцензурено ({self._restorer.name}): {self._current.name} — {len(dets)} обл." + ) + + def _toggle_restored(self) -> None: + if self._current is None or str(self._current) not in self._restored: + return + self._showing_restored = not self._showing_restored + key = str(self._current) + if self._showing_restored: + self.view.set_image(self._restored[key], []) + else: + self.view.set_image(imread_unicode(key), self._results.get(key) or []) + self._update_restore_actions() + + 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) + self.toggle_restored_action.setText( + "Показать оригинал" if self._showing_restored else "Показать результат" + ) + self.save_restored_action.setEnabled(has) + + def _save_restored(self) -> None: + if self._current is None or str(self._current) not in self._restored: + return + dest_dir = self._collection or self._current.parent + out = self._unique_dest(dest_dir, f"{self._current.stem}_restored.jpg") + if imwrite_unicode(str(out), self._restored[str(self._current)]): + self.statusBar().showMessage(f"Сохранено: {out}") + else: + QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.") + # ------------------------------------------------------------ collections def _collections_base(self) -> Path: """Where new collections are created: next to the opened folder, else home.""" @@ -377,6 +566,53 @@ class MainWindow(QMainWindow): return self._folder.parent return Path.home() / "HVideoTool" / "collections" + def _refresh_collections(self) -> None: + """Repopulate the collection combo from sibling folders of the opened folder. + + Keeps the active collection selected (and present even if it lives elsewhere). + """ + base = self._collections_base() + subdirs = [] + if base.exists(): + subdirs = sorted( + (p for p in base.iterdir() if p.is_dir() and p != self._folder), + key=lambda p: p.name.lower(), + ) + self.collection_combo.blockSignals(True) + self.collection_combo.clear() + self.collection_combo.addItem("— не выбрана —", None) + for p in subdirs: + self.collection_combo.addItem(p.name, str(p)) + # Make sure the active collection is listed even if it's outside base. + if self._collection is not None and self.collection_combo.findData(str(self._collection)) < 0: + self.collection_combo.addItem(self._collection.name, str(self._collection)) + self.collection_combo.addItem("Выбрать папку…", "__browse__") + self._select_active_in_combo() + self.collection_combo.blockSignals(False) + + def _select_active_in_combo(self) -> None: + idx = self.collection_combo.findData(str(self._collection)) if self._collection else 0 + self.collection_combo.setCurrentIndex(max(0, idx)) + + def _on_collection_selected(self, _index: int) -> None: + data = self.collection_combo.currentData() + if data == "__browse__": + self._browse_collection() + return + self._collection = Path(data) if data else None + if self._collection is not None: + self.statusBar().showMessage(f"Активная коллекция: {self._collection}") + + def _browse_collection(self) -> None: + start = str(self._collections_base()) + folder = QFileDialog.getExistingDirectory(self, "Выбрать коллекцию", start) + if folder: + self._collection = Path(folder) + self._refresh_collections() + self.statusBar().showMessage(f"Активная коллекция: {folder}") + else: + self._select_active_in_combo() # revert the combo to the current collection + def _create_collection(self) -> None: name, ok = QInputDialog.getText(self, "Создать коллекцию", "Имя коллекции:") name = name.strip() @@ -389,19 +625,14 @@ class MainWindow(QMainWindow): QMessageBox.warning(self, "Ошибка", f"Не удалось создать коллекцию:\n{exc}") return self._collection = path - self._update_collection_label() + self._refresh_collections() 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() @@ -433,6 +664,7 @@ class MainWindow(QMainWindow): self._show(Path(cur.data(Qt.UserRole))) elif self.file_list.count() == 0: self.view.set_image(None, []) + self._update_nav() @staticmethod def _unique_dest(folder: Path, name: str) -> Path: