diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1f5bf9c --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +build/ +dist/ +.venv/ +venv/ +env/ + +# Model weights & media — never committed (size / licensing) +models/ +*.pt +*.onnx +*.mp4 +*.mkv +*.avi +*.mov +*.webm + +# IDE / OS +.idea/ +.vscode/ +.DS_Store +Thumbs.db diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..7685e1a --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,201 @@ +# CLAUDE.md + +Guidance for Claude Code (and other agents) working in this repository. + +## Memory: EchoVault (read this first) + +This project uses the **EchoVault** MCP server for persistent, cross-session memory. +Prior sessions store architectural decisions, fixed bugs, and gotchas there. Follow +this protocol every session: + +1. **At session start — load context.** Call `memory_context` (project is + auto-detected from cwd) before doing any work. Use `memory_search` for specific + topics (e.g. "detector model", "classic-cv", "false positives"). +2. **During work — search before re-deciding.** When the task touches an area that + may have prior context, `memory_search` it first instead of re-deriving decisions. +3. **Before ending a session — save what matters.** Call `memory_save` when you made + a design decision, fixed a bug (include root cause + fix), found a non-obvious + gotcha, or the user corrected/clarified a requirement. Pick the right `category` + (`decision` / `bug` / `pattern` / `learning` / `context`). Do **not** save trivia, + things obvious from the code, or duplicates. + +EchoVault is the source of truth for *why* things are the way they are; this file is +the stable, high-level map. When they disagree, trust on-disk code first, then +EchoVault, then this file — and update whichever is stale. + +## What this project is + +**HVideoTool** is a Windows-first desktop GUI utility that **detects already-applied +censorship** (mosaic, pixelation, blur, black bars) in **images**, and draws outlines +over the detected censored regions. You open a folder of images; it runs each through +a detector, draws the regions, and shows a detailed per-image list of what it found. + +> **Scope was deliberately narrowed (this session).** It used to extract frames from +> video, detect, and play back with overlays (a "project model" with worker threads). +> That whole video pipeline was **removed** — the tool is now a simple **image-folder +> inspector** for viewing/debugging detector output on test images. Pre-extract video +> to frames externally if you need that. + +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.) +- 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. + +## Target environment + +- **OS:** Windows 11 x64 (primary). Use PowerShell syntax in commands. +- **Python:** 3.11+. +- **GPU:** NVIDIA + CUDA via PyTorch, only for the YOLO detector. CPU fallback works + but is slow. The `classic` detector needs no torch and no GPU. + +## Tech stack (decided) + +| Concern | Choice | +|----------------|---------------------------------------------| +| GUI | PySide6 (Qt 6) — LGPL | +| Image IO | OpenCV (`opencv-python`) + NumPy, unicode-safe via `core/imageio.py` | +| Detector | classic-CV heuristic; Ultralytics YOLO (LADA) behind a pluggable interface | + +Torch/CUDA + Ultralytics enter only with the YOLO detector. Keep that dependency +optional (the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is +installed separately per the README). The classic detector must keep running with no +torch present. + +## Architecture (as implemented) + +> This reflects the actual code on disk. It is a synchronous, single-threaded GUI app +> — no worker threads, no project/cache files. The only video touch is a one-shot +> "Создать из ролика…" that decodes a clip to a folder of JPGs via the **ffmpeg CLI** +> (cv2.VideoCapture fallback; NOT PyAV); detection still works on image folders only. +> Detection runs on the GUI +> thread (lazily per image, or via "Детектировать все"). When code and this file +> disagree, trust the code. + +``` +hvideotool/ +├── __main__.py # entry point + CLI (optional folder arg, --detector, --model) +├── app.py # QApplication bootstrap; run(config, folder=None) +├── config.py # AppConfig + DetectionConfig/OverlayConfig (thresholds live here) +├── settings_store.py # persist detector/model/threshold/last_dir to ~/HVideoTool/settings.json +├── ui/ +│ ├── main_window.py # the whole UI: toolbar + [file list | image view | detail table] +│ └── image_view.py # renders an image + draws polygon/bbox overlays (QPainter); can highlight one +└── core/ + ├── imageio.py # unicode-safe imread/imwrite (np.fromfile + imdecode) + ├── 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 + └── detection/ + ├── base.py # Detector ABC: detect(frame) -> list[Detection] + ├── factory.py # build_detector(config) -> classic | yolo | combined + ├── types.py # Detection (+ to_dict/from_dict), CensorType enum + ├── classic_cv.py # ClassicCVDetector — heuristic; accepts a `types` filter + ├── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics + └── composite.py # CompositeDetector — merge detectors + IoU dedup +``` + +### How it works + +- `MainWindow` holds the config, builds the detector lazily via `build_detector` + (cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path + -> list[Detection]]` as the detection cache. +- "Открыть папку" lists image files (`_IMAGE_EXTS`) with a progress bar (bulk insert + with `setUpdatesEnabled(False)` + periodic `processEvents`), so a big folder doesn't + freeze silently. +- **Viewing and detecting are decoupled on purpose** (so browsing stays instant even + with a slow CPU detector): selecting a file only *shows* it with its cached result + (header reads "не рассчитано" if none). Detection runs on **double-click**, the + "Рассчитать кадр" action (Space → `_recompute_current`, force-recomputes current), or + "Детектировать все" (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 + list/`_files`/cache. `_unique_dest` avoids clobbering (`foo.jpg` → `foo (1).jpg`). + Use case: sort frames into a training/example set while inspecting detections. +- `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 + the rest dim. The detail table lists ALL detections (sorted by score), so sub-threshold + hits are still visible for debugging; the threshold only affects what's drawn. + +### Separation of concerns + +- `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via + `core/detection/factory.build_detector` and talks to `core/` through the `Detector` + interface and the `Detection`/`CensorType` types. +- New detector kinds: implement `core/detection/base.Detector`, register the string in + `core/detection/factory.build_detector`, and add it to `_DETECTORS` in + `ui/main_window.py`. + +## Commands + +```powershell +python -m venv .venv; .\.venv\Scripts\Activate.ps1 +pip install -e . # classic detector needs no torch/CUDA + +python -m hvideotool # open a folder in-app +python -m hvideotool "C:\path\to\images" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt + +pip install -e ".[yolo]" # + install torch separately, see README +``` + +No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a +`MainWindow`, `open_path(folder)`, drive `file_list.setCurrentRow(...)`, and read +`detail_table` / `detail_header`. Or run `build_detector(config).detect(...)` on a +frame directly. + +## Conventions + +- Match the style of surrounding code; keep `core/` free of Qt where reasonable. +- Type hints on public functions and the `Detector` interface. +- Model weights (`.pt`) and large media are **not** committed — keep them in `models/` + and `.gitignore`d. +- User-facing strings / README are in Russian; code identifiers and this file in English. + +## Gotchas + +- **Wrong YOLO model = "noise".** The YOLO detector needs a **censorship** model + (LADA `models\lada_mosaic_detection_model_v4_accurate.pt`). If `model_path` points at + a generic COCO model (e.g. the `yolo11n-seg.pt` in the repo root, which Ultralytics + auto-downloads / is the training base), it detects people/objects and maps them to + `CensorType.UNKNOWN` → purple boxes that look like noise. This was a real user trap. +- **classic-CV is approximate and noisy on real video.** Its mosaic heuristic (low + block-reconstruction residual + 2D gradient + contrast) fires on textured real + footage (skin/hair/fabric/JPEG) → many false positives, while simultaneously missing + real mosaic after the `proc_max_dim=720` downscale softens block edges (measured: + contrast/grad fall below `mosaic_contrast_min`/`mosaic_grad_min`). For real-video + mosaic use `yolo`/`combined` + LADA. For anime there is no good public model. +- **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic + but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`), + not tuning more classic thresholds. +- **YOLO detector = LADA weights** ([HF `ladaapp/lada`](https://huggingface.co/ladaapp/lada)). + YOLO **segmentation** model, classes `{0: mosaic_nsfw, 1: mosaic_sfw_head}` → both map + to `CensorType.MOSAIC` (`_name_to_type` matches "mosaic" in the class name). Detects + mosaic only; black bars / blur stay with classic. Weights + Ultralytics are AGPL-3.0 + (accepted). `yolo.py` lazy-imports `torch`/`ultralytics`. +- **No model weights in the repo.** Code must fail with a clear, actionable message + when the model path is missing — not a raw stack trace (`factory._require_model`, + `YoloDetector.__init__`). +- **CUDA/torch install is environment-specific.** Don't add torch to core deps; it + stays out (the `yolo` extra pulls only Ultralytics) and is installed separately. +- **QImage from a numpy buffer must be `.copy()`d** (see `ImageView.set_image`), + otherwise it aliases a buffer that gets freed → garbage/crash. +- **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — + `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths. +- Don't reintroduce any generative / ControlNet dependency, nor the removed video + *pipeline* (PyAV, project/cache, worker threads, playback). The one allowed video + touch is `core/video/extract.py` (one-shot decode → JPG folder, behind "Создать из + ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by + the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is + ~10× faster than every-frame; `-hwaccel` does NOT help (GPU transfer overhead). Use + ffmpeg/cv2, not PyAV, and keep it synchronous. Decoding every frame is the inherent + cost — the speed lever is decoding *fewer* frames (keyframes). diff --git a/README.md b/README.md index 38d0acc..bf44c2f 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,170 @@ # HVideoTool +Десктопная утилита с графическим интерфейсом для **обнаружения уже наложенной +цензуры** (мозаика, пикселизация, размытие, чёрные плашки) на **картинках**. +Открываете папку с изображениями (или раскадровываете ролик кнопкой «Создать из +ролика…») — приложение прогоняет каждый кадр через детектор, **обводит найденные +области** и показывает **подробный список** того, что нашлось на каждой картинке. +Отобранные кадры можно перемещать в **коллекции** (для сбора датасета/примеров). + +> Инструмент для просмотра, отладки детекции и отбора кадров. Видео раскадровывает +> через ffmpeg (ставится автоматически с пакетом `imageio-ffmpeg`; системный ffmpeg +> из PATH используется в приоритете), либо через OpenCV. + +--- + +## Возможности + +- **Создать из ролика…** — раскадровка видео в папку-коллекцию. Два режима: + **только ключевые кадры** (в разы быстрее — декодируются лишь I-кадры) и + **каждый N-й кадр**; опциональный даунскейл (меньше файлов и нагрузки на диск/АВ). + ffmpeg идёт в комплекте (`imageio-ffmpeg`); системный ffmpeg из PATH — в приоритете. +- **Коллекции**: «Создать коллекцию…» + «В коллекцию» (Ctrl+M) перемещает выбранные + кадры в активную папку-коллекцию (мультивыбор поддерживается). +- **Открыть папку** с картинками (`.jpg/.png/.bmp/.webp/.tif`) — список слева. +- Картинка с **обводкой контуром** найденных областей — по центру. +- **Подробная таблица детекций** справа: тип, уверенность, bbox, число точек + полигона. Выбор строки **подсвечивает** конкретную область на картинке. +- **Ленивая детекция**: картинка прогоняется при первом открытии, результат + кэшируется. Кнопка **«Детектировать все»** обходит всю папку. +- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** + уверенности прямо в тулбаре — удобно сравнивать. +- Выбор файла весов модели кнопкой **«Модель…»**. + +## Что НЕ делает (осознанно вне области задачи) + +- Не **удаляет** и не **восстанавливает** зацензуренный контент. +- Не **генерирует** изображения (никакого ControlNet/SDXL/диффузии). +- Не детектирует «контент, который следовало бы зацензурить» (NSFW) — ищем + именно **уже наложенную** цензуру. +- Не декодирует видео — работает с готовыми картинками. + +--- + +## Требования + +- **ОС:** Windows 11 x64 (основная целевая платформа). +- **Python:** 3.11+. +- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора. CPU-режим работает, но + медленный. Для `classic` детектора ни torch, ни GPU не нужны. + +## Установка + +```powershell +git clone https://github.com/mrleo1nid/HVideoTool.git +cd HVideoTool +python -m venv .venv +.\.venv\Scripts\Activate.ps1 +pip install -e . +``` + +Этого достаточно для `classic` детектора — **PyTorch/CUDA не требуются**. Они +нужны только для YOLO/комбинированного детектора (см. +[Модель детектора](#модель-детектора)). + +## Запуск + +```powershell +# Без аргументов — папку открываете в приложении (тулбар → «Открыть папку…») +python -m hvideotool + +# Необязательно: сразу открыть папку / переопределить детектор и модель +python -m hvideotool "C:\path\to\images" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +``` + +Выбор детектора, путь к модели и порог сохраняются в `~/HVideoTool/settings.json` +и применяются при следующем запуске. + +--- + +## Модель детектора + +Интерфейс детектора абстрагирован (`core/detection/base.py`), детектор +выбирается в тулбаре (или флагом `--detector`): + +- `classic` — эвристический classic-CV детектор (`core/detection/classic_cv.py`). + Различает `mosaic` / `blur` / `black_bar`, без весов и без GPU. + **Приблизительный**: на реальном видео даёт много ложных срабатываний, заточен + скорее под рисованный/аниме контент — но и там ненадёжен. +- `yolo` — ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics) + (`core/detection/yolo.py`). **Сегментационная** модель — маски превращаются в + контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada). +- `combined` — `CompositeDetector`: YOLO (мозаика) + classic-CV (плашки/размытие), + результаты объединяются с дедупликацией по IoU. + +> **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры** +> (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать +> обычную COCO-модель (`yolo11n-seg.pt`), она будет детектить людей/предметы и +> помечать их как `unknown` — это и есть «шум». Файл LADA лежит в `models/`. + +> **Домен важен.** Модель LADA обучена на **реальном видео** (JAV); на части +> рисованного/аниме контента работает, на части — плохо. Хорошего публичного +> YOLO-детектора цензуры для аниме нет — это потребовало бы обучения своей модели. + +**Своя модель мозаики для аниме.** Можно обучить **YOLO11-seg** на синтетике +(накладываем мозаику на чистые кадры → авторазметка) и подключить `.pt` в наш +`YoloDetector` **без изменений кода**. Инструменты — в +[`scripts/training/`](scripts/training/README.md): + +```powershell +python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic +python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100 +# затем: тулбар → Детектор yolo → «Модель…» → runs\segment\mosaic\weights\best.pt +``` + +Установка YOLO-детектора: + +```powershell +pip install -e ".[yolo]" +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 + +curl.exe -L -o models\lada_mosaic_detection_model_v4_accurate.pt ` + "https://huggingface.co/ladaapp/lada/resolve/main/lada_mosaic_detection_model_v4_accurate.pt?download=true" + +python -m hvideotool --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +``` + +> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**. Classic-CV детектор +> от этого свободен. + +--- + +## Архитектура (кратко) + +``` +hvideotool/ +├── __main__.py # точка входа + CLI (всё опционально) +├── app.py # инициализация QApplication +├── config.py # настройки: пороги детекции, оверлей, детектор/модель +├── settings_store.py # детектор/модель/порог/последняя папка → settings.json +├── ui/ +│ ├── main_window.py # окно: список файлов | картинка | таблица детекций +│ └── image_view.py # отрисовка картинки + оверлей-контуры (QPainter) +└── core/ + ├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути) + ├── video/frame.py # Frame (картинка BGR + индекс + pts) — вход детектора + └── detection/ + ├── base.py # Detector (ABC): detect(frame) -> list[Detection] + ├── factory.py # build_detector(config) -> classic/yolo/combined + ├── types.py # Detection (+ to_dict/from_dict), CensorType + ├── classic_cv.py # эвристический детектор (mosaic/blur/black_bar) + ├── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны) + └── composite.py # CompositeDetector: объединение детекторов +``` + +Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU +заметно медленнее, чем на CUDA. + +## Технологический стек + +| Компонент | Выбор | +|----------------------|--------------------------------------------------| +| Язык | Python ≥ 3.11 | +| GUI | PySide6 (Qt 6) | +| Обработка картинок | OpenCV / NumPy | +| Детектор (без весов) | classic-CV эвристика (без GPU) | +| Детектор (ML) | Ultralytics YOLO + PyTorch/CUDA (LADA) | + +## Лицензия + +TBD. diff --git a/hvideotool/__init__.py b/hvideotool/__init__.py new file mode 100644 index 0000000..e6c3ec8 --- /dev/null +++ b/hvideotool/__init__.py @@ -0,0 +1,3 @@ +"""HVideoTool — detect and outline already-applied censorship in local videos.""" + +__version__ = "0.1.0" diff --git a/hvideotool/__main__.py b/hvideotool/__main__.py new file mode 100644 index 0000000..efc2160 --- /dev/null +++ b/hvideotool/__main__.py @@ -0,0 +1,39 @@ +"""Command-line entry point: ``python -m hvideotool``. + +Runs with no arguments — open a folder of images in-app. CLI flags are optional +overrides; choices persist to ~/HVideoTool/settings.json. +""" + +from __future__ import annotations + +import argparse +import sys + +from . import settings_store +from .app import run +from .config import AppConfig + + +def main() -> int: + parser = argparse.ArgumentParser( + prog="hvideotool", + description="Инспектор детекции уже наложенной цензуры на картинках.", + ) + parser.add_argument("folder", nargs="?", help="папка с картинками для немедленного открытия") + parser.add_argument("--detector", choices=["classic", "yolo", "combined"], default=None) + parser.add_argument("--model", dest="model_path", default=None, help="путь к весам (YOLO)") + args = parser.parse_args() + + config = AppConfig() + settings_store.apply(config) # persisted choices first + + if args.detector: + config.detector = args.detector + if args.model_path: + config.model_path = args.model_path + + return run(config, folder=args.folder) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/hvideotool/app.py b/hvideotool/app.py new file mode 100644 index 0000000..06fcf2f --- /dev/null +++ b/hvideotool/app.py @@ -0,0 +1,20 @@ +"""QApplication bootstrap.""" + +from __future__ import annotations + +import sys + +from PySide6.QtWidgets import QApplication + +from .config import AppConfig +from .ui.main_window import MainWindow + + +def run(config: AppConfig, folder: str | None = None) -> int: + app = QApplication(sys.argv) + app.setApplicationName("HVideoTool") + window = MainWindow(config) + window.show() + if folder: + window.open_path(folder) + return app.exec() diff --git a/hvideotool/config.py b/hvideotool/config.py new file mode 100644 index 0000000..64d81cd --- /dev/null +++ b/hvideotool/config.py @@ -0,0 +1,67 @@ +"""Application configuration and tunable defaults. + +Plain dataclasses. The detection thresholds matter most here — this tool is now +an image-folder inspector for tuning the detectors, so keep them easy to tweak. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class DetectionConfig: + """Parameters for the detector. Thresholds tuned for the classic-CV detector.""" + + proc_max_dim: int = 720 # downscale longer side to this before detection (speed) + min_area_frac: float = 0.0008 # ignore regions smaller than this fraction of the image + + # --- solid bars (black or white, achromatic, rectangular) --- + black_intensity: int = 40 # V below this = dark-bar candidate + white_intensity: int = 225 # V above this = light-bar candidate + bar_saturation_max: int = 45 # S below this = achromatic (excludes colored fills) + bar_min_extent: float = 0.80 # contour area / bbox area — how rectangular a bar must be + + # --- mosaic / pixelation --- + mosaic_block_sizes: tuple[int, ...] = (8, 12, 16, 24) # candidate tile sizes (px, proc space) + mosaic_residual_max: float = 6.0 # max reconstruction error to count as "blocky" + mosaic_contrast_min: float = 14.0 # min local contrast (excludes flat gradients) + mosaic_grad_min: float = 8.0 # min edge energy in BOTH x and y (excludes straight edges) + mosaic_min_side: int = 24 # reject thin regions (px) — kills edge false-positives + + # --- blur --- + blur_window: int = 31 # sliding window for local sharpness (odd) + blur_sharpness_ratio: float = 0.35 # below this fraction of median sharpness => blurry + blur_contrast_min: float = 8.0 # min local contrast (excludes flat regions) + + # --- YOLO detector (used only when detector == "yolo"/"combined") --- + yolo_conf: float = 0.2 # confidence threshold (LADA recommends ~0.2) + yolo_imgsz: int = 640 # inference image size + yolo_device: str | None = None # None => auto ("cuda" if available, else "cpu") + + +@dataclass(frozen=True) +class OverlayConfig: + """How detections are drawn over the image.""" + + # RGB per CensorType value + colors: dict[str, tuple[int, int, int]] = field( + default_factory=lambda: { + "mosaic": (231, 76, 60), # red + "blur": (241, 196, 15), # yellow + "black_bar": (26, 188, 156), # teal + "unknown": (155, 89, 182), # purple + } + ) + line_width: int = 2 + fill_alpha: int = 48 # 0..255 translucency of the region fill + show_labels: bool = True + + +@dataclass +class AppConfig: + detection: DetectionConfig = field(default_factory=DetectionConfig) + overlay: OverlayConfig = field(default_factory=OverlayConfig) + detector: str = "classic" # "classic" | "yolo" | "combined" + model_path: str | None = None # weights path, used by the YOLO detector + default_threshold: float = 0.20 # initial overlay confidence threshold diff --git a/hvideotool/core/__init__.py b/hvideotool/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hvideotool/core/detection/__init__.py b/hvideotool/core/detection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hvideotool/core/detection/base.py b/hvideotool/core/detection/base.py new file mode 100644 index 0000000..0ccdc94 --- /dev/null +++ b/hvideotool/core/detection/base.py @@ -0,0 +1,26 @@ +"""Detector interface. Implement this to plug in a new model (e.g. YOLO).""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from ..video.frame import Frame +from .types import Detection + + +class Detector(ABC): + """Abstract censorship detector. + + Implementations must be safe to call repeatedly on consecutive frames. They + receive a :class:`Frame` and return detections in *source-frame* pixel + coordinates (the same resolution as ``frame.image``). + """ + + @property + def name(self) -> str: + return type(self).__name__ + + @abstractmethod + def detect(self, frame: Frame) -> list[Detection]: + """Return detected censored regions for ``frame`` (possibly empty).""" + raise NotImplementedError diff --git a/hvideotool/core/detection/classic_cv.py b/hvideotool/core/detection/classic_cv.py new file mode 100644 index 0000000..cd52f2d --- /dev/null +++ b/hvideotool/core/detection/classic_cv.py @@ -0,0 +1,216 @@ +"""Weights-free, heuristic censorship detector (classic computer vision). + +APPROXIMATE BY DESIGN. This detector uses hand-tuned CV heuristics, not a +trained model. Its purpose is to make the whole pipeline runnable end-to-end +and to exercise the :class:`Detector` interface. For real-world accuracy, +replace it with a trained model (see ``yolo.py``, to be implemented) — the rest +of the app does not need to change. + +Heuristics: +- black_bar: large, near-uniform very dark regions (classic censor bars). +- mosaic: regions that reconstruct well from a coarse block grid (low + residual) yet have high coarse-scale contrast (i.e. blocky, not flat). +- blur: regions with local high-frequency energy far below the frame median, + while still being textured (excludes genuinely flat areas). +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from ...config import DetectionConfig +from ..video.frame import Frame +from .base import Detector +from .types import CensorType, Detection + + +class ClassicCVDetector(Detector): + def __init__( + self, + config: DetectionConfig | None = None, + types: "set[CensorType] | None" = None, + ) -> None: + self.cfg = config or DetectionConfig() + # Which censorship kinds to look for. Default: all. The composite detector + # restricts this to black_bar/blur (mosaic comes from the YOLO model). + self.types = ( + types if types is not None + else {CensorType.MOSAIC, CensorType.BLUR, CensorType.BLACK_BAR} + ) + + # ------------------------------------------------------------------ public + def detect(self, frame: Frame) -> list[Detection]: + bgr = frame.image + h0, w0 = bgr.shape[:2] + scale = self._proc_scale(w0, h0) + proc = ( + cv2.resize(bgr, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) + if scale != 1.0 + else bgr + ) + gray = cv2.cvtColor(proc, cv2.COLOR_BGR2GRAY) + ph, pw = gray.shape + min_area = self.cfg.min_area_frac * pw * ph + + dets: list[Detection] = [] + for ctype, fn, factor in ( + (CensorType.BLACK_BAR, self._detect_bars, 1.0), + (CensorType.MOSAIC, self._detect_mosaic, 4.0), + (CensorType.BLUR, self._detect_blur, 6.0), + ): + if ctype not in self.types: + continue + try: + dets += fn(proc, gray, min_area * factor) + except Exception: + # A failing heuristic must not break playback; skip it for this frame. + continue + + # Map proc-space coordinates back to source-frame pixels. + inv = 1.0 / scale + for d in dets: + x, y, w, h = d.bbox + d.bbox = (round(x * inv), round(y * inv), round(w * inv), round(h * inv)) + d.polygon = [(round(px * inv), round(py * inv)) for px, py in d.polygon] + return self._dedup(dets) + + # ----------------------------------------------------------------- helpers + def _proc_scale(self, w: int, h: int) -> float: + longest = max(w, h) + if longest <= self.cfg.proc_max_dim: + return 1.0 + return self.cfg.proc_max_dim / longest + + @staticmethod + def _local_std(g: np.ndarray, win: int) -> np.ndarray: + """Per-pixel standard deviation over a (win x win) box window.""" + mean = cv2.boxFilter(g, -1, (win, win)) + sqmean = cv2.boxFilter(g * g, -1, (win, win)) + var = np.maximum(sqmean - mean * mean, 0.0) + return np.sqrt(var) + + def _mask_to_detections( + self, + mask: np.ndarray, + ctype: CensorType, + min_area: float, + base_score: float, + min_extent: float = 0.0, + min_side: int = 0, + ) -> list[Detection]: + mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8)) + mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8)) + contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + out: list[Detection] = [] + for c in contours: + area = cv2.contourArea(c) + if area < min_area: + continue + x, y, w, h = cv2.boundingRect(c) + if min(w, h) < min_side: + continue # reject thin strips (e.g. edge false-positives) + extent = area / float(w * h + 1e-6) # how rectangular the blob is + if extent < min_extent: + continue + approx = cv2.approxPolyDP(c, 0.01 * cv2.arcLength(c, True), True) + poly = [(int(p[0][0]), int(p[0][1])) for p in approx] + score = float(np.clip(base_score + 0.25 * extent, 0.0, 1.0)) + out.append(Detection(type=ctype, score=score, bbox=(x, y, w, h), polygon=poly)) + return out + + # --------------------------------------------------------------- detectors + def _detect_bars(self, bgr, gray, min_area) -> list[Detection]: + # Solid censor bars are achromatic (black OR white) rectangles. Requiring + # low saturation + high rectangularity excludes large flat *colored* fills + # that are common in drawn/anime backgrounds. + hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + sat, val = hsv[:, :, 1], hsv[:, :, 2] + achromatic = sat < self.cfg.bar_saturation_max + dark = (val < self.cfg.black_intensity) & achromatic + light = (val > self.cfg.white_intensity) & achromatic + mask = (dark | light).astype(np.uint8) * 255 + return self._mask_to_detections( + mask, CensorType.BLACK_BAR, min_area, base_score=0.55, + min_extent=self.cfg.bar_min_extent, + ) + + def _detect_mosaic(self, bgr, gray, min_area) -> list[Detection]: + g = gray.astype(np.float32) + h, w = gray.shape + win = 17 + # Lowest reconstruction residual across candidate tile sizes AND grid phases. + # Real mosaics aren't aligned to the origin, so we try a few offsets per size + # (phase-invariant) and keep the best fit. + best_residual = np.full((h, w), np.inf, np.float32) + for b in self.cfg.mosaic_block_sizes: + half = b // 2 + for oy, ox in ((0, 0), (half, 0), (0, half), (half, half)): + sub = g[oy:, ox:] + sh, sw = sub.shape + if sh < b or sw < b: + continue + small = cv2.resize(sub, (max(1, sw // b), max(1, sh // b)), interpolation=cv2.INTER_AREA) + restored = cv2.resize(small, (sw, sh), interpolation=cv2.INTER_NEAREST) + region = best_residual[oy:oy + sh, ox:ox + sw] + np.minimum(region, np.abs(sub - restored), out=region) + best_residual = cv2.boxFilter(best_residual, -1, (win, win)) + + contrast = self._local_std(g, win) + # Mosaic has edges in BOTH directions; a lone straight boundary (flat-region + # border, bar edge) has edge energy in only one — exclude those. + gx = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)), -1, (win, win)) + gy = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)), -1, (win, win)) + both_dirs = (gx > self.cfg.mosaic_grad_min) & (gy > self.cfg.mosaic_grad_min) + + blocky = best_residual < self.cfg.mosaic_residual_max + textured = contrast > self.cfg.mosaic_contrast_min + mask = (blocky & textured & both_dirs).astype(np.uint8) * 255 + return self._mask_to_detections( + mask, CensorType.MOSAIC, min_area, base_score=0.50, min_side=self.cfg.mosaic_min_side + ) + + def _detect_blur(self, bgr, gray, min_area) -> list[Detection]: + g = gray.astype(np.float32) + win = self.cfg.blur_window | 1 # force odd + lap = cv2.Laplacian(g, cv2.CV_32F, ksize=3) + sharpness = cv2.boxFilter(lap * lap, -1, (win, win)) # local high-freq energy + median = float(np.median(sharpness)) + 1e-6 + contrast = self._local_std(g, win) + + blurry = sharpness < median * self.cfg.blur_sharpness_ratio + textured = contrast > self.cfg.blur_contrast_min + mask = (blurry & textured).astype(np.uint8) * 255 + return self._mask_to_detections( + mask, CensorType.BLUR, min_area, base_score=0.40, min_side=self.cfg.mosaic_min_side + ) + + # ----------------------------------------------------------------- dedup + def _dedup(self, dets: list[Detection]) -> list[Detection]: + """Greedy IoU suppression; prefer black_bar > mosaic > blur, then score.""" + priority = { + CensorType.BLACK_BAR: 3, + CensorType.MOSAIC: 2, + CensorType.BLUR: 1, + CensorType.UNKNOWN: 0, + } + dets = sorted(dets, key=lambda d: (priority[d.type], d.score), reverse=True) + kept: list[Detection] = [] + for d in dets: + if all(self._iou(d.bbox, k.bbox) < 0.5 for k in kept): + kept.append(d) + return kept + + @staticmethod + def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float: + ax, ay, aw, ah = a + bx, by, bw, bh = b + ix = max(ax, bx) + iy = max(ay, by) + ix2 = min(ax + aw, bx + bw) + iy2 = min(ay + ah, by + bh) + iw, ih = max(0, ix2 - ix), max(0, iy2 - iy) + inter = iw * ih + union = aw * ah + bw * bh - inter + return inter / union if union > 0 else 0.0 diff --git a/hvideotool/core/detection/composite.py b/hvideotool/core/detection/composite.py new file mode 100644 index 0000000..25b4924 --- /dev/null +++ b/hvideotool/core/detection/composite.py @@ -0,0 +1,51 @@ +"""Composite detector: runs several detectors and merges their results. + +Used for the "combined" mode = YOLO (mosaic) + classic-CV (black bars / blur). +Detections from all sub-detectors are concatenated, then de-duplicated by IoU +(higher score wins) so overlapping hits from different detectors don't stack. +""" + +from __future__ import annotations + +from ..video.frame import Frame +from .base import Detector +from .types import Detection + + +class CompositeDetector(Detector): + def __init__(self, detectors: list[Detector], iou_threshold: float = 0.6) -> None: + if not detectors: + raise ValueError("CompositeDetector requires at least one detector") + self._detectors = detectors + self._iou = iou_threshold + + @property + def name(self) -> str: + return "Composite(" + " + ".join(d.name for d in self._detectors) + ")" + + def detect(self, frame: Frame) -> list[Detection]: + merged: list[Detection] = [] + for detector in self._detectors: + try: + merged += detector.detect(frame) + except Exception: # noqa: BLE001 - one detector failing must not kill the frame + continue + return self._dedup(merged) + + def _dedup(self, dets: list[Detection]) -> list[Detection]: + dets = sorted(dets, key=lambda d: d.score, reverse=True) + kept: list[Detection] = [] + for d in dets: + if all(self._iou_of(d.bbox, k.bbox) < self._iou for k in kept): + kept.append(d) + return kept + + @staticmethod + def _iou_of(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float: + ax, ay, aw, ah = a + bx, by, bw, bh = b + ix, iy = max(ax, bx), max(ay, by) + ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh) + inter = max(0, ix2 - ix) * max(0, iy2 - iy) + union = aw * ah + bw * bh - inter + return inter / union if union > 0 else 0.0 diff --git a/hvideotool/core/detection/factory.py b/hvideotool/core/detection/factory.py new file mode 100644 index 0000000..fc6aaf2 --- /dev/null +++ b/hvideotool/core/detection/factory.py @@ -0,0 +1,41 @@ +"""Detector factory: build a Detector from AppConfig. + +Kept separate from ``app.py`` so both the app bootstrap and the UI can build +detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on +bad config so the GUI can show the message instead of exiting. +""" + +from __future__ import annotations + +from ...config import AppConfig +from .base import Detector +from .classic_cv import ClassicCVDetector +from .types import CensorType + + +def _require_model(config: AppConfig) -> str: + if not config.model_path: + raise ValueError( + "Для детектора YOLO укажите путь к весам (.pt) в Параметрах " + "или скачайте модель LADA — см. README." + ) + return config.model_path + + +def build_detector(config: AppConfig) -> Detector: + if config.detector == "classic": + return ClassicCVDetector(config.detection) + if config.detector == "yolo": + from .yolo import YoloDetector # lazy: pulls torch/ultralytics + + return YoloDetector(_require_model(config), config.detection) + if config.detector == "combined": + # YOLO handles mosaic; classic-CV handles black bars / blur. + from .composite import CompositeDetector + from .yolo import YoloDetector + + return CompositeDetector([ + YoloDetector(_require_model(config), config.detection), + ClassicCVDetector(config.detection, types={CensorType.BLACK_BAR, CensorType.BLUR}), + ]) + raise ValueError(f"Неизвестный детектор: {config.detector!r}") diff --git a/hvideotool/core/detection/types.py b/hvideotool/core/detection/types.py new file mode 100644 index 0000000..2c9a243 --- /dev/null +++ b/hvideotool/core/detection/types.py @@ -0,0 +1,42 @@ +"""Detection result types shared across detectors and the UI.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum + + +class CensorType(str, Enum): + """Kind of already-applied censorship a detection represents.""" + + MOSAIC = "mosaic" + BLUR = "blur" + BLACK_BAR = "black_bar" + UNKNOWN = "unknown" + + +@dataclass +class Detection: + """A single detected censored region, in source-frame pixel coordinates.""" + + type: CensorType + score: float # confidence, 0..1 + bbox: tuple[int, int, int, int] # x, y, w, h + polygon: list[tuple[int, int]] = field(default_factory=list) # contour points + + def to_dict(self) -> dict: + return { + "type": self.type.value, + "score": self.score, + "bbox": list(self.bbox), + "polygon": [list(p) for p in self.polygon], + } + + @classmethod + def from_dict(cls, data: dict) -> "Detection": + return cls( + type=CensorType(data["type"]), + score=float(data["score"]), + bbox=tuple(data["bbox"]), # type: ignore[arg-type] + polygon=[tuple(p) for p in data.get("polygon", [])], + ) diff --git a/hvideotool/core/detection/yolo.py b/hvideotool/core/detection/yolo.py new file mode 100644 index 0000000..e789d85 --- /dev/null +++ b/hvideotool/core/detection/yolo.py @@ -0,0 +1,104 @@ +"""Ultralytics YOLO detector. + +Wraps an Ultralytics YOLO model (detection or segmentation) behind the +:class:`Detector` interface. Designed for the LADA mosaic-detection weights +(https://huggingface.co/ladaapp/lada), which are YOLO *segmentation* models with +a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose class +names map onto :class:`CensorType`. + +Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the +rest of the app — and the classic-CV detector — never pull them in. + +Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README. +""" + +from __future__ import annotations + +import os + +import numpy as np + +from ...config import DetectionConfig +from ..video.frame import Frame +from .base import Detector +from .types import CensorType, Detection + + +def _name_to_type(name: str) -> CensorType: + n = name.lower() + if "mosaic" in n or "pixel" in n: + return CensorType.MOSAIC + if "blur" in n: + return CensorType.BLUR + if "bar" in n or "black" in n: + return CensorType.BLACK_BAR + return CensorType.UNKNOWN + + +class YoloDetector(Detector): + def __init__(self, model_path: str, config: DetectionConfig | None = None) -> None: + self.cfg = config or DetectionConfig() + if not os.path.isfile(model_path): + raise FileNotFoundError( + f"Файл весов не найден: {model_path}\n" + "Скачайте модель детекции мозаики LADA, например:\n" + " curl.exe -L -o models\\lada_mosaic_detection_model_v4_accurate.pt " + '"https://huggingface.co/ladaapp/lada/resolve/main/' + 'lada_mosaic_detection_model_v4_accurate.pt?download=true"' + ) + try: + from ultralytics import YOLO + except ImportError as exc: # pragma: no cover - environment dependent + raise ImportError( + "Не установлен ultralytics. Установите: pip install -e \".[yolo]\" " + "(и PyTorch с CUDA отдельно — см. README)." + ) from exc + + # Resolve the device: explicit override, else CUDA when available. + device = self.cfg.yolo_device + if device is None: + try: + import torch + + device = "cuda" if torch.cuda.is_available() else "cpu" + except Exception: # noqa: BLE001 + device = "cpu" + self._device = device + self._model = YOLO(model_path) + + @property + def name(self) -> str: + return f"YoloDetector(device={self._device})" + + def detect(self, frame: Frame) -> list[Detection]: + results = self._model.predict( + source=frame.image, # BGR ndarray; ultralytics handles it + conf=self.cfg.yolo_conf, + imgsz=self.cfg.yolo_imgsz, + device=self._device, + verbose=False, + ) + if not results: + return [] + res = results[0] + boxes = getattr(res, "boxes", None) + if boxes is None or len(boxes) == 0: + return [] + + names = res.names # {class_index: class_name} + xyxy = boxes.xyxy.cpu().numpy() + confs = boxes.conf.cpu().numpy() + classes = boxes.cls.cpu().numpy().astype(int) + # Segmentation polygons in source-pixel coords, one per detection (if any). + polygons = res.masks.xy if getattr(res, "masks", None) is not None else None + + out: list[Detection] = [] + for i in range(len(xyxy)): + x1, y1, x2, y2 = xyxy[i] + bbox = (int(x1), int(y1), int(x2 - x1), int(y2 - y1)) + poly: list[tuple[int, int]] = [] + if polygons is not None and i < len(polygons): + poly = [(int(px), int(py)) for px, py in polygons[i]] + ctype = _name_to_type(names.get(int(classes[i]), "")) + out.append(Detection(type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly)) + return out diff --git a/hvideotool/core/imageio.py b/hvideotool/core/imageio.py new file mode 100644 index 0000000..6f50f09 --- /dev/null +++ b/hvideotool/core/imageio.py @@ -0,0 +1,32 @@ +"""Unicode-safe image read/write. + +``cv2.imread``/``cv2.imwrite`` mishandle non-ASCII paths on Windows. These +helpers go through ``np.fromfile``/``ndarray.tofile`` + ``imdecode``/``imencode`` +so paths with Cyrillic (etc.) work regardless of the system locale. +""" + +from __future__ import annotations + +import os + +import cv2 +import numpy as np + + +def imread_unicode(path: str) -> np.ndarray | None: + try: + data = np.fromfile(path, dtype=np.uint8) + except OSError: + return None + if data.size == 0: + return None + return cv2.imdecode(data, cv2.IMREAD_COLOR) + + +def imwrite_unicode(path: str, image: np.ndarray, params: list[int] | None = None) -> bool: + ext = os.path.splitext(path)[1] or ".jpg" + ok, buf = cv2.imencode(ext, image, params or []) + if not ok: + return False + buf.tofile(path) + return True diff --git a/hvideotool/core/video/__init__.py b/hvideotool/core/video/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hvideotool/core/video/extract.py b/hvideotool/core/video/extract.py new file mode 100644 index 0000000..6462967 --- /dev/null +++ b/hvideotool/core/video/extract.py @@ -0,0 +1,168 @@ +"""Extract frames from a video into a folder of JPGs. + +Two engines: + +* **ffmpeg** (preferred, used when the ``ffmpeg`` binary is on PATH) — one + subprocess does decode + sampling + optional downscale + JPEG encode, which is + faster than pulling frames into Python one by one, and unlocks the big win: + *keyframe-only* extraction (``-skip_frame nokey`` decodes only I-frames, ~10× + faster than decoding every frame). +* **OpenCV** fallback (``cv2.VideoCapture``) when ffmpeg is absent. + +Decoding H.264/HEVC frame-by-frame is inherently the cost; hardware accel doesn't +help for this (GPU transfer overhead). The only way to be dramatically faster is +to decode fewer frames — hence the keyframe mode. +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path +from typing import Callable + +import cv2 + +from ..imageio import imwrite_unicode + +# progress(done_seconds, total_seconds) -> return False to cancel. +Progress = Callable[[float, float], bool] + + +def _find_ffmpeg() -> str | None: + """ffmpeg on PATH, else the binary bundled with imageio-ffmpeg, else None.""" + exe = shutil.which("ffmpeg") + if exe: + return exe + try: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + except Exception: # noqa: BLE001 - package missing or no bundled binary + return None + + +def _video_duration_seconds(video_path: str) -> float: + cap = cv2.VideoCapture(str(video_path)) + try: + fps = cap.get(cv2.CAP_PROP_FPS) or 0.0 + frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0 + return frames / fps if fps > 0 else 0.0 + finally: + cap.release() + + +def _quality_to_qscale(jpg_quality: int) -> int: + """Map JPEG quality 0..100 to ffmpeg -q:v (2 best .. 31 worst).""" + q = round(2 + (100 - max(0, min(100, jpg_quality))) / 100 * 29) + return max(2, min(31, q)) + + +def extract_frames( + video_path: str, + out_dir: str, + step: int = 15, + keyframes_only: bool = False, + max_dim: int = 0, + jpg_quality: int = 92, + progress: Progress | None = None, +) -> int: + """Save sampled frames of ``video_path`` into ``out_dir`` as JPGs. + + ``keyframes_only`` decodes only keyframes (fast). Otherwise keeps every + ``step``-th frame. ``max_dim`` (>0) caps the longest side. ``progress`` is + called with (done_seconds, total_seconds); returning ``False`` cancels. + Returns the number of frames written. + """ + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + ffmpeg = _find_ffmpeg() + if ffmpeg: + return _extract_ffmpeg(ffmpeg, video_path, out, step, keyframes_only, max_dim, jpg_quality, progress) + return _extract_cv2(video_path, out, step, max_dim, jpg_quality, progress) + + +# --------------------------------------------------------------------- ffmpeg +def _build_vf(step: int, keyframes_only: bool, max_dim: int) -> str | None: + filters: list[str] = [] + if not keyframes_only and step > 1: + filters.append(f"select=not(mod(n\\,{int(step)}))") + if max_dim and max_dim > 0: + # Cap the longest side to max_dim, preserve aspect, never upscale. + filters.append(f"scale='min({max_dim},iw)':'min({max_dim},ih)':force_original_aspect_ratio=decrease") + return ",".join(filters) if filters else None + + +def _extract_ffmpeg( + ffmpeg: str, video_path: str, out: Path, step: int, + keyframes_only: bool, max_dim: int, jpg_quality: int, progress: Progress | None, +) -> int: + total = _video_duration_seconds(video_path) + cmd = [ffmpeg, "-hide_banner", "-loglevel", "error", "-nostdin"] + if keyframes_only: + cmd += ["-skip_frame", "nokey"] # input option: decode only keyframes + cmd += ["-i", video_path] + vf = _build_vf(step, keyframes_only, max_dim) + if vf: + cmd += ["-vf", vf] + cmd += ["-vsync", "0", "-q:v", str(_quality_to_qscale(jpg_quality)), + "-progress", "pipe:1", str(out / "%06d.jpg")] + + proc = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, + stdin=subprocess.DEVNULL, text=True, bufsize=1, + ) + try: + assert proc.stdout is not None + for line in proc.stdout: + if progress is None: + continue + line = line.strip() + if line.startswith("out_time_us=") or line.startswith("out_time_ms="): + raw = line.split("=", 1)[1] + try: + # out_time_us is microseconds; out_time_ms is *also* microseconds + # in ffmpeg (historical misnomer). Both -> seconds via /1e6. + done = int(raw) / 1_000_000 if raw.isdigit() else 0.0 + except ValueError: + done = 0.0 + if progress(done, total) is False: + proc.terminate() + break + finally: + proc.wait() + return len(list(out.glob("*.jpg"))) + + +# ---------------------------------------------------------------------- opencv +def _extract_cv2( + video_path: str, out: Path, step: int, max_dim: int, jpg_quality: int, progress: Progress | None, +) -> int: + cap = cv2.VideoCapture(str(video_path)) + if not cap.isOpened(): + raise RuntimeError(f"Не удалось открыть видео: {video_path}") + step = max(1, int(step)) + fps = cap.get(cv2.CAP_PROP_FPS) or 0.0 + total = (cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0) / fps if fps > 0 else 0.0 + params = [cv2.IMWRITE_JPEG_QUALITY, int(jpg_quality)] + idx = saved = 0 + try: + while True: + if not cap.grab(): + break + if idx % step == 0: + ok, frame = cap.retrieve() + if ok: + if max_dim and max(frame.shape[:2]) > max_dim: + s = max_dim / max(frame.shape[:2]) + frame = cv2.resize(frame, None, fx=s, fy=s, interpolation=cv2.INTER_AREA) + imwrite_unicode(str(out / f"{idx:06d}.jpg"), frame, params) + saved += 1 + idx += 1 + if progress is not None and idx % 30 == 0: + done = idx / fps if fps > 0 else 0.0 + if progress(done, total) is False: + break + finally: + cap.release() + return saved diff --git a/hvideotool/core/video/frame.py b/hvideotool/core/video/frame.py new file mode 100644 index 0000000..4c58151 --- /dev/null +++ b/hvideotool/core/video/frame.py @@ -0,0 +1,14 @@ +"""Decoded video frame passed from the reader to the detector.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass +class Frame: + image: np.ndarray # BGR, HxWx3, uint8 (OpenCV convention) + index: int # 0-based frame counter since the last open/seek + pts: float # presentation timestamp, seconds diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py new file mode 100644 index 0000000..d4daa59 --- /dev/null +++ b/hvideotool/settings_store.py @@ -0,0 +1,58 @@ +"""Persist a handful of user choices to ~/HVideoTool/settings.json. + +Slimmed down for the image-inspector tool: it remembers the detector, the model +path, the overlay threshold, and the last opened folder. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from .config import AppConfig + +_PATH = Path.home() / "HVideoTool" / "settings.json" + + +def _read() -> dict: + try: + return json.loads(_PATH.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + + +def _write(data: dict) -> None: + _PATH.parent.mkdir(parents=True, exist_ok=True) + _PATH.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") + + +def apply(config: AppConfig) -> None: + """Overlay persisted settings onto ``config`` (mutates it in place).""" + data = _read() + if data.get("detector"): + config.detector = data["detector"] + if "model_path" in data: + config.model_path = data["model_path"] + if "threshold" in data: + config.default_threshold = float(data["threshold"]) + + +def save(config: AppConfig) -> None: + """Persist the configurable settings, preserving other keys (e.g. last_dir).""" + data = _read() + data.update( + detector=config.detector, + model_path=config.model_path, + threshold=config.default_threshold, + ) + _write(data) + + +def last_dir() -> str | None: + return _read().get("last_dir") + + +def set_last_dir(path: str) -> None: + data = _read() + data["last_dir"] = path + _write(data) diff --git a/hvideotool/ui/__init__.py b/hvideotool/ui/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hvideotool/ui/extract_dialog.py b/hvideotool/ui/extract_dialog.py new file mode 100644 index 0000000..0e0aef1 --- /dev/null +++ b/hvideotool/ui/extract_dialog.py @@ -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() diff --git a/hvideotool/ui/image_view.py b/hvideotool/ui/image_view.py new file mode 100644 index 0000000..ae8863a --- /dev/null +++ b/hvideotool/ui/image_view.py @@ -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) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py new file mode 100644 index 0000000..7e837c5 --- /dev/null +++ b/hvideotool/ui/main_window.py @@ -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"{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) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9320adf --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "hvideotool" +version = "0.1.0" +description = "Desktop GUI utility that detects and outlines already-applied censorship in images" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "TBD" } +authors = [{ name = "Leonid Pershin" }] +dependencies = [ + "PySide6>=6.6", + "opencv-python>=4.8", + "numpy>=1.24", + "imageio-ffmpeg>=0.4", # bundles an ffmpeg binary for "Создать из ролика…" +] + +[project.optional-dependencies] +# Trained-model detector (future work). PyTorch must be installed separately +# with the correct CUDA build — see README. Installing this extra only pulls in +# Ultralytics; it does NOT install torch. +yolo = ["ultralytics>=8.0"] + +[project.scripts] +hvideotool = "hvideotool.__main__:main" + +[tool.setuptools.packages.find] +include = ["hvideotool*"] diff --git a/scripts/training/README.md b/scripts/training/README.md new file mode 100644 index 0000000..c5268ca --- /dev/null +++ b/scripts/training/README.md @@ -0,0 +1,50 @@ +# Обучение детектора мозаики (аниме) + +Готовых публичных моделей детекции мозаики для рисованного/аниме контента нет, +поэтому обучаем свою на **синтетике**: берём чистые (без цензуры) кадры, случайно +накладываем мозаику и получаем разметку автоматически. + +## 1. Подготовить чистые изображения + +Сложите кадры **без цензуры** в папку (чем разнообразнее, тем лучше; от ~500 +картинок для черновой модели, тысячи — для нормальной). Кадры можно нарезать из +чистых видео — например, любым плеером или нашим извлечением (папка проекта +`frames/`). + +## 2. Сгенерировать датасет + +```powershell +.\.venv\Scripts\Activate.ps1 +python scripts\training\gen_mosaic_dataset.py --input C:\clean_frames --output dataset_mosaic --variants 3 +``` + +Получите `dataset_mosaic/` с `images/`, `labels/` и `data.yaml` +(класс `0: mosaic`, формат сегментации YOLO). Часть кадров остаётся чистыми +(негативы) — это снижает ложные срабатывания. + +Полезные флаги: `--neg-frac 0.2`, `--tile-min/--tile-max` (размер плиток мозаики), +`--area-min/--area-max` (доля площади под мозаику), `--shapes rect,ellipse`. + +## 3. Обучить + +```powershell +pip install -e ".[yolo]" +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 # CUDA +python scripts\training\train_mosaic.py --data dataset_mosaic\data.yaml --epochs 100 +``` + +Веса появятся в `runs/segment/mosaic/weights/best.pt`. + +## 4. Подключить в приложении + +**Параметры → Детектор: YOLO → Обзор…** и выбрать `best.pt`, затем +**Файл → Детектировать заново**. Класс `mosaic` маппится в `CensorType.MOSAIC`, +маски превращаются в контуры обводки. + +## Советы по качеству + +- Разнообразьте источник (разные стили, сцены, освещение). +- Варьируйте размер плиток и областей (флаги выше) — мозаика в реальности разная. +- Доля негативов 15–25 % обычно хорошо снижает ложные срабатывания. +- GPU NVIDIA сильно ускоряет; на CPU обучение очень медленное. +- Дальше можно добавить классы `bar` (плашки) — генератор легко расширить. diff --git a/scripts/training/gen_mosaic_dataset.py b/scripts/training/gen_mosaic_dataset.py new file mode 100644 index 0000000..c45ba1d --- /dev/null +++ b/scripts/training/gen_mosaic_dataset.py @@ -0,0 +1,164 @@ +"""Generate a synthetic YOLO-seg dataset for MOSAIC detection. + +Takes a folder of CLEAN (uncensored) images — anime frames work best for the +anime domain — and produces censored copies with random mosaic regions plus +matching YOLO segmentation labels (class 0 = mosaic). Some outputs are left +clean (negatives / background) so the model learns what is *not* mosaic. + +The model only needs to recognise mosaic *texture*, so random placement is fine +(we detect already-applied mosaic anywhere, not "where to censor"). + +Output layout (Ultralytics format): + /images/train/*.jpg /labels/train/*.txt + /images/val/*.jpg /labels/val/*.txt + /data.yaml + +Usage: + python scripts/training/gen_mosaic_dataset.py --input clean_frames --output dataset_mosaic +""" + +from __future__ import annotations + +import argparse +import random +from pathlib import Path + +import cv2 +import numpy as np + +IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} + + +# --- unicode-safe IO (self-contained, no hvideotool import needed) ------------ +def imread(path: Path) -> "np.ndarray | None": + data = np.fromfile(str(path), dtype=np.uint8) + if data.size == 0: + return None + return cv2.imdecode(data, cv2.IMREAD_COLOR) + + +def imwrite(path: Path, img: np.ndarray, quality: int = 92) -> None: + ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, quality]) + if ok: + buf.tofile(str(path)) + + +# --- mosaic + polygon --------------------------------------------------------- +def pixelate_region(img: np.ndarray, poly: np.ndarray, tile: int) -> None: + """Apply mosaic inside the polygon (in place). Clamps to image bounds.""" + H, W = img.shape[:2] + x, y, w, h = cv2.boundingRect(poly) + x, y = max(0, x), max(0, y) + x2, y2 = min(x + w, W), min(y + h, H) + w, h = x2 - x, y2 - y + if w < 1 or h < 1: + return + roi = img[y:y2, x:x2] + small = cv2.resize(roi, (max(1, w // tile), max(1, h // tile)), interpolation=cv2.INTER_LINEAR) + mosaic = cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST) + mask = np.zeros((h, w), np.uint8) + cv2.fillPoly(mask, [poly - [x, y]], 255) + roi[mask > 0] = mosaic[mask > 0] + + +def make_region(W: int, H: int, area_min: float, area_max: float, shape: str) -> np.ndarray: + """Return an Nx2 int polygon for a random mosaic region within the image.""" + area = random.uniform(area_min, area_max) * W * H + aspect = random.uniform(0.5, 2.0) + w = int(min(W * 0.9, max(24, (area * aspect) ** 0.5))) + h = int(min(H * 0.9, max(24, area / max(1, w)))) + x = random.randint(0, max(0, W - w)) + y = random.randint(0, max(0, H - h)) + if shape == "ellipse": + cx, cy = x + w // 2, y + h // 2 + pts = cv2.ellipse2Poly((cx, cy), (w // 2, h // 2), random.randint(0, 180), 0, 360, 20) + pts[:, 0] = np.clip(pts[:, 0], 0, W - 1) + pts[:, 1] = np.clip(pts[:, 1], 0, H - 1) + return pts.astype(np.int32) + return np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]], np.int32) + + +def poly_to_label(poly: np.ndarray, W: int, H: int) -> str: + coords = [] + for px, py in poly: + coords.append(f"{np.clip(px / W, 0, 1):.6f}") + coords.append(f"{np.clip(py / H, 0, 1):.6f}") + return "0 " + " ".join(coords) + + +def main() -> None: + ap = argparse.ArgumentParser(description="Synthetic mosaic YOLO-seg dataset generator") + ap.add_argument("--input", required=True, help="folder of clean (uncensored) images") + ap.add_argument("--output", required=True, help="output dataset folder") + ap.add_argument("--variants", type=int, default=3, help="augmented copies per source image") + ap.add_argument("--val-split", type=float, default=0.15) + ap.add_argument("--neg-frac", type=float, default=0.2, help="fraction of outputs left clean") + ap.add_argument("--min-regions", type=int, default=1) + ap.add_argument("--max-regions", type=int, default=3) + ap.add_argument("--tile-min", type=int, default=6) + ap.add_argument("--tile-max", type=int, default=22) + ap.add_argument("--area-min", type=float, default=0.02) + ap.add_argument("--area-max", type=float, default=0.22) + ap.add_argument("--max-dim", type=int, default=1280, help="downscale clean images larger than this") + ap.add_argument("--shapes", default="rect,ellipse") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + random.seed(args.seed) + np.random.seed(args.seed) + shapes = [s.strip() for s in args.shapes.split(",") if s.strip()] + + sources = sorted(p for p in Path(args.input).rglob("*") if p.suffix.lower() in IMG_EXTS) + if not sources: + raise SystemExit(f"Не найдено изображений в {args.input}") + random.shuffle(sources) + n_val = max(1, int(len(sources) * args.val_split)) + val_set = set(sources[:n_val]) + + out = Path(args.output) + for split in ("train", "val"): + (out / "images" / split).mkdir(parents=True, exist_ok=True) + (out / "labels" / split).mkdir(parents=True, exist_ok=True) + + counts = {"train": 0, "val": 0, "neg": 0, "pos": 0} + for src in sources: + img0 = imread(src) + if img0 is None: + continue + H0, W0 = img0.shape[:2] + scale = min(1.0, args.max_dim / max(H0, W0)) + if scale < 1.0: + img0 = cv2.resize(img0, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) + H, W = img0.shape[:2] + split = "val" if src in val_set else "train" + + for v in range(args.variants): + img = img0.copy() + lines: list[str] = [] + if random.random() >= args.neg_frac: + for _ in range(random.randint(args.min_regions, args.max_regions)): + shape = random.choice(shapes) + poly = make_region(W, H, args.area_min, args.area_max, shape) + tile = random.randint(args.tile_min, args.tile_max) + pixelate_region(img, poly, tile) + lines.append(poly_to_label(poly, W, H)) + stem = f"{src.stem}_{v:02d}" + imwrite(out / "images" / split / f"{stem}.jpg", img) + (out / "labels" / split / f"{stem}.txt").write_text("\n".join(lines), encoding="utf-8") + counts[split] += 1 + counts["neg" if not lines else "pos"] += 1 + + (out / "data.yaml").write_text( + f"path: {out.resolve().as_posix()}\n" + "train: images/train\n" + "val: images/val\n" + "names:\n 0: mosaic\n", + encoding="utf-8", + ) + print(f"Готово: train={counts['train']} val={counts['val']} " + f"(с мозаикой={counts['pos']}, чистых={counts['neg']})") + print(f"data.yaml: {(out / 'data.yaml').resolve()}") + + +if __name__ == "__main__": + main() diff --git a/scripts/training/train_mosaic.py b/scripts/training/train_mosaic.py new file mode 100644 index 0000000..2a0d076 --- /dev/null +++ b/scripts/training/train_mosaic.py @@ -0,0 +1,52 @@ +"""Train a YOLO11-seg mosaic detector on a generated synthetic dataset. + +Prereqs: pip install -e ".[yolo]" plus PyTorch (CUDA build recommended — see README). + +Usage: + python scripts/training/train_mosaic.py --data dataset_mosaic/data.yaml --epochs 100 + +The resulting weights (runs/segment//weights/best.pt) drop straight into the +app: Параметры → Детектор «YOLO» → выбрать этот .pt. The class is named "mosaic", +which YoloDetector maps to CensorType.MOSAIC. +""" + +from __future__ import annotations + +import argparse + + +def main() -> None: + ap = argparse.ArgumentParser(description="Train YOLO11-seg mosaic detector") + ap.add_argument("--data", required=True, help="path to data.yaml from the generator") + ap.add_argument("--model", default="yolo11n-seg.pt", help="base model (n/s/m...-seg)") + ap.add_argument("--epochs", type=int, default=100) + ap.add_argument("--imgsz", type=int, default=640) + ap.add_argument("--batch", default="-1", help="batch size (-1 = auto)") + ap.add_argument("--device", default=None, help="cuda / 0 / cpu (default: auto)") + ap.add_argument("--name", default="mosaic", help="run name under the project dir") + ap.add_argument("--project", default=None, help="output dir for runs (default: runs/segment)") + args = ap.parse_args() + + try: + from ultralytics import YOLO + except ImportError as exc: # pragma: no cover + raise SystemExit('Не установлен ultralytics: pip install -e ".[yolo]"') from exc + + batch = int(args.batch) if str(args.batch).lstrip("-").isdigit() else args.batch + model = YOLO(args.model) + results = model.train( + data=args.data, + epochs=args.epochs, + imgsz=args.imgsz, + batch=batch, + device=args.device, + name=args.name, + project=args.project, + ) + save_dir = getattr(results, "save_dir", "runs/segment/" + args.name) + print(f"\nГотово. Веса: {save_dir}/weights/best.pt") + print("Подключите их в приложении: Параметры → Детектор YOLO → выбрать best.pt") + + +if __name__ == "__main__": + main()