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