Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.
This commit is contained in:
Leonid Pershin
@@ -47,12 +47,15 @@ what it found.
|
||||
Keep this scope sharp:
|
||||
|
||||
- Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить")
|
||||
step was added later (user-requested): per-frame, on-demand, behind a `Restorer`
|
||||
interface. Two engines: a cv2 **inpaint baseline** (fills, does NOT reconstruct) and
|
||||
**DeepMosaics** — real generative mosaic removal, its GPL-3.0 network code **vendored**
|
||||
under `core/restore/_deepmosaics/` and run in-process (user supplies only the weights).
|
||||
Because of that vendoring the **whole project is GPL-3.0**. LADA (BasicVSR++) is a
|
||||
possible future engine, not wired.
|
||||
step was added later (user-requested): on-demand (single frame **or** whole-project
|
||||
batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
|
||||
restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
|
||||
composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were
|
||||
**removed** as "works poorly". DeepMosaics has two engines: **image** (per-frame) and
|
||||
**video** (BVDNet, temporal — uses neighbour frames). Its GPL-3.0 network code is
|
||||
**vendored** under `core/restore/_deepmosaics/` and run in-process (user supplies only
|
||||
the weights). Because of that vendoring the **whole project is GPL-3.0**. LADA
|
||||
(BasicVSR++) is a possible future engine, not wired.
|
||||
- Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was
|
||||
rejected early — a generative *conditioning* model, not a censorship restorer. Don't
|
||||
reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/
|
||||
@@ -67,8 +70,8 @@ Keep this scope sharp:
|
||||
|
||||
- **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.
|
||||
- **GPU:** NVIDIA + CUDA via PyTorch, for the YOLO detector and the DeepMosaics
|
||||
restorer. CPU fallback works but is slow (esp. DeepMosaics / the temporal BVDNet).
|
||||
|
||||
## Tech stack (decided)
|
||||
|
||||
@@ -76,12 +79,12 @@ Keep this scope sharp:
|
||||
|----------------|---------------------------------------------|
|
||||
| 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 |
|
||||
| Detector | Ultralytics YOLO (LADA weights) behind a pluggable interface (YOLO-only) |
|
||||
|
||||
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.
|
||||
Torch/CUDA + Ultralytics enter with the YOLO detector. Keep that dependency optional
|
||||
(the `yolo` extra in `pyproject.toml` pulls only Ultralytics; torch is installed
|
||||
separately per the README). Both detection (YOLO) and restoration (DeepMosaics) now
|
||||
require torch — there's no longer a torch-free detector.
|
||||
|
||||
## Architecture (as implemented)
|
||||
|
||||
@@ -115,21 +118,17 @@ hvideotool/
|
||||
├── video/
|
||||
│ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale
|
||||
│ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
|
||||
├── restore/ # "un-censor" detected regions (per-frame)
|
||||
│ ├── base.py # Restorer ABC: restore(image, detections, should_cancel=None) -> image; Cancelled exc
|
||||
│ ├── factory.py # build_restorer(name, config) -> inpaint | deepmosaics (lada = TODO)
|
||||
│ ├── inpaint.py # InpaintRestorer (cv2) — baseline, fills not reconstructs
|
||||
│ ├── deepmosaics.py # DeepMosaicsRestorer — in-process, loads models once; uses _deepmosaics/
|
||||
│ ├── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
|
||||
│ └── mask.py # detections_to_mask(shape, dets, dilate)
|
||||
└── detection/
|
||||
├── restore/ # "un-censor" (DeepMosaics only; per-frame OR temporal)
|
||||
│ ├── base.py # Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc
|
||||
│ ├── factory.py # build_restorer(name, config) -> deepmosaics | deepmosaics_video (lada = TODO)
|
||||
│ ├── deepmosaics.py # DeepMosaicsRestorer (image, per-frame) + DeepMosaicsVideoRestorer (BVDNet, temporal); in-process, load once; uses _deepmosaics/
|
||||
│ └── _deepmosaics/ # VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
|
||||
└── detection/ # YOLO only
|
||||
├── base.py # Detector ABC: detect(frame) -> list[Detection]
|
||||
├── factory.py # build_detector(config) -> classic | yolo | combined
|
||||
├── factory.py # build_detector(config) -> yolo (the only kind)
|
||||
├── types.py # Detection (+ to_dict/from_dict), CensorType enum
|
||||
├── cache.py # save/load the project detection cache (detections.json): cache_file + base_dir args
|
||||
├── 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
|
||||
└── yolo.py # YoloDetector — Ultralytics YOLO-seg; lazy-imports torch/ultralytics
|
||||
```
|
||||
|
||||
### How it works
|
||||
@@ -146,15 +145,21 @@ hvideotool/
|
||||
`engine.detect/restore` all run **inside the worker** (`_compute` is the pure
|
||||
read+detect helper); the `fn` must touch NO Qt widgets — it emits plain data that the
|
||||
GUI-thread slots (`_apply_detection`, restore `tick`) apply. `_begin_busy` disables
|
||||
`detector_combo`/`model_action` for the duration (they'd race the running detector).
|
||||
This is the deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics
|
||||
per-call latency can't be hidden with `processEvents`).
|
||||
- **Device badge.** A clickable status-bar chip (`device_badge`) shows "⚡ CUDA" (green)
|
||||
or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in a background
|
||||
`Job` at startup (importing torch is slow, so it's off the GUI thread) → `_set_device_badge`.
|
||||
Clicking (`_show_device_info`) opens a diagnostic dialog: `torch_info.reason()` explains
|
||||
why CPU (no torch / CPU-only `+cpu` build / built-with-CUDA-but-no-GPU) plus
|
||||
`install_hint()` (the cu121 pip command). `core/torch_info.py` is pure (no Qt).
|
||||
`model_action` for the duration (it'd race the running detector). This is the
|
||||
deliberate exception to the old single-threaded rule (CPU YOLO/DeepMosaics per-call
|
||||
latency can't be hidden with `processEvents`).
|
||||
- **Device badge + CUDA diagnostics.** A clickable status-bar chip (`device_badge`) shows
|
||||
"⚡ CUDA" (green) or "🖥 CPU" (orange). `_probe_device` runs `core/torch_info.gather()` in
|
||||
a background `Job` at startup (it imports torch AND shells out to `nvidia-smi`, so it's off
|
||||
the GUI thread) → `_set_device_badge`. `gather()` collects torch facts (version, built_cuda,
|
||||
cuda_available, device_name) **and** NVIDIA facts (gpus, driver_version, max cuda_driver).
|
||||
Clicking (`_show_device_info`) opens a diagnostic **QDialog** (not QMessageBox — its text
|
||||
wasn't copyable): a read-only monospace `QPlainTextEdit` with `torch_info.analyze(info)` →
|
||||
`{summary, details, steps, command}` — a verdict on *why* it's on CPU (CPU-only `+cpu`
|
||||
build / no GPU / driver-too-old-for-built-CUDA) and the exact pip fix. Buttons: "Скопировать
|
||||
команду установки" (`_copy_install_command` → the recommended cu121/cu118 command, picked by
|
||||
`recommend_channel` from the driver's CUDA) and "Проверить заново" (re-runs `_probe_device`).
|
||||
`core/torch_info.py` is pure (no Qt); subprocess uses `CREATE_NO_WINDOW` on Windows.
|
||||
- **Projects (`core/project.py`).** `MainWindow._project` is the open `Project`;
|
||||
`_folder` is kept as a synonym for `project.frames_dir` so the rest of the code
|
||||
(navigation, cache, tags) didn't need rewiring. Entry points: "Создать проект…"
|
||||
@@ -179,7 +184,7 @@ hvideotool/
|
||||
"Рассчитать кадр" 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`).
|
||||
Switching the model clears the cache (`_choose_model` → `_invalidate_results`).
|
||||
- **Detection cache (persisted).** `_results` is mirrored to the project's
|
||||
`detections.json` (`core/detection/cache.py`, at `project.cache_path`; keyed by
|
||||
**basename** so it survives moving the project). `cache.save_results`/`load_results`
|
||||
@@ -203,27 +208,43 @@ hvideotool/
|
||||
created lazily on first move), removing them from list/`_files`/cache. `_unique_dest`
|
||||
avoids clobbering (`foo.jpg` → `foo (1).jpg`). Use case: flag good frames into a
|
||||
training/example set while inspecting detections.
|
||||
- **Restoration ("Расцензурить кадр").** Toolbar action runs `self._restorer` (built via
|
||||
`build_restorer`) on the current frame's detections (computing them first if needed),
|
||||
**on a background job** (`_restore_current` builds an `fn` that detects-if-needed +
|
||||
restores in the worker; a `tick` caches freshly-computed detections, `done` stores
|
||||
`_restored[path]` + shows it). "Показать оригинал/
|
||||
результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
||||
`<stem>_restored.jpg` beside the frame. The baseline
|
||||
is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), run
|
||||
**in-process** from the vendored `_deepmosaics/` code: it loads the BiSeNet locator +
|
||||
clean generator **once** (lazy, cached on the instance) and per frame reproduces their
|
||||
`cleanmosaic_img_server` (locate mosaic → run generator on the crop → feather back),
|
||||
~0.3 s/frame cached on CPU vs ~7 s when it spawned a subprocess. Use the **image**
|
||||
model `clean_youknow_resnet_9blocks.pth` — the video model (BVDNet) is rejected per
|
||||
frame (needs a neighbour). `should_cancel` (= `lambda: job.cancelled`) is polled so
|
||||
"■ Стоп" stops it; the engine raises `Cancelled`, which `Job.run` reports as a clean cancel.
|
||||
The engine + weights are set in `RestoreDialog` (Файл → Движок восстановления…),
|
||||
persisted, and built lazily/cached in `_make_restorer` (like `_make_detector`). NOTE:
|
||||
DeepMosaics locates mosaics itself (its `mosaic_position.pth`, expected beside the
|
||||
clean weights) — our detections aren't passed to it. `_show` resets `_showing_restored`
|
||||
+ `_update_restore_actions`. To add another engine (e.g. LADA), implement
|
||||
`core/restore/base.Restorer` and register it in `restore/factory.build_restorer`.
|
||||
- **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it
|
||||
**locates the mosaic itself** — so restoration is **fully decoupled from detection**: no
|
||||
detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar
|
||||
action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on
|
||||
a background job** (`_restore_current`; `done` stores `_restored[path]` + shows it). "Показать
|
||||
оригинал/результат" toggles (`_showing_restored`); "Сохранить результат" writes
|
||||
`<stem>_restored.jpg` beside the frame. **Batch ("Расцензурить все" / "Все заново",
|
||||
`_restore_all(force)`)** mirrors `_detect_all`: a single background job restores every
|
||||
frame and writes results to the project's **`restored/`** dir (`Project.restored_dir`,
|
||||
basename-mirrored, kept OUT of `frames/` so outputs aren't re-listed/re-restored); the
|
||||
per-frame engine **skips frames already in `restored/`** unless `force` (resume). The
|
||||
engines are **DeepMosaics** (`restore/deepmosaics.py`), run **in-process** from the
|
||||
vendored `_deepmosaics/` code, loading the BiSeNet locator + generator **once** (lazy,
|
||||
cached on the instance):
|
||||
- `deepmosaics` (image, per-frame): reproduces `cleanmosaic_img_server` (locate mosaic
|
||||
→ run generator on the crop → feather back), ~0.3 s/frame cached on CPU. Image model
|
||||
`clean_youknow_resnet_9blocks.pth`.
|
||||
- `deepmosaics_video` (**temporal, BVDNet**): `DeepMosaicsVideoRestorer`, `.temporal=True`.
|
||||
Reproduces `cleanmosaic_video_fusion` — per target frame it feeds the net a window of
|
||||
`T=5` neighbour frames sampled at step `S=3` around it (`N=2` each side, clamped at the
|
||||
sequence edges) **plus its own previous output (recurrent)**, for temporal coherence.
|
||||
Because of that recurrence it must run a **contiguous, ordered range** — it implements
|
||||
`restore_sequence(count, get_frame, get_dets, emit, should_cancel)` (the batch run uses
|
||||
it; single-frame `restore` degrades to a window of the same frame). Needs the **video**
|
||||
weights `clean_youknow_video.pth` (+ `mosaic_position.pth` beside). `INPUT_SIZE=256`.
|
||||
|
||||
`restore_sequence` is on the `Restorer` ABC (default = independent per-frame loop);
|
||||
`_restore_all` dispatches on `restorer.temporal` (temporal → `restore_sequence` over the
|
||||
whole range; per-frame → resumable loop with skip-existing). `should_cancel`
|
||||
(= `lambda: job.cancelled`) is polled so "■ Стоп" stops it; engines raise `Cancelled`,
|
||||
which `Job.run` reports as a clean cancel. Engine + weights are set in `RestoreDialog`
|
||||
(Файл → Движок восстановления…) — the model dropdown shows image vs video weights per
|
||||
selected engine — persisted, and built lazily/cached in `_make_restorer`. NOTE:
|
||||
DeepMosaics locates mosaics itself (its `mosaic_position.pth`) — our detections aren't
|
||||
passed to it. To add another engine (e.g. LADA), implement `core/restore/base.Restorer`
|
||||
(set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
|
||||
`restore/factory.build_restorer`.
|
||||
- **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys
|
||||
`,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label`
|
||||
("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`,
|
||||
@@ -255,20 +276,20 @@ hvideotool/
|
||||
- `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`.
|
||||
- Detection is YOLO-only and restoration is DeepMosaics-only. If you re-add an engine
|
||||
kind, implement `core/detection/base.Detector` / `core/restore/base.Restorer`, register
|
||||
the string in the respective `factory`, and add it to `DETECTORS`/`RESTORERS` in
|
||||
`config.py` (and `normalize_config`). There's no detector dropdown anymore — the toolbar
|
||||
just shows "Детектор: YOLO"; restoration engines are chosen in `RestoreDialog`.
|
||||
|
||||
## Commands
|
||||
|
||||
```powershell
|
||||
python -m venv .venv; .\.venv\Scripts\Activate.ps1
|
||||
pip install -e . # classic detector needs no torch/CUDA
|
||||
pip install -e ".[yolo]" # YOLO needs ultralytics; install torch separately (README)
|
||||
|
||||
python -m hvideotool # reopen the last project (or create/open one in-app)
|
||||
python -m hvideotool "C:\path\to\MyProject" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
|
||||
pip install -e ".[yolo]" # + install torch separately, see README
|
||||
python -m hvideotool "C:\path\to\MyProject" --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
```
|
||||
|
||||
No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a
|
||||
@@ -292,24 +313,22 @@ frame directly.
|
||||
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.
|
||||
**Switching to yolo/combined without a model auto-picks one** via
|
||||
`MainWindow._auto_find_model()`: it scans `./models/**.pt` and matches only filenames
|
||||
containing `lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; it
|
||||
falls back to the "Модель…" file dialog only when nothing suitable is found.
|
||||
- **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.
|
||||
**A model is auto-picked on project open** via `MainWindow._ensure_model()` →
|
||||
`_auto_find_model()`: it scans `./models/**.pt` and matches only filenames containing
|
||||
`lada`/`mosaic` (so it skips the COCO `yolo11n-seg.pt` trap), no prompt; otherwise the
|
||||
user picks via "Модель…" (`_choose_model`).
|
||||
- **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was
|
||||
noisy/approximate on real footage (false positives on skin/hair/fabric/JPEG; missed
|
||||
real mosaic after downscale) and the `combined` mode + cv2 `inpaint` baseline went with
|
||||
it. Detection is YOLO-only, restoration is DeepMosaics-only. `normalize_config` coerces
|
||||
any leftover `classic`/`combined`/`inpaint` in old settings/projects to `yolo`/`deepmosaics`.
|
||||
- **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.
|
||||
but not all. The real anime fix is *retraining* a YOLO11-seg (see `scripts/training/`).
|
||||
- **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`.
|
||||
mosaic only. 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__`).
|
||||
@@ -318,10 +337,11 @@ frame directly.
|
||||
- **CPU-only torch must not request CUDA.** A `+cpu` torch build raises "Torch not
|
||||
compiled with CUDA enabled" the moment something calls `.cuda()`. Both engines guard
|
||||
for this: `YoloDetector` picks `cuda` only when `torch.cuda.is_available()` (even an
|
||||
explicit `yolo_device="cuda"` is downgraded to cpu); `DeepMosaicsRestorer._ensure_loaded`
|
||||
forces `gpu_id="-1"` when CUDA is absent (its vendored `model_util.todevice` /
|
||||
`data.im2tensor` call `.cuda()` for any `gpu_id != "-1"`, e.g. the `dm_gpu="0"` default).
|
||||
So a wrong/CPU-only torch falls back to CPU instead of crashing.
|
||||
explicit `yolo_device="cuda"` is downgraded to cpu); both `DeepMosaicsRestorer._ensure_loaded`
|
||||
and `DeepMosaicsVideoRestorer._ensure_loaded` force `gpu_id="-1"` when CUDA is absent (the
|
||||
vendored `model_util.todevice` / `data.im2tensor`/`to_tensor` call `.cuda()` for any
|
||||
`gpu_id != "-1"`, e.g. the `dm_gpu="0"` default). So a wrong/CPU-only torch falls back to
|
||||
CPU instead of crashing (the temporal BVDNet engine is heavy on CPU, though).
|
||||
- **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 —
|
||||
|
||||
@@ -64,9 +64,10 @@
|
||||
кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры
|
||||
с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена,
|
||||
🟢 проверено и чисто).
|
||||
- Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога**
|
||||
уверенности прямо в тулбаре — удобно сравнивать.
|
||||
- Выбор файла весов модели кнопкой **«Модель…»**.
|
||||
- Детекция — через **YOLO** (модель LADA для мозаики). Порог уверенности
|
||||
настраивается прямо в тулбаре.
|
||||
- Выбор файла весов модели кнопкой **«Модель…»** (нужная модель ищется в `models/`
|
||||
автоматически при открытии проекта).
|
||||
- **Индикатор устройства** в строке состояния: «⚡ CUDA» или «🖥 CPU». Клик по «CPU»
|
||||
показывает диагностику (почему GPU не задействован) и команды установки PyTorch с
|
||||
CUDA. Если CUDA недоступна, YOLO и DeepMosaics автоматически работают на CPU
|
||||
@@ -96,14 +97,24 @@
|
||||
`<имя>_restored.jpg` рядом с кадром. Движок выбирается в
|
||||
меню **Файл → Движок восстановления…**.
|
||||
|
||||
Два движка:
|
||||
**Пакетный прогон по диапазону.** Кнопки **«Расцензурить все»** (дозапуск — пропускает
|
||||
уже сделанные) и **«Все заново»** обрабатывают **все кадры проекта** в фоне и пишут
|
||||
результаты в папку **`restored/`** проекта (имена кадров сохраняются; папка держится
|
||||
отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс,
|
||||
отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции.
|
||||
|
||||
- **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по
|
||||
окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает).
|
||||
- **DeepMosaics** — реальное генеративное удаление мозаики. Код **встроен** в
|
||||
приложение (vendored, GPL-3.0), ставить его отдельно не нужно — требуются только
|
||||
**веса** и (желательно) **GPU NVIDIA/CUDA**. На аниме качество ограничено (модели
|
||||
обучены на реальном видео).
|
||||
Расцензуривание — только через **DeepMosaics** (код **встроен** в приложение, vendored,
|
||||
GPL-3.0; ставить отдельно не нужно — требуются только **веса** и желательно **GPU
|
||||
NVIDIA/CUDA**). Два движка:
|
||||
|
||||
- **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**.
|
||||
- **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры**
|
||||
(окно ±2 кадра с шагом 3) и собственный предыдущий результат для когерентности на
|
||||
роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е.
|
||||
запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к
|
||||
окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`.
|
||||
|
||||
На аниме качество ограничено (модели обучены на реальном видео).
|
||||
|
||||
### Настройка DeepMosaics
|
||||
|
||||
@@ -112,18 +123,20 @@
|
||||
([Google Drive](https://drive.google.com/drive/folders/1LTERcN33McoiztYEwBxMuRjjgxh4DEPs),
|
||||
Baidu код `1x0a`):
|
||||
|
||||
- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель;
|
||||
- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом).
|
||||
- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель (движок «картинка»);
|
||||
- **`clean_youknow_video.pth`** — видеомодель BVDNet (движок «видео», соседние кадры);
|
||||
- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом, нужен обоим).
|
||||
|
||||
Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, выберите **модель
|
||||
из выпадающего списка** (наполняется из `models/deepmosaics`; есть «Обзор…» для файла в
|
||||
другом месте) и GPU id (`-1` = CPU). Если веса в `models/deepmosaics` — работает сразу;
|
||||
модель грузится один раз, дальше кадры считаются быстро.
|
||||
Затем в приложении: **Файл → Движок восстановления…**, выберите движок (**DeepMosaics —
|
||||
картинка** или **видео**) и **модель из выпадающего списка** (наполняется из
|
||||
`models/deepmosaics` — для видеодвижка показываются только `clean_*_video.pth`; есть
|
||||
«Обзор…» для файла в другом месте) и GPU id (`-1` = CPU). Если веса в `models/deepmosaics`
|
||||
— работает сразу; модель грузится один раз, дальше кадры считаются быстро.
|
||||
|
||||
> **Важно:** берите именно **картиночную** модель `clean_youknow_resnet_9blocks.pth`.
|
||||
> Видеомодель `clean_youknow_video.pth` (BVDNet) покадрово **не работает** — ей нужен
|
||||
> соседний кадр (приложение это распознаёт и подскажет). Если на кадре нет мозаики,
|
||||
> результат = исходный кадр.
|
||||
> **Какую модель брать:** для покадрового движка — **картиночную**
|
||||
> `clean_youknow_resnet_9blocks.pth`; для временно́го — **видео** `clean_youknow_video.pth`
|
||||
> (она запускается только пакетно, «Расцензурить все», т.к. ей нужны соседние кадры). Если
|
||||
> на кадре нет мозаики, результат = исходный кадр.
|
||||
>
|
||||
> Код DeepMosaics (GPL-3.0) лежит в `core/restore/_deepmosaics/` и поэтому **весь
|
||||
> проект распространяется под GPL-3.0**. Запускается на современных `torch 2.x`/
|
||||
@@ -143,8 +156,8 @@ Baidu код `1x0a`):
|
||||
|
||||
- **ОС:** Windows 11 x64 (основная целевая платформа).
|
||||
- **Python:** 3.11+.
|
||||
- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора. CPU-режим работает, но
|
||||
медленный. Для `classic` детектора ни torch, ни GPU не нужны.
|
||||
- **GPU (опционально):** NVIDIA + CUDA для YOLO-детектора и DeepMosaics. CPU-режим
|
||||
работает, но медленный (особенно DeepMosaics / временно́й BVDNet).
|
||||
|
||||
## Установка
|
||||
|
||||
@@ -153,12 +166,11 @@ git clone https://github.com/mrleo1nid/HVideoTool.git
|
||||
cd HVideoTool
|
||||
python -m venv .venv
|
||||
.\.venv\Scripts\Activate.ps1
|
||||
pip install -e .
|
||||
pip install -e ".[yolo]"
|
||||
```
|
||||
|
||||
Этого достаточно для `classic` детектора — **PyTorch/CUDA не требуются**. Они
|
||||
нужны только для YOLO/комбинированного детектора (см.
|
||||
[Модель детектора](#модель-детектора)).
|
||||
Детекция (YOLO) и расцензуривание (DeepMosaics) требуют **PyTorch** — установите его
|
||||
отдельно под вашу CUDA (см. [Модель детектора](#модель-детектора)).
|
||||
|
||||
## Запуск
|
||||
|
||||
@@ -166,29 +178,23 @@ pip install -e .
|
||||
# Без аргументов — открывается последний проект (или создайте/откройте новый в тулбаре)
|
||||
python -m hvideotool
|
||||
|
||||
# Необязательно: сразу открыть проект / переопределить детектор и модель (по умолчанию)
|
||||
python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
# Необязательно: сразу открыть проект и указать модель YOLO
|
||||
python -m hvideotool "C:\path\to\МойПроект" --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
```
|
||||
|
||||
Выбор детектора, путь к модели и порог сохраняются в `~/HVideoTool/settings.json`
|
||||
и применяются при следующем запуске.
|
||||
Путь к модели и порог сохраняются в `~/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** (интерфейс абстрагирован в `core/detection/base.py`):
|
||||
ML-детектор на базе [Ultralytics](https://github.com/ultralytics/ultralytics)
|
||||
(`core/detection/yolo.py`), **сегментационная** модель — маски превращаются в
|
||||
контуры. Рекомендуемые веса — [**LADA mosaic detection**](https://huggingface.co/ladaapp/lada).
|
||||
(Старый эвристический classic-CV детектор и комбинированный режим удалены — давали
|
||||
много ложных срабатываний.)
|
||||
|
||||
> **⚠️ Берите правильную модель.** Для YOLO нужна модель **детекции цензуры**
|
||||
> (LADA `lada_mosaic_detection_model_v4_accurate.pt`). Если по ошибке указать
|
||||
@@ -207,7 +213,7 @@ python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model mod
|
||||
```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
|
||||
# затем: тулбар → «Модель…» → runs\segment\mosaic\weights\best.pt
|
||||
```
|
||||
|
||||
Установка YOLO-детектора:
|
||||
@@ -219,11 +225,11 @@ 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
|
||||
python -m hvideotool --model models\lada_mosaic_detection_model_v4_accurate.pt
|
||||
```
|
||||
|
||||
> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**. Classic-CV детектор
|
||||
> от этого свободен.
|
||||
> **⚠️ Лицензия.** Ultralytics YOLO и веса LADA — **AGPL-3.0**; код DeepMosaics —
|
||||
> **GPL-3.0**. Поэтому весь проект распространяется под **GPL-3.0**.
|
||||
|
||||
---
|
||||
|
||||
@@ -243,14 +249,17 @@ hvideotool/
|
||||
├── imageio.py # unicode-safe чтение/запись картинок (Windows-пути)
|
||||
├── project.py # Project: раскладка (project.json/frames/detections.json/collections) + настройки
|
||||
├── 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
|
||||
├── cache.py # сохранение/загрузка кэша детекций (detections.json в проекте)
|
||||
├── classic_cv.py # эвристический детектор (mosaic/blur/black_bar)
|
||||
├── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны)
|
||||
└── composite.py # CompositeDetector: объединение детекторов
|
||||
├── detection/ # только YOLO
|
||||
│ ├── base.py # Detector (ABC): detect(frame) -> list[Detection]
|
||||
│ ├── factory.py # build_detector(config) -> yolo
|
||||
│ ├── types.py # Detection (+ to_dict/from_dict), CensorType
|
||||
│ ├── cache.py # сохранение/загрузка кэша детекций (detections.json)
|
||||
│ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны)
|
||||
└── restore/ # только DeepMosaics
|
||||
├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal
|
||||
├── factory.py # build_restorer -> deepmosaics | deepmosaics_video
|
||||
├── deepmosaics.py # движки «картинка» (покадрово) и «видео» (BVDNet)
|
||||
└── _deepmosaics/ # встроенный код DeepMosaics (GPL-3.0)
|
||||
```
|
||||
|
||||
Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU
|
||||
@@ -264,9 +273,9 @@ hvideotool/
|
||||
| Язык | Python ≥ 3.11 |
|
||||
| GUI | PySide6 (Qt 6) |
|
||||
| Обработка картинок | OpenCV / NumPy |
|
||||
| Детектор (без весов) | classic-CV эвристика (без GPU) |
|
||||
| Детектор (ML) | Ultralytics YOLO + PyTorch/CUDA (LADA) |
|
||||
| Детектор | Ultralytics YOLO + PyTorch/CUDA (веса LADA) |
|
||||
| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA |
|
||||
|
||||
## Лицензия
|
||||
|
||||
TBD.
|
||||
**GPL-3.0** — из-за встроенного кода DeepMosaics (GPL-3.0); веса/код YOLO LADA — AGPL-3.0.
|
||||
|
||||
@@ -12,7 +12,7 @@ import sys
|
||||
|
||||
from . import settings_store
|
||||
from .app import run
|
||||
from .config import AppConfig
|
||||
from .config import AppConfig, normalize_config
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -21,15 +21,13 @@ def main() -> int:
|
||||
description="Инспектор детекции уже наложенной цензуры на картинках.",
|
||||
)
|
||||
parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)")
|
||||
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 defaults first
|
||||
normalize_config(config) # drop any legacy classic/inpaint values
|
||||
|
||||
if args.detector:
|
||||
config.detector = args.detector
|
||||
if args.model_path:
|
||||
config.model_path = args.model_path
|
||||
|
||||
|
||||
+21
-27
@@ -1,40 +1,21 @@
|
||||
"""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.
|
||||
Plain dataclasses. Detection is YOLO-only and restoration is DeepMosaics-only, so the
|
||||
knobs here are the YOLO inference params, the overlay style, and the DeepMosaics weights.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
DETECTORS = ("yolo",)
|
||||
RESTORERS = ("deepmosaics", "deepmosaics_video")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DetectionConfig:
|
||||
"""Parameters for the detector. Thresholds tuned for the classic-CV detector."""
|
||||
"""Parameters for the YOLO 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")
|
||||
@@ -62,13 +43,26 @@ class OverlayConfig:
|
||||
class AppConfig:
|
||||
detection: DetectionConfig = field(default_factory=DetectionConfig)
|
||||
overlay: OverlayConfig = field(default_factory=OverlayConfig)
|
||||
detector: str = "classic" # "classic" | "yolo" | "combined"
|
||||
detector: str = "yolo" # only "yolo"
|
||||
model_path: str | None = None # weights path, used by the YOLO detector
|
||||
default_threshold: float = 0.20 # initial overlay confidence threshold
|
||||
|
||||
# --- restoration ("расцензурить") ---
|
||||
restorer: str = "inpaint" # "inpaint" | "deepmosaics"
|
||||
restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
|
||||
dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py)
|
||||
dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
|
||||
dm_python: str | None = None # python exe for DeepMosaics (None = current)
|
||||
dm_gpu: str = "0" # CUDA device id, "-1" for CPU
|
||||
|
||||
|
||||
def normalize_config(cfg: AppConfig) -> None:
|
||||
"""Coerce legacy/removed settings to supported values (mutates ``cfg``).
|
||||
|
||||
Old projects / settings.json may carry the removed ``classic``/``combined``
|
||||
detectors or the ``inpaint`` restorer — map those onto the survivors so loading
|
||||
them doesn't blow up at build time.
|
||||
"""
|
||||
if cfg.detector not in DETECTORS:
|
||||
cfg.detector = "yolo"
|
||||
if cfg.restorer not in RESTORERS:
|
||||
cfg.restorer = "deepmosaics"
|
||||
|
||||
@@ -1,216 +0,0 @@
|
||||
"""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
|
||||
@@ -1,51 +0,0 @@
|
||||
"""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
|
||||
@@ -3,14 +3,15 @@
|
||||
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.
|
||||
|
||||
Only the YOLO detector is supported — the classic-CV heuristic (and the composite
|
||||
mode that combined them) were removed: they were noisy/approximate on real footage.
|
||||
"""
|
||||
|
||||
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:
|
||||
@@ -23,19 +24,6 @@ def _require_model(config: AppConfig) -> str:
|
||||
|
||||
|
||||
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
|
||||
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}")
|
||||
return YoloDetector(_require_model(config), config.detection)
|
||||
|
||||
@@ -7,7 +7,7 @@ a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose cl
|
||||
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.
|
||||
rest of the app never pulls them in until detection actually runs.
|
||||
|
||||
Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
|
||||
"""
|
||||
|
||||
@@ -30,6 +30,7 @@ FRAMES_DIR = "frames"
|
||||
CACHE_FILE = "detections.json"
|
||||
COLLECTIONS_DIR = "collections"
|
||||
FAVORITES_DIR = "Избранное" # the single default collection ("в избранное")
|
||||
RESTORED_DIR = "restored" # batch "расцензурить все" output (kept out of frames/)
|
||||
_VERSION = 1
|
||||
|
||||
# The subset of AppConfig fields a project remembers (mirrored to/from project.json).
|
||||
@@ -78,6 +79,12 @@ class Project:
|
||||
"""The single default collection — frames moved "to favorites" land here."""
|
||||
return self.collections_dir / FAVORITES_DIR
|
||||
|
||||
@property
|
||||
def restored_dir(self) -> Path:
|
||||
"""Batch restoration output ("расцензурить все") — mirrors frame basenames.
|
||||
Kept out of ``frames/`` so results aren't listed/re-detected/re-restored."""
|
||||
return self.root / RESTORED_DIR
|
||||
|
||||
# ------------------------------------------------------------- lifecycle
|
||||
@classmethod
|
||||
def create(
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
"""Restorer interface — "un-censor" detected regions of an image.
|
||||
|
||||
A Restorer takes an image plus the detected censored regions and returns a new
|
||||
image with those regions reconstructed/filled. This mirrors the ``Detector``
|
||||
abstraction so different engines (classic inpaint now; a generative model like
|
||||
DeepMosaics / LADA later) plug in behind the same interface.
|
||||
image with those regions reconstructed. This mirrors the ``Detector`` abstraction
|
||||
so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug
|
||||
in behind the same interface. Note DeepMosaics locates the mosaic itself, so the
|
||||
``detections`` argument is currently advisory (unused by the DeepMosaics engines).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -24,7 +25,18 @@ class Cancelled(Exception):
|
||||
"""Raised by a Restorer when ``should_cancel`` asked it to stop."""
|
||||
|
||||
|
||||
# Sequence (batch) restore callbacks — see ``Restorer.restore_sequence``.
|
||||
FrameGetter = Callable[[int], np.ndarray] # index -> BGR image
|
||||
DetGetter = Callable[[int], list[Detection]] # index -> that frame's detections
|
||||
ResultSink = Callable[[int, np.ndarray], None] # (index, restored image) -> None
|
||||
|
||||
|
||||
class Restorer(ABC):
|
||||
#: Whether this engine uses *neighbouring* frames (so a batch run must feed it a
|
||||
#: contiguous, ordered sequence — see :meth:`restore_sequence`). Per-frame engines
|
||||
#: leave this False; the temporal DeepMosaics (BVDNet) sets it True.
|
||||
temporal: bool = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return type(self).__name__
|
||||
@@ -42,3 +54,24 @@ class Restorer(ABC):
|
||||
True the engine should abort and raise :class:`Cancelled`.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def restore_sequence(
|
||||
self,
|
||||
count: int,
|
||||
get_frame: FrameGetter,
|
||||
get_dets: DetGetter,
|
||||
emit: ResultSink,
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> None:
|
||||
"""Restore ``count`` frames *in order*, calling ``emit(i, restored)`` for each.
|
||||
|
||||
The default treats every frame independently (just loops :meth:`restore`).
|
||||
Temporal engines override this to pull neighbouring frames via ``get_frame``
|
||||
and carry recurrent state across the sequence. ``get_frame``/``get_dets`` are
|
||||
lazy so the engine only reads the frames it needs; ``emit`` lets the caller
|
||||
stream results to disk instead of holding them all in memory.
|
||||
"""
|
||||
for i in range(count):
|
||||
if should_cancel is not None and should_cancel():
|
||||
raise Cancelled("Восстановление отменено")
|
||||
emit(i, self.restore(get_frame(i), get_dets(i), should_cancel))
|
||||
|
||||
@@ -26,18 +26,29 @@ from types import SimpleNamespace
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Cancelled, Restorer
|
||||
from .base import (
|
||||
CancelCheck,
|
||||
Cancelled,
|
||||
DetGetter,
|
||||
FrameGetter,
|
||||
ResultSink,
|
||||
Restorer,
|
||||
)
|
||||
|
||||
_VENDOR = Path(__file__).parent / "_deepmosaics"
|
||||
# Default place to drop DeepMosaics clean weights (gitignored — see models/).
|
||||
DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"
|
||||
|
||||
|
||||
def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
|
||||
"""Find usable per-frame clean models: (display_name, full_path).
|
||||
def discover_models(
|
||||
extra_dir: str | None = None, include_video: bool = False
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Find usable clean models: (display_name, full_path).
|
||||
|
||||
Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given)
|
||||
for ``clean_*.pth``. The video model is skipped — it can't run per-frame.
|
||||
for ``clean_*.pth``. By default the video model is skipped — it can't run
|
||||
per-frame; pass ``include_video=True`` for the temporal (BVDNet) engine, which
|
||||
*needs* ``clean_*_video.pth``.
|
||||
"""
|
||||
dirs = [DEFAULT_WEIGHTS_DIR]
|
||||
if extra_dir:
|
||||
@@ -48,13 +59,23 @@ def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
|
||||
if not d.is_dir():
|
||||
continue
|
||||
for p in sorted(d.glob("clean_*.pth")):
|
||||
if "video" in p.name.lower() or p.name in seen:
|
||||
if p.name in seen:
|
||||
continue
|
||||
if "video" in p.name.lower() and not include_video:
|
||||
continue
|
||||
seen.add(p.name)
|
||||
out.append((p.stem, str(p)))
|
||||
return out
|
||||
|
||||
|
||||
def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
|
||||
"""Locate ``mosaic_position.pth`` (the BiSeNet mosaic locator) for ``model``."""
|
||||
candidates = [model.parent / "mosaic_position.pth"]
|
||||
if dm_dir:
|
||||
candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
|
||||
return next((p for p in candidates if p.is_file()), None)
|
||||
|
||||
|
||||
def _netg_kind(model_name: str) -> str:
|
||||
"""Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
|
||||
n = model_name.lower()
|
||||
@@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer):
|
||||
)
|
||||
model = Path(model_path)
|
||||
self._netg = _netg_kind(model.name) # raises on a video model
|
||||
pos = self._find_mosaic_position(model, deepmosaics_dir)
|
||||
pos = _find_mosaic_position(model, deepmosaics_dir)
|
||||
if pos is None:
|
||||
raise ValueError(
|
||||
"Рядом с clean-моделью не найден mosaic_position.pth.\n"
|
||||
@@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer):
|
||||
self._gpu = gpu_id
|
||||
self._loaded = False # models loaded lazily on first restore
|
||||
|
||||
@staticmethod
|
||||
def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
|
||||
candidates = [model.parent / "mosaic_position.pth"]
|
||||
if dm_dir:
|
||||
candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
|
||||
return next((p for p in candidates if p.is_file()), None)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"DeepMosaics(gpu={self._gpu})"
|
||||
@@ -167,3 +181,147 @@ class DeepMosaicsRestorer(Restorer):
|
||||
img_mosaic = work[y - size:y + size, x - size:x + size]
|
||||
img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
|
||||
return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather)
|
||||
|
||||
|
||||
class DeepMosaicsVideoRestorer(Restorer):
|
||||
"""Temporal DeepMosaics (BVDNet) — un-censors using *neighbouring* frames.
|
||||
|
||||
Reproduces DeepMosaics' ``cleanmosaic_video_fusion`` per target frame: for frame
|
||||
``i`` it feeds the network a temporal window of ``T`` frames (sampled at step ``S``
|
||||
around ``i``) plus its own previous output (recurrent), so the reconstruction is
|
||||
temporally coherent. Because of that recurrence the frames MUST be processed in
|
||||
order over a contiguous range — see :meth:`restore_sequence` (the batch run).
|
||||
|
||||
Needs the **video** weights ``clean_youknow_video.pth`` + ``mosaic_position.pth``
|
||||
(beside it). Single-frame :meth:`restore` degrades to a window of the same frame.
|
||||
"""
|
||||
|
||||
temporal = True
|
||||
|
||||
# DeepMosaics fusion window: N before/after at step S → T = 2N+1 frames, INPUT_SIZE px.
|
||||
_N, _T, _S = 2, 5, 3
|
||||
_INPUT_SIZE = 256
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
deepmosaics_dir: str | None,
|
||||
model_path: str | None,
|
||||
python_exe: str | None = None, # unused (in-process); kept for factory parity
|
||||
gpu_id: str = "0",
|
||||
) -> None:
|
||||
chosen: Path | None = None
|
||||
if model_path and Path(model_path).is_file() and "video" in Path(model_path).name.lower():
|
||||
chosen = Path(model_path)
|
||||
else: # configured model missing or not a video model → auto-pick a video model
|
||||
vids = [p for _n, p in discover_models(include_video=True) if "video" in Path(p).name.lower()]
|
||||
if vids:
|
||||
chosen = Path(vids[0])
|
||||
if chosen is None:
|
||||
raise ValueError(
|
||||
"Не найдены веса видеомодели DeepMosaics (clean_*_video.pth).\n"
|
||||
"Положите clean_youknow_video.pth + mosaic_position.pth в models/deepmosaics "
|
||||
"(или выберите в «Восстановление…»). См. README."
|
||||
)
|
||||
pos = _find_mosaic_position(chosen, deepmosaics_dir)
|
||||
if pos is None:
|
||||
raise ValueError(
|
||||
"Рядом с видеомоделью не найден mosaic_position.pth.\n"
|
||||
"Положите mosaic_position.pth в ту же папку, что и clean_*_video.pth. См. README."
|
||||
)
|
||||
self._model = str(chosen)
|
||||
self._pos = str(pos)
|
||||
self._gpu = gpu_id
|
||||
self._loaded = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"DeepMosaicsVideo(gpu={self._gpu})"
|
||||
|
||||
# ------------------------------------------------------------------ engine
|
||||
def _ensure_loaded(self) -> None:
|
||||
if self._loaded:
|
||||
return
|
||||
if str(_VENDOR) not in sys.path:
|
||||
sys.path.insert(0, str(_VENDOR))
|
||||
|
||||
import torch # noqa: E402
|
||||
if self._gpu != "-1" and not torch.cuda.is_available():
|
||||
self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why)
|
||||
|
||||
from models import loadmodel, runmodel # type: ignore # noqa: E402
|
||||
import util.data as data # type: ignore # noqa: E402
|
||||
import util.image_processing as impro # type: ignore # noqa: E402
|
||||
|
||||
self._torch = torch
|
||||
self._runmodel = runmodel
|
||||
self._data = data
|
||||
self._impro = impro
|
||||
self._opt = SimpleNamespace(
|
||||
gpu_id=self._gpu,
|
||||
model_path=self._model,
|
||||
mosaic_position_model_path=self._pos,
|
||||
mask_threshold=64,
|
||||
all_mosaic_area=False,
|
||||
ex_mult=1.5,
|
||||
no_feather=False,
|
||||
)
|
||||
self._netM = loadmodel.bisenet(self._opt, "mosaic")
|
||||
self._netG = loadmodel.video(self._opt) # BVDNet
|
||||
self._loaded = True
|
||||
|
||||
def restore(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
detections: list[Detection],
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Single-frame restore — no neighbours, so the window is the same frame."""
|
||||
out: dict[int, np.ndarray] = {}
|
||||
self.restore_sequence(
|
||||
1,
|
||||
lambda _i: image,
|
||||
lambda _i: detections,
|
||||
lambda i, r: out.__setitem__(i, r),
|
||||
should_cancel,
|
||||
)
|
||||
return out.get(0, image.copy())
|
||||
|
||||
def restore_sequence(
|
||||
self,
|
||||
count: int,
|
||||
get_frame: FrameGetter,
|
||||
get_dets: DetGetter,
|
||||
emit: ResultSink,
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> None:
|
||||
self._ensure_loaded()
|
||||
torch, data, impro, opt = self._torch, self._data, self._impro, self._opt
|
||||
N, T, S, SZ = self._N, self._T, self._S, self._INPUT_SIZE
|
||||
|
||||
previous = None # recurrent state: the network's previous output (a tensor)
|
||||
for i in range(count):
|
||||
if should_cancel is not None and should_cancel():
|
||||
raise Cancelled("Восстановление отменено")
|
||||
img_origin = get_frame(i)
|
||||
x, y, size, mask = self._runmodel.get_mosaic_position(img_origin, self._netM, opt)
|
||||
if size <= 50:
|
||||
emit(i, img_origin.copy()) # no mosaic here; recurrence carries over
|
||||
continue
|
||||
|
||||
stream = []
|
||||
for k in range(T):
|
||||
j = min(max(i + (k - N) * S, 0), count - 1) # clamp window to range edges
|
||||
frame = img_origin if j == i else get_frame(j)
|
||||
crop = frame[y - size:y + size, x - size:x + size]
|
||||
stream.append(impro.resize(crop, SZ)[:, :, ::-1]) # BGR→RGB, SZ×SZ
|
||||
|
||||
if previous is None: # seed recurrence with the (centre) input crop
|
||||
previous = data.im2tensor(stream[N], bgr2rgb=False, gpu_id=opt.gpu_id)
|
||||
|
||||
arr = np.array(stream).reshape(1, T, SZ, SZ, 3).transpose((0, 4, 1, 2, 3))
|
||||
tensor = data.to_tensor(data.normalize(arr), gpu_id=opt.gpu_id)
|
||||
with torch.no_grad():
|
||||
pred = self._netG(tensor, previous)
|
||||
previous = pred
|
||||
img_fake = data.tensor2im(pred, rgb2bgr=True)
|
||||
emit(i, impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather))
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
"""Restorer factory: build a Restorer from the app config.
|
||||
|
||||
- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct).
|
||||
- ``deepmosaics``: real generative mosaic removal. The DeepMosaics network code is
|
||||
vendored (``_deepmosaics/``, GPL-3.0) and run in-process; the user supplies only the
|
||||
clean weights (+ ``mosaic_position.pth`` alongside). A CUDA GPU is recommended.
|
||||
Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but
|
||||
did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``,
|
||||
GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth``
|
||||
alongside). A CUDA GPU is recommended.
|
||||
|
||||
- ``deepmosaics``: per-frame generative mosaic removal (image model).
|
||||
- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for
|
||||
coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,25 +15,28 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .base import Restorer
|
||||
from .inpaint import InpaintRestorer
|
||||
|
||||
if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed)
|
||||
from ...config import AppConfig
|
||||
|
||||
|
||||
def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer:
|
||||
if name == "inpaint":
|
||||
return InpaintRestorer()
|
||||
def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer:
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
if name == "deepmosaics":
|
||||
from .deepmosaics import DeepMosaicsRestorer
|
||||
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
return DeepMosaicsRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "deepmosaics_video":
|
||||
from .deepmosaics import DeepMosaicsVideoRestorer
|
||||
|
||||
return DeepMosaicsVideoRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "lada":
|
||||
raise ValueError(
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README."
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics. См. README."
|
||||
)
|
||||
raise ValueError(f"Неизвестный режим восстановления: {name!r}")
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Classic inpainting restorer (cv2) — the always-available baseline.
|
||||
|
||||
HONEST LIMITATION: cv2 inpainting fills the masked region by propagating
|
||||
surrounding pixels. It removes the mosaic/bar but does NOT reconstruct the hidden
|
||||
detail — it smooths/guesses. For real reconstruction a generative model
|
||||
(DeepMosaics / LADA) is needed; this is the no-weights, no-GPU fallback so the
|
||||
"Расцензурить кадр" flow works end-to-end today.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Restorer
|
||||
from .mask import detections_to_mask
|
||||
|
||||
|
||||
class InpaintRestorer(Restorer):
|
||||
def __init__(self, radius: int = 3, dilate: int = 2, method: str = "telea") -> None:
|
||||
self.radius = radius
|
||||
self.dilate = dilate
|
||||
self.method = method
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"InpaintRestorer({self.method})"
|
||||
|
||||
def restore(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
detections: list[Detection],
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> np.ndarray:
|
||||
# Single cv2.inpaint call — effectively instant, so cancellation is moot.
|
||||
if not detections:
|
||||
return image.copy()
|
||||
mask = detections_to_mask(image.shape, detections, dilate=self.dilate)
|
||||
flags = cv2.INPAINT_TELEA if self.method == "telea" else cv2.INPAINT_NS
|
||||
return cv2.inpaint(image, mask, self.radius, flags)
|
||||
@@ -1,26 +0,0 @@
|
||||
"""Build a binary mask of the censored regions from detections."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
|
||||
|
||||
def detections_to_mask(
|
||||
shape: tuple[int, int], detections: list[Detection], dilate: int = 0
|
||||
) -> np.ndarray:
|
||||
"""White (255) over every detected region (polygon if present, else bbox)."""
|
||||
h, w = shape[:2]
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
for d in detections:
|
||||
if len(d.polygon) >= 3:
|
||||
cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255)
|
||||
else:
|
||||
x, y, bw, bh = d.bbox
|
||||
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1)
|
||||
if dilate > 0:
|
||||
k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8)
|
||||
mask = cv2.dilate(mask, k)
|
||||
return mask
|
||||
+191
-48
@@ -1,18 +1,66 @@
|
||||
"""Probe the PyTorch / CUDA situation, so the UI can show a device badge.
|
||||
"""Diagnose the PyTorch / CUDA situation so the UI can explain *why* it's on CPU.
|
||||
|
||||
Pure (no Qt). ``gather()`` imports torch (slow / heavy) — call it off the GUI
|
||||
thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and
|
||||
how to enable the GPU.
|
||||
Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to
|
||||
``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on
|
||||
the gathered dict and figures out the most likely cause + concrete fix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# pip index for the CUDA build (matches the README).
|
||||
CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121"
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# pip indexes for the CUDA builds (see README). cu121 needs a driver with CUDA >= 12.1.
|
||||
CUDA_WHEELS = {
|
||||
"cu121": "https://download.pytorch.org/whl/cu121",
|
||||
"cu118": "https://download.pytorch.org/whl/cu118",
|
||||
}
|
||||
|
||||
_NO_WINDOW = 0x08000000 if sys.platform == "win32" else 0 # CREATE_NO_WINDOW
|
||||
|
||||
|
||||
def _run_nvidia_smi() -> dict:
|
||||
"""Probe the NVIDIA driver/GPU via nvidia-smi. Never raises."""
|
||||
out: dict = {"found": False, "gpus": [], "driver_version": None, "cuda_driver": None}
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe and sys.platform == "win32":
|
||||
candidate = r"C:\Windows\System32\nvidia-smi.exe"
|
||||
exe = candidate if shutil.os.path.isfile(candidate) else None
|
||||
if not exe:
|
||||
return out
|
||||
# GPU names + driver version (robust CSV form).
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[exe, "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"],
|
||||
capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
||||
)
|
||||
if r.returncode == 0:
|
||||
out["found"] = True
|
||||
for line in r.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if parts and parts[0]:
|
||||
out["gpus"].append(parts[0])
|
||||
if len(parts) > 1 and parts[1]:
|
||||
out["driver_version"] = parts[1]
|
||||
except Exception: # noqa: BLE001 - any failure => "not found"
|
||||
return out
|
||||
# Max CUDA version the driver supports (only in the plain header).
|
||||
try:
|
||||
r2 = subprocess.run(
|
||||
[exe], capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
||||
)
|
||||
m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout)
|
||||
if m:
|
||||
out["cuda_driver"] = m.group(1)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
return out
|
||||
|
||||
|
||||
def gather() -> dict:
|
||||
"""Collect torch/CUDA facts. Never raises — missing torch is a valid result."""
|
||||
"""Collect torch + NVIDIA facts. Never raises — missing torch/GPU are valid."""
|
||||
info: dict = {
|
||||
"installed": False,
|
||||
"version": None, # torch.__version__ (e.g. "2.12.0+cpu")
|
||||
@@ -20,27 +68,38 @@ def gather() -> dict:
|
||||
"cuda_available": False,
|
||||
"device_name": None, # the active GPU's name, if any
|
||||
"import_error": None,
|
||||
# filled by nvidia-smi:
|
||||
"nvidia_smi": False,
|
||||
"gpus": [],
|
||||
"driver_version": None,
|
||||
"cuda_driver": None,
|
||||
}
|
||||
try:
|
||||
import torch
|
||||
except Exception as exc: # noqa: BLE001 - report any import failure, not just ImportError
|
||||
except Exception as exc: # noqa: BLE001 - report any import failure
|
||||
info["import_error"] = str(exc)
|
||||
return info
|
||||
info["installed"] = True
|
||||
info["version"] = getattr(torch, "__version__", None)
|
||||
try:
|
||||
info["built_cuda"] = torch.version.cuda
|
||||
except Exception: # noqa: BLE001
|
||||
info["built_cuda"] = None
|
||||
try:
|
||||
info["cuda_available"] = bool(torch.cuda.is_available())
|
||||
except Exception: # noqa: BLE001
|
||||
info["cuda_available"] = False
|
||||
if info["cuda_available"]:
|
||||
else:
|
||||
info["installed"] = True
|
||||
info["version"] = getattr(torch, "__version__", None)
|
||||
try:
|
||||
info["device_name"] = torch.cuda.get_device_name(0)
|
||||
info["built_cuda"] = torch.version.cuda
|
||||
except Exception: # noqa: BLE001
|
||||
info["device_name"] = None
|
||||
info["built_cuda"] = None
|
||||
try:
|
||||
info["cuda_available"] = bool(torch.cuda.is_available())
|
||||
except Exception: # noqa: BLE001
|
||||
info["cuda_available"] = False
|
||||
if info["cuda_available"]:
|
||||
try:
|
||||
info["device_name"] = torch.cuda.get_device_name(0)
|
||||
except Exception: # noqa: BLE001
|
||||
info["device_name"] = None
|
||||
|
||||
smi = _run_nvidia_smi()
|
||||
info["nvidia_smi"] = smi["found"]
|
||||
info["gpus"] = smi["gpus"]
|
||||
info["driver_version"] = smi["driver_version"]
|
||||
info["cuda_driver"] = smi["cuda_driver"]
|
||||
return info
|
||||
|
||||
|
||||
@@ -48,33 +107,117 @@ def device_label(info: dict) -> str:
|
||||
return "CUDA" if info.get("cuda_available") else "CPU"
|
||||
|
||||
|
||||
def reason(info: dict) -> str:
|
||||
"""One-sentence human explanation of the current device choice."""
|
||||
if not info.get("installed"):
|
||||
return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics "
|
||||
"работают на CPU (классический детектор torch не требует).")
|
||||
if info.get("cuda_available"):
|
||||
name = info.get("device_name") or "GPU"
|
||||
return f"PyTorch использует CUDA: {name}. Вычисления идут на видеокарте."
|
||||
version = info.get("version") or "?"
|
||||
built = info.get("built_cuda")
|
||||
if not built:
|
||||
return (f"Установлена CPU-сборка PyTorch ({version}) — без поддержки CUDA, "
|
||||
"поэтому вычисления идут на процессоре (медленно).")
|
||||
return (f"PyTorch собран с CUDA {built} ({version}), но GPU недоступен: нет "
|
||||
"NVIDIA-видеокарты, не установлен/устарел драйвер, либо версия CUDA "
|
||||
"несовместима с драйвером.")
|
||||
def _ver_tuple(v: str | None) -> tuple[int, ...]:
|
||||
try:
|
||||
return tuple(int(x) for x in str(v).split(".")[:2])
|
||||
except (ValueError, AttributeError):
|
||||
return ()
|
||||
|
||||
|
||||
def install_hint() -> str:
|
||||
"""Steps to enable the GPU (shown when running on CPU)."""
|
||||
def recommend_channel(info: dict) -> str:
|
||||
"""Pick the pip CUDA wheel index that matches the driver (cu118 for older)."""
|
||||
cd = _ver_tuple(info.get("cuda_driver"))
|
||||
if cd and cd < (12, 1):
|
||||
return "cu118"
|
||||
return "cu121"
|
||||
|
||||
|
||||
def install_command(channel: str = "cu121") -> str:
|
||||
"""The pip commands to (re)install the chosen CUDA build.
|
||||
|
||||
Targets the **running interpreter** (``sys.executable -m pip``) so the command
|
||||
hits the same venv that runs the app — not whatever ``pip`` is on PATH. (A common
|
||||
trap: running bare ``pip`` in a global shell while torch lives in the project venv.)
|
||||
"""
|
||||
index = CUDA_WHEELS.get(channel, CUDA_WHEELS["cu121"])
|
||||
py = sys.executable or "python"
|
||||
q = f'"{py}"' if " " in py else py
|
||||
return (
|
||||
"Как включить GPU (NVIDIA):\n"
|
||||
"1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n"
|
||||
"2. Переустановите PyTorch со сборкой CUDA:\n\n"
|
||||
" pip uninstall -y torch torchvision\n"
|
||||
f" pip install torch torchvision --index-url {CUDA_WHEEL_INDEX}\n\n"
|
||||
"3. Перезапустите приложение.\n\n"
|
||||
"Классический детектор работает и без CUDA. На CPU детекция и расцензуривание "
|
||||
"просто медленнее."
|
||||
f"{q} -m pip uninstall -y torch torchvision\n"
|
||||
f"{q} -m pip install torch torchvision --index-url {index}"
|
||||
)
|
||||
|
||||
|
||||
def analyze(info: dict) -> dict:
|
||||
"""Turn the raw facts into {summary, details[list], steps, command}."""
|
||||
installed = info.get("installed")
|
||||
cuda = info.get("cuda_available")
|
||||
built = info.get("built_cuda")
|
||||
gpus = info.get("gpus") or []
|
||||
smi = info.get("nvidia_smi")
|
||||
driver = info.get("driver_version")
|
||||
cuda_driver = info.get("cuda_driver")
|
||||
channel = recommend_channel(info)
|
||||
command = install_command(channel)
|
||||
|
||||
if not installed:
|
||||
build_str = "не установлен"
|
||||
elif built:
|
||||
build_str = f"CUDA {built}"
|
||||
else:
|
||||
build_str = "CPU-only (+cpu)"
|
||||
gpu_str = (
|
||||
", ".join(gpus) if gpus
|
||||
else ("не обнаружена" if smi or info.get("nvidia_smi") is False else "nvidia-smi не найден")
|
||||
)
|
||||
if not gpus and not smi:
|
||||
gpu_str = "nvidia-smi не найден (нет драйвера NVIDIA?)"
|
||||
details = [
|
||||
f"PyTorch: {info.get('version') or 'не установлен'}",
|
||||
f"Сборка PyTorch: {build_str}",
|
||||
f"CUDA доступна в PyTorch: {'да' if cuda else 'нет'}",
|
||||
f"Видеокарта (nvidia-smi): {gpu_str}",
|
||||
f"Драйвер NVIDIA: {driver or '—'}",
|
||||
f"Макс. CUDA драйвера: {cuda_driver or '—'}",
|
||||
f"Интерпретатор (venv): {sys.executable}",
|
||||
]
|
||||
if info.get("import_error"):
|
||||
details.append(f"Ошибка импорта torch: {info['import_error']}")
|
||||
|
||||
if not installed:
|
||||
summary = "PyTorch не установлен — детекция YOLO и DeepMosaics идут на CPU."
|
||||
steps = ("Установите PyTorch (CUDA-сборку, если есть NVIDIA-видеокарта):\n\n"
|
||||
+ command + "\n\nДля YOLO также: pip install -e \".[yolo]\"")
|
||||
elif cuda:
|
||||
gpu = info.get("device_name") or (gpus[0] if gpus else "GPU")
|
||||
summary = f"Всё в порядке: PyTorch использует CUDA. Активный GPU: {gpu}."
|
||||
steps = "GPU уже задействован — ничего делать не нужно."
|
||||
elif not built: # CPU-only torch build — the usual case
|
||||
if gpus:
|
||||
summary = (
|
||||
"Главная причина: установлена CPU-сборка PyTorch (+cpu) — она физически "
|
||||
f"не умеет в CUDA. Видеокарта ({gpus[0]}) и драйвер {driver or '?'} на месте, "
|
||||
"поэтому достаточно переустановить PyTorch со сборкой CUDA."
|
||||
)
|
||||
steps = ("Переустановите PyTorch под CUDA, затем перезапустите приложение:\n\n"
|
||||
+ command)
|
||||
else:
|
||||
summary = (
|
||||
"Установлена CPU-сборка PyTorch (+cpu), и видеокарта NVIDIA не обнаружена "
|
||||
"(nvidia-smi не отвечает). Либо нет NVIDIA GPU, либо не установлен драйвер."
|
||||
)
|
||||
steps = ("1. Проверьте видеокарту и драйвер: в консоли выполните nvidia-smi\n"
|
||||
"2. Если NVIDIA GPU есть — переустановите PyTorch под CUDA:\n\n"
|
||||
+ command +
|
||||
"\n\nБез NVIDIA GPU всё работает на CPU — просто медленнее.")
|
||||
else: # CUDA-enabled torch build, but CUDA still not available
|
||||
if not gpus:
|
||||
summary = (
|
||||
f"PyTorch собран под CUDA {built}, но видеокарта/драйвер NVIDIA не найдены. "
|
||||
"Скорее всего не установлен драйвер NVIDIA или нет GPU."
|
||||
)
|
||||
steps = "Установите свежий драйвер NVIDIA и перезапустите. Проверка: nvidia-smi."
|
||||
elif cuda_driver and _ver_tuple(cuda_driver) < _ver_tuple(built):
|
||||
summary = (
|
||||
f"Драйвер поддерживает CUDA {cuda_driver}, а PyTorch собран под CUDA {built} — "
|
||||
"версия драйвера слишком старая."
|
||||
)
|
||||
steps = ("Вариант A — обновите драйвер NVIDIA (рекомендуется).\n"
|
||||
"Вариант B — поставьте PyTorch под CUDA вашего драйвера:\n\n" + command)
|
||||
else:
|
||||
summary = (
|
||||
f"PyTorch собран под CUDA {built}, GPU ({gpus[0]}) есть, но CUDA недоступна — "
|
||||
"вероятен конфликт версий или повреждённая установка."
|
||||
)
|
||||
steps = "Переустановите PyTorch под CUDA:\n\n" + command
|
||||
return {"summary": summary, "details": details, "steps": steps, "command": command}
|
||||
|
||||
+156
-68
@@ -26,11 +26,10 @@ import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from PySide6.QtCore import Qt, QThreadPool
|
||||
from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
|
||||
from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcut
|
||||
from PySide6.QtWidgets import (
|
||||
QAbstractItemView,
|
||||
QApplication,
|
||||
QComboBox,
|
||||
QDialog,
|
||||
QDoubleSpinBox,
|
||||
QFileDialog,
|
||||
@@ -41,6 +40,7 @@ from PySide6.QtWidgets import (
|
||||
QListWidgetItem,
|
||||
QMainWindow,
|
||||
QMessageBox,
|
||||
QPlainTextEdit,
|
||||
QProgressBar,
|
||||
QPushButton,
|
||||
QSplitter,
|
||||
@@ -51,7 +51,7 @@ from PySide6.QtWidgets import (
|
||||
)
|
||||
|
||||
from .. import settings_store
|
||||
from ..config import AppConfig
|
||||
from ..config import AppConfig, normalize_config
|
||||
from ..core.detection import cache as detection_cache
|
||||
from ..core.detection.factory import build_detector
|
||||
from ..core.detection.types import Detection
|
||||
@@ -68,7 +68,6 @@ from .workers import Job
|
||||
|
||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||||
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
|
||||
_DETECTORS = ["classic", "yolo", "combined"]
|
||||
|
||||
|
||||
class MainWindow(QMainWindow):
|
||||
@@ -120,6 +119,8 @@ class MainWindow(QMainWindow):
|
||||
file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True))
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Движок восстановления…", self._open_restore_settings)
|
||||
file_menu.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
|
||||
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M")
|
||||
file_menu.addSeparator()
|
||||
@@ -136,14 +137,9 @@ class MainWindow(QMainWindow):
|
||||
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)
|
||||
|
||||
tb.addWidget(QLabel(" Детектор: YOLO "))
|
||||
self.model_action = QAction("Модель…", self, triggered=self._choose_model)
|
||||
self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики")
|
||||
tb.addAction(self.model_action)
|
||||
|
||||
tb.addSeparator()
|
||||
@@ -166,6 +162,12 @@ class MainWindow(QMainWindow):
|
||||
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
|
||||
restore.setToolTip("Восстановить найденные области на текущем кадре")
|
||||
tb.addAction(restore)
|
||||
restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False))
|
||||
restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)")
|
||||
tb.addAction(restore_all)
|
||||
restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True))
|
||||
restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)")
|
||||
tb.addAction(restore_regen)
|
||||
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
|
||||
self.toggle_restored_action.setEnabled(False)
|
||||
tb.addAction(self.toggle_restored_action)
|
||||
@@ -359,27 +361,58 @@ class MainWindow(QMainWindow):
|
||||
|
||||
info = self._device_info if self._device_info else torch_info.gather()
|
||||
cuda = bool(info.get("cuda_available"))
|
||||
a = torch_info.analyze(info)
|
||||
self._install_command = a["command"] # what the Copy button will copy
|
||||
|
||||
lines = [
|
||||
torch_info.reason(info),
|
||||
"ВЕРДИКТ:",
|
||||
a["summary"],
|
||||
"",
|
||||
"Диагностика:",
|
||||
f" • PyTorch: {info.get('version') or 'не установлен'}",
|
||||
f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}",
|
||||
f" • CUDA доступна: {'да' if cuda else 'нет'}",
|
||||
*(f" • {d}" for d in a["details"]),
|
||||
"",
|
||||
"Что делать:",
|
||||
a["steps"],
|
||||
]
|
||||
if info.get("device_name"):
|
||||
lines.append(f" • GPU: {info['device_name']}")
|
||||
if info.get("import_error"):
|
||||
lines.append(f" • Ошибка импорта torch: {info['import_error']}")
|
||||
|
||||
box = QMessageBox(self)
|
||||
box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning)
|
||||
box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU"))
|
||||
box.setText("\n".join(lines))
|
||||
# A real dialog (not QMessageBox) so the text — incl. the install command — is
|
||||
# selectable, and a Copy button drops the pip command straight onto the clipboard.
|
||||
dlg = QDialog(self)
|
||||
dlg.setWindowTitle("Почему " + ("GPU" if cuda else "CPU") + " — диагностика PyTorch/CUDA")
|
||||
dlg.resize(620, 480)
|
||||
layout = QVBoxLayout(dlg)
|
||||
|
||||
text = QPlainTextEdit()
|
||||
text.setReadOnly(True)
|
||||
text.setPlainText("\n".join(lines))
|
||||
mono = QFont("Consolas")
|
||||
mono.setStyleHint(QFont.Monospace)
|
||||
text.setFont(mono)
|
||||
layout.addWidget(text, 1)
|
||||
|
||||
buttons = QHBoxLayout()
|
||||
if not cuda:
|
||||
box.setInformativeText(torch_info.install_hint())
|
||||
box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands
|
||||
box.exec()
|
||||
copy_btn = QPushButton("Скопировать команду установки")
|
||||
copy_btn.clicked.connect(self._copy_install_command)
|
||||
buttons.addWidget(copy_btn)
|
||||
recheck = QPushButton("Проверить заново")
|
||||
recheck.setToolTip("Перепроверить torch/CUDA (например, после переустановки)")
|
||||
recheck.clicked.connect(lambda: (self._probe_device(), dlg.accept()))
|
||||
buttons.addWidget(recheck)
|
||||
buttons.addStretch(1)
|
||||
close_btn = QPushButton("Закрыть")
|
||||
close_btn.clicked.connect(dlg.accept)
|
||||
buttons.addWidget(close_btn)
|
||||
layout.addLayout(buttons)
|
||||
dlg.exec()
|
||||
|
||||
def _copy_install_command(self) -> None:
|
||||
command = getattr(self, "_install_command", None)
|
||||
if not command:
|
||||
from ..core import torch_info
|
||||
command = torch_info.install_command()
|
||||
QApplication.clipboard().setText(command)
|
||||
self.statusBar().showMessage("Команда установки скопирована в буфер обмена")
|
||||
|
||||
# ------------------------------------------------------------- cancellation
|
||||
def _begin_busy(self, total: int | None = None) -> None:
|
||||
@@ -388,7 +421,6 @@ class MainWindow(QMainWindow):
|
||||
self._cancel = False
|
||||
self.stop_action.setEnabled(True)
|
||||
# Disable inputs that would race a running job (they clear cache / rebuild engines).
|
||||
self.detector_combo.setEnabled(False)
|
||||
self.model_action.setEnabled(False)
|
||||
if total is None:
|
||||
self.progress.setRange(0, 0) # indeterminate
|
||||
@@ -400,7 +432,6 @@ class MainWindow(QMainWindow):
|
||||
def _end_busy(self) -> None:
|
||||
self._busy = False
|
||||
self.stop_action.setEnabled(False)
|
||||
self.detector_combo.setEnabled(True)
|
||||
self.model_action.setEnabled(True)
|
||||
self.progress.setVisible(False)
|
||||
self.progress.setRange(0, 100) # leave it determinate for the next user
|
||||
@@ -459,19 +490,19 @@ class MainWindow(QMainWindow):
|
||||
self._detector_key = key
|
||||
return self._detector
|
||||
|
||||
def _on_detector_changed(self, name: str) -> None:
|
||||
self._cfg.detector = name
|
||||
# YOLO/combined need a model. Auto-pick a known one from models/ if we have it;
|
||||
# only prompt when nothing suitable is found (don't nag when the path is obvious).
|
||||
if name in ("yolo", "combined") and not self._cfg.model_path:
|
||||
found = self._auto_find_model()
|
||||
if found:
|
||||
self._cfg.model_path = found
|
||||
self.statusBar().showMessage(f"Модель найдена автоматически: {found}")
|
||||
else:
|
||||
self._choose_model()
|
||||
self._persist_settings()
|
||||
self._invalidate_results()
|
||||
def _ensure_model(self) -> None:
|
||||
"""Make sure the YOLO detector has weights — auto-pick from ./models silently.
|
||||
|
||||
Called on project open. Doesn't prompt (the user can pick via "Модель…"); the
|
||||
detector factory raises a clear message if a detect is attempted without one.
|
||||
"""
|
||||
if self._cfg.model_path and Path(self._cfg.model_path).is_file():
|
||||
return
|
||||
found = self._auto_find_model()
|
||||
if found:
|
||||
self._cfg.model_path = found
|
||||
self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}")
|
||||
self._persist_settings()
|
||||
|
||||
@staticmethod
|
||||
def _auto_find_model() -> str | None:
|
||||
@@ -667,6 +698,8 @@ class MainWindow(QMainWindow):
|
||||
self._project = project
|
||||
project.frames_dir.mkdir(parents=True, exist_ok=True)
|
||||
project.apply_to_config(self._cfg) # per-project settings -> live config
|
||||
normalize_config(self._cfg) # coerce any legacy classic/inpaint values
|
||||
self._ensure_model() # YOLO needs weights — auto-pick if missing
|
||||
self._sync_settings_ui()
|
||||
self._detector_key = None
|
||||
self._restorer_key = None
|
||||
@@ -679,9 +712,6 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def _sync_settings_ui(self) -> None:
|
||||
"""Reflect the (project's) config onto the toolbar widgets without signal loops."""
|
||||
self.detector_combo.blockSignals(True)
|
||||
self.detector_combo.setCurrentText(self._cfg.detector)
|
||||
self.detector_combo.blockSignals(False)
|
||||
self.threshold_spin.blockSignals(True)
|
||||
self.threshold_spin.setValue(self._cfg.default_threshold)
|
||||
self.threshold_spin.blockSignals(False)
|
||||
@@ -912,41 +942,32 @@ class MainWindow(QMainWindow):
|
||||
key, dets = payload
|
||||
self._results[key] = dets
|
||||
self._tag_file(Path(key), len(dets))
|
||||
# If the frame being viewed was just computed, show its overlay live.
|
||||
if not self._showing_restored and self._current is not None and str(self._current) == key:
|
||||
self._show(self._current)
|
||||
self._tick_count += 1
|
||||
if self._tick_count % 25 == 0:
|
||||
self._refresh_marks() # let marks appear progressively (throttled)
|
||||
|
||||
# ------------------------------------------------------------- restoration
|
||||
def _restore_current(self) -> None:
|
||||
"""Restore the current frame's regions on a background thread, then show it.
|
||||
"""Restore the current frame on a background thread, then show it.
|
||||
|
||||
Detections are computed first (in the same job) if not cached. The DeepMosaics
|
||||
engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
|
||||
DeepMosaics locates the mosaic itself, so no detection step is needed — we just
|
||||
run the engine on the frame (if there's no mosaic the frame comes back unchanged).
|
||||
The engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
|
||||
if self._current is None or self._busy:
|
||||
return
|
||||
path = self._current
|
||||
key = str(path)
|
||||
|
||||
def fn(job):
|
||||
dets = self._results.get(key)
|
||||
if dets is None:
|
||||
dets = self._compute(self._make_detector(), path)
|
||||
job.tick(("dets", key, dets)) # cache them on the GUI thread
|
||||
if not dets:
|
||||
return ("empty", key)
|
||||
img = imread_unicode(key)
|
||||
if img is None:
|
||||
raise RuntimeError(f"Не удалось прочитать: {path.name}")
|
||||
restorer = self._make_restorer()
|
||||
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
|
||||
return ("restored", key, restored, len(dets), restorer.name)
|
||||
|
||||
def tick(payload):
|
||||
if payload[0] == "dets":
|
||||
_, k, dets = payload
|
||||
self._results[k] = dets
|
||||
self._tag_file(Path(k), len(dets))
|
||||
self._refresh_marks()
|
||||
restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled)
|
||||
return ("restored", key, restored, restorer.name)
|
||||
|
||||
def done(result, cancelled):
|
||||
if cancelled:
|
||||
@@ -954,19 +975,86 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
if result is None:
|
||||
return
|
||||
if result[0] == "empty":
|
||||
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
|
||||
return
|
||||
_, k, restored, n, engine = result
|
||||
_, k, restored, engine = result
|
||||
self._restored[k] = restored
|
||||
if self._current is not None and str(self._current) == k:
|
||||
self._showing_restored = True
|
||||
self.view.set_image(restored, [])
|
||||
self._update_restore_actions()
|
||||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name} — {n} обл.")
|
||||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}")
|
||||
|
||||
self.statusBar().showMessage(f"Восстановление: {path.name}…")
|
||||
self._start_job(fn, None, on_tick=tick, on_done=done)
|
||||
self._start_job(fn, None, on_done=done)
|
||||
|
||||
def _restore_all(self, force: bool = False) -> None:
|
||||
"""Restore every frame on a background thread, writing results to ``restored/``.
|
||||
|
||||
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
|
||||
engine skips frames already restored (resume) unless ``force``. The temporal
|
||||
engine (DeepMosaics-video) runs the whole contiguous sequence in order via
|
||||
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
|
||||
"""
|
||||
if not self._files or self._project is None or self._busy:
|
||||
return
|
||||
files = list(self._files) # snapshot — favorites/move mutate self._files
|
||||
total = len(files)
|
||||
out_dir = self._project.restored_dir
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def out_path(p: Path) -> Path:
|
||||
return out_dir / f"{p.stem}.jpg"
|
||||
|
||||
def fn(job):
|
||||
restorer = self._make_restorer() # built on the worker (may raise)
|
||||
frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads
|
||||
|
||||
def get_frame(i):
|
||||
img = frame_cache.get(i)
|
||||
if img is None:
|
||||
img = imread_unicode(str(files[i]))
|
||||
if img is None:
|
||||
raise RuntimeError(f"Не удалось прочитать: {files[i].name}")
|
||||
if len(frame_cache) > 24:
|
||||
frame_cache.clear()
|
||||
frame_cache[i] = img
|
||||
return img
|
||||
|
||||
def emit(i, restored):
|
||||
imwrite_unicode(str(out_path(files[i])), restored)
|
||||
job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}")
|
||||
|
||||
if restorer.temporal:
|
||||
restorer.restore_sequence(
|
||||
total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled
|
||||
)
|
||||
else:
|
||||
for i, p in enumerate(files):
|
||||
if job.cancelled:
|
||||
break
|
||||
if not force and out_path(p).is_file():
|
||||
job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}")
|
||||
continue
|
||||
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
||||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||||
return None
|
||||
|
||||
def done(_result, cancelled):
|
||||
if self._current is not None: # live-preview the current frame's result, if any
|
||||
rp = out_path(self._current)
|
||||
if rp.is_file():
|
||||
img = imread_unicode(str(rp))
|
||||
if img is not None:
|
||||
self._restored[str(self._current)] = img
|
||||
self._showing_restored = True
|
||||
self.view.set_image(img, [])
|
||||
self._update_restore_actions()
|
||||
self.statusBar().showMessage(
|
||||
"Расцензуривание отменено" if cancelled
|
||||
else f"Готово: результаты в {out_dir.name}/ ({total} кадров)"
|
||||
)
|
||||
|
||||
self.statusBar().showMessage("Пакетное расцензуривание…")
|
||||
self._start_job(fn, total, on_done=done)
|
||||
|
||||
def _make_restorer(self):
|
||||
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
"""Configure the restoration ("расцензурить") engine.
|
||||
|
||||
inpaint — no setup. deepmosaics — the network code is vendored (built-in); the user
|
||||
picks a clean model from a dropdown of the bundled ``models/deepmosaics`` weights
|
||||
(or browses to another ``clean_*.pth``). ``mosaic_position.pth`` must sit beside the
|
||||
chosen model. A CUDA GPU is strongly recommended (GPU id, -1 = CPU/slow).
|
||||
Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео")
|
||||
engine. The network code is vendored (built-in); the user picks a clean model from a
|
||||
dropdown of the bundled ``models/deepmosaics`` weights (or browses to another
|
||||
``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered.
|
||||
``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly
|
||||
recommended (GPU id, -1 = CPU/slow).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -35,9 +37,9 @@ class RestoreDialog(QDialog):
|
||||
self.setMinimumWidth(560)
|
||||
|
||||
self.engine = QComboBox()
|
||||
self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint")
|
||||
self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics")
|
||||
self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0)
|
||||
self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics")
|
||||
self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video")
|
||||
self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer)))
|
||||
self.engine.currentIndexChanged.connect(self._sync)
|
||||
|
||||
# Model dropdown — bundled clean models, plus the configured one if external.
|
||||
@@ -52,9 +54,11 @@ class RestoreDialog(QDialog):
|
||||
form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model))
|
||||
form.addRow("GPU id:", self.dm_gpu)
|
||||
hint = QLabel(
|
||||
"Модели берутся из models/deepmosaics. Нужны clean_youknow_resnet_9blocks.pth "
|
||||
"и mosaic_position.pth (рядом). Видеомодель clean_*_video.pth покадрово не "
|
||||
"работает и в списке не показывается. На CPU медленно — лучше GPU. См. README."
|
||||
"Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n"
|
||||
"• Картинка: clean_youknow_resnet_9blocks.pth — покадрово.\n"
|
||||
"• Видео: clean_youknow_video.pth — использует соседние кадры (когерентнее на "
|
||||
"роликах), требует прогона по диапазону («Расцензурить все»).\n"
|
||||
"На CPU медленно — лучше GPU. См. README."
|
||||
)
|
||||
hint.setWordWrap(True)
|
||||
form.addRow(hint)
|
||||
@@ -67,7 +71,12 @@ class RestoreDialog(QDialog):
|
||||
|
||||
def _populate_models(self, current: str | None) -> None:
|
||||
self.model_combo.clear()
|
||||
for name, path in discover_models():
|
||||
is_video = self.engine.currentData() == "deepmosaics_video"
|
||||
if is_video: # temporal engine: only the video weights (clean_*_video.pth)
|
||||
models = [(n, p) for n, p in discover_models(include_video=True) if "video" in n.lower()]
|
||||
else:
|
||||
models = discover_models() # per-frame engine: image clean models only
|
||||
for name, path in models:
|
||||
self.model_combo.addItem(name, path)
|
||||
# Keep an externally-configured model selectable even if it's outside the folder.
|
||||
if current and self.model_combo.findData(current) < 0:
|
||||
@@ -88,7 +97,9 @@ class RestoreDialog(QDialog):
|
||||
return w
|
||||
|
||||
def _sync(self) -> None:
|
||||
is_dm = self.engine.currentData() == "deepmosaics"
|
||||
is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video")
|
||||
# The model list differs per engine (image vs video weights) — repopulate.
|
||||
self._populate_models(self._cfg.dm_model)
|
||||
self.model_combo.setEnabled(is_dm)
|
||||
self.dm_gpu.setEnabled(is_dm)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user