diff --git a/CLAUDE.md b/CLAUDE.md index 97ac7e3..9ab9ba8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -27,30 +27,41 @@ EchoVault, then this file — and update whichever is stale. **HVideoTool** is a Windows-first desktop GUI utility that **detects already-applied censorship** (mosaic, pixelation, blur, black bars) in **images**, and draws outlines -over the detected censored regions. You open a folder of images; it runs each through -a detector, draws the regions, and shows a detailed per-image list of what it found. +over the detected censored regions. You open a **project** (see below); it runs each +image through a detector, draws the regions, and shows a detailed per-image list of +what it found. -> **Scope was deliberately narrowed (this session).** It used to extract frames from -> video, detect, and play back with overlays (a "project model" with worker threads). -> That whole video pipeline was **removed** — the tool is now a simple **image-folder -> inspector** for viewing/debugging detector output on test images. Pre-extract video -> to frames externally if you need that. +> **Projects (this session).** Work is organized into **projects** — a project is a +> folder holding `project.json` (metadata + per-project settings) · `frames/` (the +> images) · `detections.json` (the detection cache) · `collections/Избранное` (the +> single default "favorites" collection). See `core/project.py`. The per-project settings (detector, model, +> threshold, restore engine) live in `project.json`; the global `settings.json` only +> seeds the **defaults for new projects**. The old "open a bare folder" flow is now +> "Импортировать папку как проект…" (copies images into a new project's `frames/`). + +> **Scope is still narrow.** It used to extract frames from video, detect, and play +> back with overlays (a thread-based "project model"). That video *playback* pipeline +> was **removed**; the new "projects" are just an on-disk layout, NOT worker threads or +> playback. The tool remains a synchronous, single-threaded **image inspector**. 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. The shipped engine is a cv2 **inpaint baseline** (fills, does NOT truly - reconstruct); a generative engine (DeepMosaics / LADA BasicVSR++) is the intended - real engine but needs weights + a CUDA GPU and is not wired yet. + 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. - Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was rejected early — a generative *conditioning* model, not a censorship restorer. Don't reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/ LADA), not a general text-to-image diffusion pipeline. - It detects **already-censored** regions, not "content that should be censored" (i.e. not an NSFW classifier). -- Video is only a one-shot frame-extraction convenience (see below); detection and - restoration operate on image folders. +- Video is only a one-shot frame-extraction convenience (see below): "Создать из + ролика…" makes a new project and decodes the clip into its `frames/`. Detection and + restoration operate on the project's images. ## Target environment @@ -75,37 +86,42 @@ torch present. ## Architecture (as implemented) > This reflects the actual code on disk. It is a synchronous, single-threaded GUI app -> — no worker threads, no project/cache files. The only video touch is a one-shot -> "Создать из ролика…" that decodes a clip to a folder of JPGs via the **ffmpeg CLI** -> (cv2.VideoCapture fallback; NOT PyAV); detection still works on image folders only. -> Detection runs on the GUI -> thread (lazily per image, or via "Детектировать все"). When code and this file -> disagree, trust the code. +> — no worker threads. Work is organized into **projects** (`core/project.py`): a +> project folder holds `project.json` (metadata + per-project settings), `frames/` (the +> images), `detections.json` (the detection cache, at the project root — no longer a +> sidecar next to the images), and `collections/Избранное` (the favorites collection). The only +> video touch is a one-shot "Создать из ролика…" that creates a new project and decodes +> a clip into its `frames/` via the **ffmpeg CLI** (cv2.VideoCapture fallback; NOT +> PyAV). Detection runs on the GUI thread (lazily per image, or via "Детектировать +> все"). When code and this file disagree, trust the code. ``` hvideotool/ -├── __main__.py # entry point + CLI (optional folder arg, --detector, --model) -├── app.py # QApplication bootstrap; run(config, folder=None) +├── __main__.py # entry point + CLI (optional project path, --detector, --model) +├── app.py # QApplication bootstrap; run(config, target=None) — opens/auto-reopens a project ├── config.py # AppConfig + DetectionConfig/OverlayConfig (thresholds live here) -├── settings_store.py # persist detector/model/threshold/last_dir to ~/HVideoTool/settings.json +├── settings_store.py # new-project DEFAULTS + last/recent projects to ~/HVideoTool/settings.json ├── ui/ │ ├── main_window.py # the whole UI: toolbar + [file list | image view | detail table] │ └── image_view.py # renders an image + draws polygon/bbox overlays (QPainter); can highlight one └── core/ ├── imageio.py # unicode-safe imread/imwrite (np.fromfile + imdecode) + ├── project.py # Project: layout (project.json/frames/detections.json/collections) + per-project settings ├── video/ │ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type ├── restore/ # "un-censor" detected regions (per-frame) - │ ├── base.py # Restorer ABC: restore(image, detections) -> image + │ ├── 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 — shells out to user-installed DeepMosaics + │ ├── 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/ ├── base.py # Detector ABC: detect(frame) -> list[Detection] ├── factory.py # build_detector(config) -> classic | yolo | combined ├── types.py # Detection (+ to_dict/from_dict), CensorType enum + ├── 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 @@ -116,8 +132,23 @@ hvideotool/ - `MainWindow` holds the config, builds the detector lazily via `build_detector` (cached by detector+model+conf in `_make_detector`), and keeps `_results: dict[path -> list[Detection]]` as the detection cache. -- "Открыть папку" lists image files (`_IMAGE_EXTS`) with a progress bar (bulk insert - with `setUpdatesEnabled(False)` + periodic `processEvents`), so a big folder doesn't +- **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: "Создать проект…" + (`_create_project`), "Открыть проект…" (`_open_project_dialog`), "Импортировать папку + как проект…" (`_import_folder_as_project` — copies images into a new project's + `frames/`, carries over an old `.hvideotool_detections.json` sidecar if present), and + a "Недавние проекты" submenu. `_open_project(project)` is the core open: it + `apply_to_config`s the project's settings, syncs the toolbar widgets without signal + loops (`_sync_settings_ui`), titles the window, records last/recent, and lists + `frames/`. On startup `app.run` opens the CLI `target` or auto-reopens + `settings_store.last_project()` (`_auto_open_last`). +- **Per-project settings.** Detector/model/threshold/restore engine live in + `project.json` (`Project.settings`, `_SETTING_KEYS`). `_persist_settings()` writes both + the global defaults (for new projects) **and** the open project. Global `settings.json` + is now only defaults + last/recent projects. +- Listing image files (`_IMAGE_EXTS`) from `frames/` uses a progress bar (bulk insert + with `setUpdatesEnabled(False)` + periodic `processEvents`), so a big project doesn't freeze silently. - **Viewing and detecting are decoupled on purpose** (so browsing stays instant even with a slow CPU detector): selecting a file only *shows* it with its cached result @@ -126,38 +157,68 @@ hvideotool/ "Детектировать все" (whole folder, progress bar). Do NOT re-add auto-detect-on-select. Results cache in `_results`; the file-list row gets a count suffix when computed. Switching detector/model clears the cache (`_invalidate_results`). -- **Collections (curation).** A combo in the left pane (next to the file list, since it - acts on the list selection — not on the toolbar, to keep that uncluttered) - (`collection_combo`) picks the active destination; `_refresh_collections()` repopulates it from sibling folders of the - opened folder (`_collections_base()` = the opened folder's parent, else - `~/HVideoTool/collections`) on load/create, so previously-made collections are - reselectable. Items carry the path in itemData; "— не выбрана —" = None, - "Выбрать папку…" = `"__browse__"` sentinel → `_browse_collection()` for an arbitrary - folder (kept in the combo even if outside base). "Создать…" makes a new one and - selects it. The file list is `ExtendedSelection`; "В коллекцию" / Ctrl+M **moves** - (`shutil.move`, not copy) the selected frames there, removing them from - list/`_files`/cache. `_unique_dest` avoids clobbering (`foo.jpg` → `foo (1).jpg`). - Use case: sort frames into a training/example set while inspecting detections. +- **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` + take the cache file and the image `base_dir` (= `project.frames_dir`) separately, since + the cache lives at the project root, not next to the images. It's tagged with + the detector identity (`_results_key` = detector + model + conf/imgsz); on open, + `_load_cached_results` reloads it **only on a key match** (else ignored, not shown as + current). Saved (`_save_results`, skipped when `_results` is empty so it never clobbers + a good cache with nothing) after detect-all (incl. cancel → partial), single + detect/recompute, move-to-collection, and on `closeEvent`. + **"Детектировать все" is incremental** (skips already-cached frames → resumes/top-ups); + **"Все заново"** (`_detect_all(force=True)`) clears the cache first (full regen); + **"Рассчитать кадр"/Space** always recomputes the one current frame. Empty list in the + cache = "checked, clean" (tinted green, no mark); `in _results` distinguishes it from + "not computed". +- **Favorites (curation).** Curation was simplified (user request) to a **single + default collection** — no create/select/browse UI. The "★ В избранное" button (left + pane, under the file list) / "В избранное" menu item / Ctrl+M → `_move_to_favorites` + **moves** (`shutil.move`, not copy) the `ExtendedSelection`-selected frames into + `project.favorites_dir` (`collections/Избранное`, `FAVORITES_DIR` in `core/project.py`, + 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), caches the result in `_restored[path]`, and shows it overlay-free. "Показать оригинал/ результат" toggles (`_showing_restored`); "Сохранить результат" writes - `_restored.jpg` into the active collection (or beside the frame). The baseline - is cv2 inpaint; the real engine is **DeepMosaics** (`restore/deepmosaics.py`), which - shells out to a user-installed DeepMosaics repo (`deepmosaic.py --mode clean`, stdin - closed so its error `input()` can't hang; reads the newest image from a temp - result_dir). The engine + paths are set in `RestoreDialog` (Файл → Движок - восстановления…), persisted, and built lazily/cached in `_make_restorer` (like - `_make_detector`). NOTE: DeepMosaics locates mosaics itself (its `mosaic_position.pth`) - — our detections aren't passed to it. `_show` resets `_showing_restored` + - `_update_restore_actions`. To add another engine (e.g. LADA), implement + `_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` is polled at entry (raises `Cancelled`). + 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`. - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys `,`/`.`), a scrubber `frame_slider` across the whole sequence, a `pos_label` ("row / n"), and jump-to-detection (◀ детекция / детекция ▶, keys `[`/`]`, `_step_hit` scans `_results` for the next non-empty frame). The slider and file list are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all - navigation ultimately drives `file_list.setCurrentRow`. + navigation ultimately drives `file_list.setCurrentRow`. The scrubber is a custom + `MarkerSlider` (`ui/marker_slider.py`) that paints cyan ticks at frames with + detections (`_refresh_marks` projects `_results` onto row indices; per-pixel deduped + so big folders stay cheap). File-list rows are tinted too (`_tag_file`): red = + censorship found, green = checked & clean. Both reset on `_invalidate_results`. +- **Cancellation (cooperative, no threads).** A single "■ Стоп" toolbar action (Esc) + cancels the running long op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + + the Stop button + the progress bar (`total=None` → indeterminate); `_request_cancel` + sets `self._cancel`; the loop-based ops (`_detect_all`, `_load_folder`) and video + extraction (its `progress` cb returns `not self._cancel`) check the flag between + `processEvents` ticks. Single-image restore passes `should_cancel=self._poll_cancel` + (which pumps `processEvents` then returns the flag) into `Restorer.restore`; only + DeepMosaics actually polls it (kills its subprocess + raises `Cancelled`) — cv2 ops are + instant. Entry points guard with `if self._busy: return` (notably `_move_to_collection`, + which mutates `_files` that `_detect_all` iterates). This keeps the synchronous, + single-threaded model — do NOT reintroduce worker threads for cancellation. - `image_view.ImageView` draws the image scaled-to-fit plus overlays. Overlay visibility/threshold are applied at paint time. Selecting a row in the detail table calls `set_highlight(i)` — that detection is drawn boldly (even below threshold) and @@ -179,14 +240,15 @@ hvideotool/ python -m venv .venv; .\.venv\Scripts\Activate.ps1 pip install -e . # classic detector needs no torch/CUDA -python -m hvideotool # open a folder in-app -python -m hvideotool "C:\path\to\images" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +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 ``` No formal test suite. Headless sanity check: set `QT_QPA_PLATFORM=offscreen`, build a -`MainWindow`, `open_path(folder)`, drive `file_list.setCurrentRow(...)`, and read +`MainWindow`, `Project.create(tmp)` + copy a few images into `frames/`, +`_open_project(project)`, drive `file_list.setCurrentRow(...)`, and read `detail_table` / `detail_header`. Or run `build_detector(config).detect(...)` on a frame directly. @@ -229,9 +291,10 @@ frame directly. - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths. - Don't reintroduce any generative / ControlNet dependency, nor the removed video - *pipeline* (PyAV, project/cache, worker threads, playback). The one allowed video - touch is `core/video/extract.py` (one-shot decode → JPG folder, behind "Создать из - ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by + *playback pipeline* (PyAV, worker threads, player). (The new `core/project.py` is an + on-disk layout, not that thread-based "project model".) The one allowed video touch is + `core/video/extract.py` (one-shot decode → a new project's `frames/`, behind "Создать + из ролика…"): ffmpeg CLI — `_find_ffmpeg()` prefers PATH, else the binary bundled by the `imageio-ffmpeg` dep, else cv2 fallback. Keyframe-only `-skip_frame nokey` is ~10× faster than every-frame; `-hwaccel` does NOT help (GPU transfer overhead). Use ffmpeg/cv2, not PyAV, and keep it synchronous. Decoding every frame is the inherent diff --git a/README.md b/README.md index 3c72887..db36b07 100644 --- a/README.md +++ b/README.md @@ -2,10 +2,29 @@ Десктопная утилита с графическим интерфейсом для **обнаружения уже наложенной цензуры** (мозаика, пикселизация, размытие, чёрные плашки) на **картинках**. -Открываете папку с изображениями (или раскадровываете ролик кнопкой «Создать из -ролика…») — приложение прогоняет каждый кадр через детектор, **обводит найденные -области** и показывает **подробный список** того, что нашлось на каждой картинке. -Отобранные кадры можно перемещать в **коллекции** (для сбора датасета/примеров). +Работа организована в **проекты**: создаёте проект (или раскадровываете ролик кнопкой +«Создать из ролика…») — приложение прогоняет каждый кадр через детектор, **обводит +найденные области** и показывает **подробный список** того, что нашлось на каждой +картинке. Отобранные кадры можно перемещать в **избранное** (для сбора датасета/ +примеров). + +### Проект + +Проект — это папка со всем необходимым: + +``` +МойПроект/ +├── project.json # настройки проекта (детектор, модель, порог, движок) + метаданные +├── frames/ # картинки проекта +├── detections.json # кэш детекций +└── collections/ # коллекции; «Избранное» — папка для отобранных кадров +``` + +Настройки (детектор / модель / порог / движок восстановления) хранятся **внутри +проекта** — каждый проект помнит, как его настраивали. Глобальные настройки +(`~/HVideoTool/settings.json`) задают лишь **значения по умолчанию для новых проектов** +и список недавних. При запуске без аргументов автоматически открывается последний +проект. > Инструмент для просмотра, отладки детекции и отбора кадров. Видео раскадровывает > через ffmpeg (ставится автоматически с пакетом `imageio-ffmpeg`; системный ffmpeg @@ -15,62 +34,97 @@ ## Возможности -- **Создать из ролика…** — раскадровка видео в папку-коллекцию. Два режима: - **только ключевые кадры** (в разы быстрее — декодируются лишь I-кадры) и +- **Проекты**: «Создать проект…», «Открыть проект…», «Импортировать папку как + проект…» (копирует картинки из обычной папки в `frames/` нового проекта) и подменю + **«Недавние проекты»**. +- **Создать из ролика…** — раскадровка видео в новый проект (кадры в `frames/`). Два + режима: **только ключевые кадры** (в разы быстрее — декодируются лишь I-кадры) и **каждый N-й кадр**; опциональный даунскейл (меньше файлов и нагрузки на диск/АВ). ffmpeg идёт в комплекте (`imageio-ffmpeg`); системный ffmpeg из PATH — в приоритете. -- **Коллекции**: «Создать коллекцию…» + «В коллекцию» (Ctrl+M) перемещает выбранные - кадры в активную папку-коллекцию (мультивыбор поддерживается). -- **Открыть папку** с картинками (`.jpg/.png/.bmp/.webp/.tif`) — список слева. +- **Избранное**: «★ В избранное» (Ctrl+M) перемещает выбранные кадры в папку + `collections/Избранное` внутри проекта — для отбора кадров (мультивыбор + поддерживается). Папка создаётся автоматически. +- Картинки проекта (`.jpg/.png/.bmp/.webp/.tif`) — список слева. - Картинка с **обводкой контуром** найденных областей — по центру. - **Подробная таблица детекций** справа: тип, уверенность, bbox, число точек полигона. Выбор строки **подсвечивает** конкретную область на картинке. -- **Ленивая детекция**: картинка прогоняется при первом открытии, результат - кэшируется. Кнопка **«Детектировать все»** обходит всю папку. +- **Ленивая детекция**: картинка прогоняется по двойному клику / кнопке «Рассчитать + кадр», результат кэшируется. +- **Кэш детекций сохраняется в проект** (см. [ниже](#кэш-детекций)) — при повторном + открытии проекта результаты подхватываются, не нужно считать заново. +- **Три режима пересчёта**: + - **«Детектировать все»** — *дозапуск*: считает только ещё не посчитанные кадры + (можно прерывать и продолжать); + - **«Все заново»** — полная регенерация: очищает кэш и пересчитывает весь проект; + - **«Рассчитать кадр»** (Space / двойной клик) — всегда пересчитывает текущий кадр. +- **Кнопка «■ Стоп» (Esc)** отменяет любую текущую длинную операцию (детекция всего + проекта, раскадровка ролика, импорт папки, восстановление DeepMosaics). Уже + посчитанное при отмене сохраняется в кэш. +- **Навигация под картинкой**: ◀ ▶ (`,`/`.`), ползунок-перемотка и переходы к + кадрам с детекцией ◀/▶ (`[`/`]`). На ползунке **бирюзовыми метками** отмечены кадры + с найденной цензурой; строки списка **подсвечиваются цветом** (🔴 цензура найдена, + 🟢 проверено и чисто). - Переключение **детектора** (`classic` / `yolo` / `combined`) и **порога** уверенности прямо в тулбаре — удобно сравнивать. - Выбор файла весов модели кнопкой **«Модель…»**. +## Кэш детекций + +Результаты детекции сохраняются в корне проекта в файл `detections.json`. Это даёт: + +- **возобновление между сессиями** — открыли проект повторно, готовые детекции сразу + на месте (метки на ползунке и подсветка строк восстанавливаются); +- **дозапуск** — «Детектировать все» пропускает уже посчитанные кадры; +- **отказоустойчивость** — при отмене/закрытии посчитанное не теряется. + +Кэш помечен «удостоверением» детектора (детектор + модель + порог `conf`/`imgsz`). +Если открыть проект **другим** детектором, чужой кэш не загружается (чтобы не выдавать +старые результаты за текущие). Полностью пересчитать — кнопка **«Все заново»**. + +> Ключи внутри файла — **имена файлов**, поэтому кэш переживает перемещение/ +> переименование проекта. Хранится результат **одного** детектора за раз: посчитали +> одним, переключились на другой и посчитали — кэш перезапишется. + ## Восстановление (расцензуривание) Кнопка **«Расцензурить кадр»** восстанавливает найденные области на текущем кадре, **«Показать оригинал/результат»** переключает вид, **«Сохранить результат»** пишет -`<имя>_restored.jpg` (в активную коллекцию или рядом с кадром). Движок выбирается в +`<имя>_restored.jpg` рядом с кадром. Движок выбирается в меню **Файл → Движок восстановления…**. Два движка: - **Инпейнт (cv2)** — по умолчанию, без модели и GPU. ⚠️ *Заполняет* область по окружению, но **не реконструирует** скрытые детали (замазывает, а не раскрывает). -- **DeepMosaics** — реальное генеративное удаление мозаики. Требует **GPU NVIDIA/CUDA** - (на CPU очень медленно) и отдельной установки. На аниме качество ограничено - (модели обучены на реальном видео). +- **DeepMosaics** — реальное генеративное удаление мозаики. Код **встроен** в + приложение (vendored, GPL-3.0), ставить его отдельно не нужно — требуются только + **веса** и (желательно) **GPU NVIDIA/CUDA**. На аниме качество ограничено (модели + обучены на реальном видео). ### Настройка DeepMosaics -```powershell -git clone https://github.com/HypoX64/DeepMosaics -# установите зависимости DeepMosaics (см. его README; нужен torch с CUDA) -# скачайте веса (clean_youknow_resnet_9blocks.pth + mosaic_position.pth) в pretrained_models/mosaic -``` +Нужны только **веса** — положите их в **`models/deepmosaics/`** (папка в `.gitignore`, +веса большие, ~92 МБ). Скачать: официальная папка +([Google Drive](https://drive.google.com/drive/folders/1LTERcN33McoiztYEwBxMuRjjgxh4DEPs), +Baidu код `1x0a`): -Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, укажите папку -DeepMosaics (с `deepmosaic.py`), файл весов и GPU id (`-1` = CPU). Приложение вызывает -DeepMosaics на текущем кадре и показывает результат. +- **`clean_youknow_resnet_9blocks.pth`** — картиночная clean-модель; +- **`mosaic_position.pth`** — локатор мозаики (должен лежать рядом). -> **Важно:** для покадрового режима берите **картиночную** модель -> `clean_youknow_resnet_9blocks.pth`. Видеомодель `clean_youknow_video.pth` (BVDNet) -> покадрово **не работает** — ей нужен соседний кадр. Если кадр без мозаики, движок -> вернёт его без изменений. +Затем в приложении: **Файл → Движок восстановления… → DeepMosaics**, выберите **модель +из выпадающего списка** (наполняется из `models/deepmosaics`; есть «Обзор…» для файла в +другом месте) и GPU id (`-1` = CPU). Если веса в `models/deepmosaics` — работает сразу; +модель грузится один раз, дальше кадры считаются быстро. + +> **Важно:** берите именно **картиночную** модель `clean_youknow_resnet_9blocks.pth`. +> Видеомодель `clean_youknow_video.pth` (BVDNet) покадрово **не работает** — ей нужен +> соседний кадр (приложение это распознаёт и подскажет). Если на кадре нет мозаики, +> результат = исходный кадр. > -> DeepMosaics 2021 года рассчитан на старые версии (`torch 1.7`, `numpy 1.19`). -> Проверено: на современных `torch 2.x`/`numpy 2.x` картиночная модель запускается -> (CPU ~7 с/кадр), но если столкнётесь с несовместимостью — заведите для DeepMosaics -> отдельное окружение по его `requirements.txt` и укажите его `python.exe` в диалоге. - -> Движок подключается через интерфейс `core/restore/base.Restorer` (`build_restorer`). -> [LADA](https://github.com/ladaapp/lada) (BasicVSR++, лучшее качество на реальном -> видео, но видеомодель) пока не подключён. +> Код DeepMosaics (GPL-3.0) лежит в `core/restore/_deepmosaics/` и поэтому **весь +> проект распространяется под GPL-3.0**. Запускается на современных `torch 2.x`/ +> `numpy 2.x` (проверено). [LADA](https://github.com/ladaapp/lada) (видеомодель, +> лучшее качество на реальном видео) пока не подключён. ## Что НЕ делает (осознанно вне области задачи) @@ -105,11 +159,11 @@ pip install -e . ## Запуск ```powershell -# Без аргументов — папку открываете в приложении (тулбар → «Открыть папку…») +# Без аргументов — открывается последний проект (или создайте/откройте новый в тулбаре) python -m hvideotool -# Необязательно: сразу открыть папку / переопределить детектор и модель -python -m hvideotool "C:\path\to\images" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt +# Необязательно: сразу открыть проект / переопределить детектор и модель (по умолчанию) +python -m hvideotool "C:\path\to\МойПроект" --detector yolo --model models\lada_mosaic_detection_model_v4_accurate.pt ``` Выбор детектора, путь к модели и порог сохраняются в `~/HVideoTool/settings.json` @@ -176,24 +230,28 @@ hvideotool/ ├── __main__.py # точка входа + CLI (всё опционально) ├── app.py # инициализация QApplication ├── config.py # настройки: пороги детекции, оверлей, детектор/модель -├── settings_store.py # детектор/модель/порог/последняя папка → settings.json +├── settings_store.py # дефолты новых проектов + последний/недавние → settings.json ├── ui/ │ ├── main_window.py # окно: список файлов | картинка | таблица детекций -│ └── image_view.py # отрисовка картинки + оверлей-контуры (QPainter) +│ ├── image_view.py # отрисовка картинки + оверлей-контуры (QPainter) +│ └── marker_slider.py # ползунок-перемотка с метками кадров с детекцией └── core/ ├── 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: объединение детекторов ``` Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU -заметно медленнее, чем на CUDA. +заметно медленнее, чем на CUDA. Длинные операции можно прервать кнопкой «■ Стоп» +(Esc), а результаты кэшируются на диск — см. [Кэш детекций](#кэш-детекций). ## Технологический стек diff --git a/hvideotool/__main__.py b/hvideotool/__main__.py index efc2160..ad8b735 100644 --- a/hvideotool/__main__.py +++ b/hvideotool/__main__.py @@ -1,7 +1,8 @@ """Command-line entry point: ``python -m hvideotool``. -Runs with no arguments — open a folder of images in-app. CLI flags are optional -overrides; choices persist to ~/HVideoTool/settings.json. +Runs with no arguments — reopens the last project (if any). An optional path opens +that project. CLI flags are optional overrides for the new-project defaults; they +persist to ~/HVideoTool/settings.json. """ from __future__ import annotations @@ -19,20 +20,20 @@ def main() -> int: prog="hvideotool", description="Инспектор детекции уже наложенной цензуры на картинках.", ) - parser.add_argument("folder", nargs="?", help="папка с картинками для немедленного открытия") + 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 choices first + settings_store.apply(config) # persisted defaults first if args.detector: config.detector = args.detector if args.model_path: config.model_path = args.model_path - return run(config, folder=args.folder) + return run(config, target=args.target) if __name__ == "__main__": diff --git a/hvideotool/app.py b/hvideotool/app.py index 06fcf2f..36ff229 100644 --- a/hvideotool/app.py +++ b/hvideotool/app.py @@ -10,11 +10,13 @@ from .config import AppConfig from .ui.main_window import MainWindow -def run(config: AppConfig, folder: str | None = None) -> int: +def run(config: AppConfig, target: str | None = None) -> int: app = QApplication(sys.argv) app.setApplicationName("HVideoTool") window = MainWindow(config) window.show() - if folder: - window.open_path(folder) + if target: + window.open_path(target) # open the given project + else: + window._auto_open_last() # reopen the last project, if any return app.exec() diff --git a/hvideotool/core/detection/cache.py b/hvideotool/core/detection/cache.py new file mode 100644 index 0000000..fb3dfe8 --- /dev/null +++ b/hvideotool/core/detection/cache.py @@ -0,0 +1,77 @@ +"""Persist detection results for a project. + +The detection cache is a JSON file (``detections.json`` at the project root) that +maps each image to its detections so reopening a project doesn't have to re-run +the detector. The cache file lives apart from the images (which sit in the +project's ``frames/`` sub-folder), so the file location and the image base +directory are passed separately. + +The cache is tagged with the detector identity (name + model + conf/imgsz); a +mismatch means the cache was produced by a different detector and is ignored +(``load_results`` returns ``None``) rather than shown as if current. + +Keys are stored as **basenames**, so the cache survives moving/renaming the +project folder. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from .types import Detection + +_VERSION = 1 + + +def make_key(detector: str, model_path: str | None, yolo_conf: float, yolo_imgsz: int) -> dict: + """Identity of the detector that produced a cache; cache is only valid for a match.""" + return { + "detector": detector, + "model_path": model_path or "", + "yolo_conf": round(float(yolo_conf), 4), + "yolo_imgsz": int(yolo_imgsz), + } + + +def save_results(cache_file: Path, key: dict, results: dict[str, list[Detection]]) -> bool: + """Write the cache (basename -> detections) to ``cache_file``. False on failure.""" + payload = { + "version": _VERSION, + "key": key, + "results": { + Path(p).name: [d.to_dict() for d in dets] + for p, dets in results.items() + }, + } + try: + cache_file.parent.mkdir(parents=True, exist_ok=True) + cache_file.write_text( + json.dumps(payload, ensure_ascii=False), encoding="utf-8" + ) + return True + except OSError: + return False + + +def load_results(cache_file: Path, key: dict, base_dir: Path) -> dict[str, list[Detection]] | None: + """Load cached detections from ``cache_file`` if present and the detector matches. + + Returns a dict keyed by **full path** (``base_dir / basename``), or ``None`` if + there is no cache, it's unreadable, or it was made by a different detector. + """ + if not cache_file.is_file(): + return None + try: + payload = json.loads(cache_file.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if payload.get("version") != _VERSION or payload.get("key") != key: + return None + out: dict[str, list[Detection]] = {} + for name, dets in payload.get("results", {}).items(): + try: + out[str(base_dir / name)] = [Detection.from_dict(d) for d in dets] + except (KeyError, TypeError, ValueError): + continue # skip a corrupt entry, keep the rest + return out diff --git a/hvideotool/core/project.py b/hvideotool/core/project.py new file mode 100644 index 0000000..e7ee9f5 --- /dev/null +++ b/hvideotool/core/project.py @@ -0,0 +1,148 @@ +"""A HVideoTool *project*: a self-contained folder on disk. + +This replaces the old "open a bare folder of images" model. A project is a folder +that holds: + +``` +MyProject/ +├── project.json # version, name, created, source, settings{detector/model/threshold/restore} +├── frames/ # the images (what used to be "the folder") +├── detections.json # the detection cache (basename -> detections, tagged with detector key) +└── collections/ # curation sub-folders (created lazily) +``` + +The project file carries the **per-project** settings (detector, model, overlay +threshold, restore engine). Global ``settings.json`` only seeds the defaults for +*new* projects; once a project exists it remembers how it was last inspected. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path + +from ..config import AppConfig + +PROJECT_FILE = "project.json" +FRAMES_DIR = "frames" +CACHE_FILE = "detections.json" +COLLECTIONS_DIR = "collections" +FAVORITES_DIR = "Избранное" # the single default collection ("в избранное") +_VERSION = 1 + +# The subset of AppConfig fields a project remembers (mirrored to/from project.json). +_SETTING_KEYS = ( + "detector", + "model_path", + "default_threshold", + "restorer", + "dm_dir", + "dm_model", + "dm_python", + "dm_gpu", +) + + +@dataclass +class Project: + """A HVideoTool project rooted at ``root`` — the single source of truth for its + on-disk layout and its per-project settings.""" + + root: Path + name: str + settings: dict = field(default_factory=dict) # subset of AppConfig fields + source: str | None = None # originating video/folder, if any + created: str | None = None + + # ----------------------------------------------------------------- paths + @property + def project_file(self) -> Path: + return self.root / PROJECT_FILE + + @property + def frames_dir(self) -> Path: + return self.root / FRAMES_DIR + + @property + def cache_path(self) -> Path: + return self.root / CACHE_FILE + + @property + def collections_dir(self) -> Path: + return self.root / COLLECTIONS_DIR + + @property + def favorites_dir(self) -> Path: + """The single default collection — frames moved "to favorites" land here.""" + return self.collections_dir / FAVORITES_DIR + + # ------------------------------------------------------------- lifecycle + @classmethod + def create( + cls, + root: Path | str, + name: str | None = None, + settings: dict | None = None, + source: str | None = None, + ) -> "Project": + """Create a new project folder (with ``frames/``) and write ``project.json``.""" + root = Path(root) + proj = cls( + root=root, + name=name or root.name, + settings={k: v for k, v in (settings or {}).items() if k in _SETTING_KEYS}, + source=source, + created=datetime.now().isoformat(timespec="seconds"), + ) + proj.frames_dir.mkdir(parents=True, exist_ok=True) + proj.save() + return proj + + @classmethod + def load(cls, path: Path | str) -> "Project": + """Load a project from its folder or directly from its ``project.json``.""" + path = Path(path) + root = path.parent if path.name == PROJECT_FILE else path + data = json.loads((root / PROJECT_FILE).read_text(encoding="utf-8")) + return cls( + root=root, + name=data.get("name", root.name), + settings={k: v for k, v in data.get("settings", {}).items() if k in _SETTING_KEYS}, + source=data.get("source"), + created=data.get("created"), + ) + + @staticmethod + def is_project(path: Path | str) -> bool: + """True if ``path`` is a project folder (or a ``project.json``).""" + path = Path(path) + if path.name == PROJECT_FILE: + return path.is_file() + return (path / PROJECT_FILE).is_file() + + def save(self) -> None: + """Write ``project.json``.""" + payload = { + "version": _VERSION, + "name": self.name, + "created": self.created, + "source": self.source, + "settings": self.settings, + } + self.root.mkdir(parents=True, exist_ok=True) + self.project_file.write_text( + json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8" + ) + + # ----------------------------------------------------- settings <-> config + def apply_to_config(self, cfg: AppConfig) -> None: + """Overlay this project's stored settings onto ``cfg`` (mutates it).""" + for key in _SETTING_KEYS: + if key in self.settings: + setattr(cfg, key, self.settings[key]) + + def update_from_config(self, cfg: AppConfig) -> None: + """Capture the per-project settings from ``cfg`` into ``self.settings``.""" + self.settings = {key: getattr(cfg, key) for key in _SETTING_KEYS} diff --git a/hvideotool/core/restore/_deepmosaics/LICENSE b/hvideotool/core/restore/_deepmosaics/LICENSE new file mode 100644 index 0000000..f288702 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/LICENSE @@ -0,0 +1,674 @@ + GNU GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU General Public License is a free, copyleft license for +software and other kinds of works. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +the GNU General Public License is intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. We, the Free Software Foundation, use the +GNU General Public License for most of our software; it applies also to +any other work released this way by its authors. You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + To protect your rights, we need to prevent others from denying you +these rights or asking you to surrender the rights. Therefore, you have +certain responsibilities if you distribute copies of the software, or if +you modify it: responsibilities to respect the freedom of others. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must pass on to the recipients the same +freedoms that you received. You must make sure that they, too, receive +or can get the source code. And you must show them these terms so they +know their rights. + + Developers that use the GNU GPL protect your rights with two steps: +(1) assert copyright on the software, and (2) offer you this License +giving you legal permission to copy, distribute and/or modify it. + + For the developers' and authors' protection, the GPL clearly explains +that there is no warranty for this free software. For both users' and +authors' sake, the GPL requires that modified versions be marked as +changed, so that their problems will not be attributed erroneously to +authors of previous versions. + + Some devices are designed to deny users access to install or run +modified versions of the software inside them, although the manufacturer +can do so. This is fundamentally incompatible with the aim of +protecting users' freedom to change the software. The systematic +pattern of such abuse occurs in the area of products for individuals to +use, which is precisely where it is most unacceptable. Therefore, we +have designed this version of the GPL to prohibit the practice for those +products. If such problems arise substantially in other domains, we +stand ready to extend this provision to those domains in future versions +of the GPL, as needed to protect the freedom of users. + + Finally, every program is threatened constantly by software patents. +States should not allow patents to restrict development and use of +software on general-purpose computers, but in those that do, we wish to +avoid the special danger that patents applied to a free program could +make it effectively proprietary. To prevent this, the GPL assures that +patents cannot be used to render the program non-free. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Use with the GNU Affero General Public License. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU Affero General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the special requirements of the GNU Affero General Public License, +section 13, concerning interaction through a network will apply to the +combination as such. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If the program does terminal interaction, make it output a short +notice like this when it starts in an interactive mode: + + Copyright (C) + This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, your program's commands +might be different; for a GUI interface, you would use an "about box". + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU GPL, see +. + + The GNU General Public License does not permit incorporating your program +into proprietary programs. If your program is a subroutine library, you +may consider it more useful to permit linking proprietary applications with +the library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. But first, please read +. diff --git a/hvideotool/core/restore/_deepmosaics/NOTICE.md b/hvideotool/core/restore/_deepmosaics/NOTICE.md new file mode 100644 index 0000000..248c441 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/NOTICE.md @@ -0,0 +1,16 @@ +# Vendored DeepMosaics code + +`models/` and `util/` here are copied **verbatim** from DeepMosaics +(https://github.com/HypoX64/DeepMosaics) by HypoX64, licensed **GPL-3.0** +(see `LICENSE`). We use only the per-image mosaic-clean path +(`models.loadmodel`, `models.runmodel`, `util.image_processing`) — loaded +in-process by `hvideotool/core/restore/deepmosaics.py`. + +Because this GPL-3.0 code is combined into HVideoTool, the project as a whole is +distributed under **GPL-3.0**. + +Model weights (`clean_*.pth`, `mosaic_position.pth`) are NOT included — the user +points the app at their own downloaded weights. + +These files are unmodified; this directory is added to `sys.path` at import time +so the original `from models import …` / `import util.…` statements resolve. diff --git a/hvideotool/core/restore/_deepmosaics/util/__init__.py b/hvideotool/core/restore/_deepmosaics/util/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hvideotool/core/restore/_deepmosaics/util/clean_cache.py b/hvideotool/core/restore/_deepmosaics/util/clean_cache.py new file mode 100644 index 0000000..80c47b9 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/clean_cache.py @@ -0,0 +1,51 @@ +import os +import shutil + +def findalldir(rootdir): + dir_list = [] + for root,dirs,files in os.walk(rootdir): + for dir in dirs: + dir_list.append(os.path.join(root,dir)) + return(dir_list) + +def Traversal(filedir): + file_list=[] + dir_list = [] + for root,dirs,files in os.walk(filedir): + for file in files: + file_list.append(os.path.join(root,file)) + for dir in dirs: + dir_list.append(os.path.join(root,dir)) + Traversal(dir) + return file_list,dir_list + +def is_img(path): + ext = os.path.splitext(path)[1] + ext = ext.lower() + if ext in ['.jpg','.png','.jpeg','.bmp']: + return True + else: + return False + +def is_video(path): + ext = os.path.splitext(path)[1] + ext = ext.lower() + if ext in ['.mp4','.flv','.avi','.mov','.mkv','.wmv','.rmvb']: + return True + else: + return False + +def cleanall(): + file_list,dir_list = Traversal('./') + for file in file_list: + if ('tmp' in file) | ('pth' in file)|('pycache' in file) | is_video(file) | is_img(file): + if os.path.exists(file): + if 'imgs' not in file: + os.remove(file) + print('remove file:',file) + + for dir in dir_list: + if ('tmp'in dir)|('pycache'in dir): + if os.path.exists(dir): + shutil.rmtree(dir) + print('remove dir:',dir) \ No newline at end of file diff --git a/hvideotool/core/restore/_deepmosaics/util/data.py b/hvideotool/core/restore/_deepmosaics/util/data.py new file mode 100644 index 0000000..8a7865e --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/data.py @@ -0,0 +1,160 @@ +import random +import os +from util.mosaic import get_random_parameter +import numpy as np +import torch +import torchvision.transforms as transforms +import cv2 +from . import image_processing as impro +from . import degradater + +def to_tensor(data,gpu_id): + data = torch.from_numpy(data) + if gpu_id != '-1': + data = data.cuda() + return data + +def normalize(data): + ''' + normalize to -1 ~ 1 + ''' + return (data.astype(np.float32)/255.0-0.5)/0.5 + +def anti_normalize(data): + return np.clip((data*0.5+0.5)*255,0,255).astype(np.uint8) + +def tensor2im(image_tensor, gray=False, rgb2bgr = True ,is0_1 = False, batch_index=0): + image_tensor =image_tensor.data + image_numpy = image_tensor[batch_index].cpu().float().numpy() + + if not is0_1: + image_numpy = (image_numpy + 1)/2.0 + image_numpy = np.clip(image_numpy * 255.0,0,255) + + # gray -> output 1ch + if gray: + h, w = image_numpy.shape[1:] + image_numpy = image_numpy.reshape(h,w) + return image_numpy.astype(np.uint8) + + # output 3ch + if image_numpy.shape[0] == 1: + image_numpy = np.tile(image_numpy, (3, 1, 1)) + image_numpy = image_numpy.transpose((1, 2, 0)) + if rgb2bgr and not gray: + image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy) + return image_numpy.astype(np.uint8) + + +def im2tensor(image_numpy, gray=False,bgr2rgb = True, reshape = True, gpu_id = '-1',is0_1 = False): + if gray: + h, w = image_numpy.shape + image_numpy = (image_numpy/255.0-0.5)/0.5 + image_tensor = torch.from_numpy(image_numpy).float() + if reshape: + image_tensor = image_tensor.reshape(1,1,h,w) + else: + h, w ,ch = image_numpy.shape + if bgr2rgb: + image_numpy = image_numpy[...,::-1]-np.zeros_like(image_numpy) + if is0_1: + image_numpy = image_numpy/255.0 + else: + image_numpy = (image_numpy/255.0-0.5)/0.5 + image_numpy = image_numpy.transpose((2, 0, 1)) + image_tensor = torch.from_numpy(image_numpy).float() + if reshape: + image_tensor = image_tensor.reshape(1,ch,h,w) + if gpu_id != '-1': + image_tensor = image_tensor.cuda() + return image_tensor + +def shuffledata(data,target): + state = np.random.get_state() + np.random.shuffle(data) + np.random.set_state(state) + np.random.shuffle(target) + +def random_transform_single_mask(img,out_shape): + out_h,out_w = out_shape + img = cv2.resize(img,(int(out_w*random.uniform(1.1, 1.5)),int(out_h*random.uniform(1.1, 1.5)))) + h,w = img.shape[:2] + h_move = int((h-out_h)*random.random()) + w_move = int((w-out_w)*random.random()) + img = img[h_move:h_move+out_h,w_move:w_move+out_w] + if random.random()<0.5: + if random.random()<0.5: + img = img[:,::-1] + else: + img = img[::-1,:] + if img.shape[0] != out_h or img.shape[1]!= out_w : + img = cv2.resize(img,(out_w,out_h)) + return img + +def get_transform_params(): + crop_flag = True + rotat_flag = np.random.random()<0.2 + color_flag = True + flip_flag = np.random.random()<0.2 + degradate_flag = np.random.random()<0.5 + flag_dict = {'crop':crop_flag,'rotat':rotat_flag,'color':color_flag,'flip':flip_flag,'degradate':degradate_flag} + + crop_rate = [np.random.random(),np.random.random()] + rotat_rate = np.random.random() + color_rate = [np.random.uniform(-0.05,0.05),np.random.uniform(-0.05,0.05),np.random.uniform(-0.05,0.05), + np.random.uniform(-0.05,0.05),np.random.uniform(-0.05,0.05)] + flip_rate = np.random.random() + degradate_params = degradater.get_random_degenerate_params(mod='weaker_2') + rate_dict = {'crop':crop_rate,'rotat':rotat_rate,'color':color_rate,'flip':flip_rate,'degradate':degradate_params} + + return {'flag':flag_dict,'rate':rate_dict} + +def random_transform_single_image(img,finesize,params=None,test_flag = False): + if params is None: + params = get_transform_params() + + if params['flag']['degradate']: + img = degradater.degradate(img,params['rate']['degradate']) + + if params['flag']['crop']: + h,w = img.shape[:2] + h_move = int((h-finesize)*params['rate']['crop'][0]) + w_move = int((w-finesize)*params['rate']['crop'][1]) + img = img[h_move:h_move+finesize,w_move:w_move+finesize] + + if test_flag: + return img + + if params['flag']['rotat']: + h,w = img.shape[:2] + M = cv2.getRotationMatrix2D((w/2,h/2),90*int(4*params['rate']['rotat']),1) + img = cv2.warpAffine(img,M,(w,h)) + + if params['flag']['color']: + img = impro.color_adjust(img,params['rate']['color'][0],params['rate']['color'][1], + params['rate']['color'][2],params['rate']['color'][3],params['rate']['color'][4]) + + if params['flag']['flip']: + img = img[:,::-1] + + #check shape + if img.shape[0]!= finesize or img.shape[1]!= finesize: + img = cv2.resize(img,(finesize,finesize)) + print('warning! shape error.') + return img + +def random_transform_pair_image(img,mask,finesize,test_flag = False): + params = get_transform_params() + img = random_transform_single_image(img,finesize,params) + params['flag']['degradate'] = False + params['flag']['color'] = False + mask = random_transform_single_image(mask,finesize,params) + return img,mask + +def showresult(img1,img2,img3,name,is0_1 = False): + size = img1.shape[3] + showimg=np.zeros((size,size*3,3)) + showimg[0:size,0:size] = tensor2im(img1,rgb2bgr = False, is0_1 = is0_1) + showimg[0:size,size:size*2] = tensor2im(img2,rgb2bgr = False, is0_1 = is0_1) + showimg[0:size,size*2:size*3] = tensor2im(img3,rgb2bgr = False, is0_1 = is0_1) + cv2.imwrite(name, showimg) diff --git a/hvideotool/core/restore/_deepmosaics/util/dataloader.py b/hvideotool/core/restore/_deepmosaics/util/dataloader.py new file mode 100644 index 0000000..334b38f --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/dataloader.py @@ -0,0 +1,141 @@ +import os +import random +import numpy as np +from multiprocessing import Process, Queue +from . import image_processing as impro +from . import mosaic,data + +class VideoLoader(object): + """docstring for VideoLoader + Load a single video(Converted to images) + How to use: + 1.Init VideoLoader as loader + 2.Get data by loader.ori_stream + 3.loader.next() to get next stream + """ + def __init__(self, opt, video_dir, test_flag=False): + super(VideoLoader, self).__init__() + self.opt = opt + self.test_flag = test_flag + self.video_dir = video_dir + self.t = 0 + self.n_iter = self.opt.M -self.opt.S*(self.opt.T+1) + self.transform_params = data.get_transform_params() + self.ori_load_pool = [] + self.mosaic_load_pool = [] + self.previous_pred = None + feg_ori = impro.imread(os.path.join(video_dir,'origin_image','00001.jpg'),loadsize=self.opt.loadsize,rgb=True) + feg_mask = impro.imread(os.path.join(video_dir,'mask','00001.png'),mod='gray',loadsize=self.opt.loadsize) + self.mosaic_size,self.mod,self.rect_rat,self.feather = mosaic.get_random_parameter(feg_ori,feg_mask) + self.startpos = [random.randint(0,self.mosaic_size),random.randint(0,self.mosaic_size)] + self.loadsize = self.opt.loadsize + #Init load pool + for i in range(self.opt.S*self.opt.T): + _ori_img = impro.imread(os.path.join(video_dir,'origin_image','%05d' % (i+1)+'.jpg'),loadsize=self.loadsize,rgb=True) + _mask = impro.imread(os.path.join(video_dir,'mask','%05d' % (i+1)+'.png' ),mod='gray',loadsize=self.loadsize) + _mosaic_img = mosaic.addmosaic_base(_ori_img, _mask, self.mosaic_size,0, self.mod,self.rect_rat,self.feather,self.startpos) + _ori_img = data.random_transform_single_image(_ori_img,opt.finesize,self.transform_params) + _mosaic_img = data.random_transform_single_image(_mosaic_img,opt.finesize,self.transform_params) + + self.ori_load_pool.append(self.normalize(_ori_img)) + self.mosaic_load_pool.append(self.normalize(_mosaic_img)) + self.ori_load_pool = np.array(self.ori_load_pool) + self.mosaic_load_pool = np.array(self.mosaic_load_pool) + + #Init frist stream + self.ori_stream = self.ori_load_pool [np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() + self.mosaic_stream = self.mosaic_load_pool[np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() + # stream B,T,H,W,C -> B,C,T,H,W + self.ori_stream = self.ori_stream.reshape (1,self.opt.T,opt.finesize,opt.finesize,3).transpose((0,4,1,2,3)) + self.mosaic_stream = self.mosaic_stream.reshape(1,self.opt.T,opt.finesize,opt.finesize,3).transpose((0,4,1,2,3)) + + #Init frist previous frame + self.previous_pred = self.ori_load_pool[self.opt.S*self.opt.N-1].copy() + # previous B,C,H,W + self.previous_pred = self.previous_pred.reshape(1,opt.finesize,opt.finesize,3).transpose((0,3,1,2)) + + def normalize(self,data): + ''' + normalize to -1 ~ 1 + ''' + return (data.astype(np.float32)/255.0-0.5)/0.5 + + def anti_normalize(self,data): + return np.clip((data*0.5+0.5)*255,0,255).astype(np.uint8) + + def next(self): + # random + if np.random.random()<0.05: + self.startpos = [random.randint(0,self.mosaic_size),random.randint(0,self.mosaic_size)] + if np.random.random()<0.02: + self.transform_params['rate']['crop'] = [np.random.random(),np.random.random()] + if np.random.random()<0.02: + self.loadsize = np.random.randint(self.opt.finesize,self.opt.loadsize) + + if self.t != 0: + self.previous_pred = None + self.ori_load_pool [:self.opt.S*self.opt.T-1] = self.ori_load_pool [1:self.opt.S*self.opt.T] + self.mosaic_load_pool[:self.opt.S*self.opt.T-1] = self.mosaic_load_pool[1:self.opt.S*self.opt.T] + #print(os.path.join(self.video_dir,'origin_image','%05d' % (self.opt.S*self.opt.T+self.t)+'.jpg')) + _ori_img = impro.imread(os.path.join(self.video_dir,'origin_image','%05d' % (self.opt.S*self.opt.T+self.t)+'.jpg'),loadsize=self.loadsize,rgb=True) + _mask = impro.imread(os.path.join(self.video_dir,'mask','%05d' % (self.opt.S*self.opt.T+self.t)+'.png' ),mod='gray',loadsize=self.loadsize) + _mosaic_img = mosaic.addmosaic_base(_ori_img, _mask, self.mosaic_size,0, self.mod,self.rect_rat,self.feather,self.startpos) + _ori_img = data.random_transform_single_image(_ori_img,self.opt.finesize,self.transform_params) + _mosaic_img = data.random_transform_single_image(_mosaic_img,self.opt.finesize,self.transform_params) + + _ori_img,_mosaic_img = self.normalize(_ori_img),self.normalize(_mosaic_img) + self.ori_load_pool [self.opt.S*self.opt.T-1] = _ori_img + self.mosaic_load_pool[self.opt.S*self.opt.T-1] = _mosaic_img + + self.ori_stream = self.ori_load_pool [np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() + self.mosaic_stream = self.mosaic_load_pool[np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() + + # stream B,T,H,W,C -> B,C,T,H,W + self.ori_stream = self.ori_stream.reshape (1,self.opt.T,self.opt.finesize,self.opt.finesize,3).transpose((0,4,1,2,3)) + self.mosaic_stream = self.mosaic_stream.reshape(1,self.opt.T,self.opt.finesize,self.opt.finesize,3).transpose((0,4,1,2,3)) + + self.t += 1 + +class VideoDataLoader(object): + """VideoDataLoader""" + def __init__(self, opt, videolist, test_flag=False): + super(VideoDataLoader, self).__init__() + self.videolist = [] + self.opt = opt + self.test_flag = test_flag + for i in range(self.opt.n_epoch): + self.videolist += videolist.copy() + random.shuffle(self.videolist) + self.each_video_n_iter = self.opt.M -self.opt.S*(self.opt.T+1) + self.n_iter = len(self.videolist)//self.opt.load_thread//self.opt.batchsize*self.each_video_n_iter*self.opt.load_thread + self.queue = Queue(self.opt.load_thread) + self.ori_stream = np.zeros((self.opt.batchsize,3,self.opt.T,self.opt.finesize,self.opt.finesize),dtype=np.float32)# B,C,T,H,W + self.mosaic_stream = np.zeros((self.opt.batchsize,3,self.opt.T,self.opt.finesize,self.opt.finesize),dtype=np.float32)# B,C,T,H,W + self.previous_pred = np.zeros((self.opt.batchsize,3,self.opt.finesize,self.opt.finesize),dtype=np.float32) + self.load_init() + + def load(self,videolist): + for load_video_iter in range(len(videolist)//self.opt.batchsize): + iter_videolist = videolist[load_video_iter*self.opt.batchsize:(load_video_iter+1)*self.opt.batchsize] + videoloaders = [VideoLoader(self.opt,os.path.join(self.opt.dataset,iter_videolist[i]),self.test_flag) for i in range(self.opt.batchsize)] + for each_video_iter in range(self.each_video_n_iter): + for i in range(self.opt.batchsize): + self.ori_stream[i] = videoloaders[i].ori_stream + self.mosaic_stream[i] = videoloaders[i].mosaic_stream + if each_video_iter == 0: + self.previous_pred[i] = videoloaders[i].previous_pred + videoloaders[i].next() + if each_video_iter == 0: + self.queue.put([self.ori_stream.copy(),self.mosaic_stream.copy(),self.previous_pred]) + else: + self.queue.put([self.ori_stream.copy(),self.mosaic_stream.copy(),None]) + + def load_init(self): + ptvn = len(self.videolist)//self.opt.load_thread #pre_thread_video_num + for i in range(self.opt.load_thread): + p = Process(target=self.load,args=(self.videolist[i*ptvn:(i+1)*ptvn],)) + p.daemon = True + p.start() + + def get_data(self): + return self.queue.get() \ No newline at end of file diff --git a/hvideotool/core/restore/_deepmosaics/util/degradater.py b/hvideotool/core/restore/_deepmosaics/util/degradater.py new file mode 100644 index 0000000..9e4d227 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/degradater.py @@ -0,0 +1,119 @@ +''' +https://github.com/sonack/GFRNet_pytorch_new +''' +import random +import cv2 +import numpy as np + +def gaussian_blur(img, sigma=3, size=13): + if sigma > 0: + if isinstance(size, int): + size = (size, size) + img = cv2.GaussianBlur(img, size, sigma) + return img + +def down(img, scale, shape): + if scale > 1: + h, w, _ = shape + scaled_h, scaled_w = int(h / scale), int(w / scale) + img = cv2.resize(img, (scaled_w, scaled_h), interpolation = cv2.INTER_CUBIC) + return img + +def up(img, scale, shape): + if scale > 1: + h, w, _ = shape + img = cv2.resize(img, (w, h), interpolation = cv2.INTER_CUBIC) + return img + +def awgn(img, level): + if level > 0: + noise = np.random.randn(*img.shape) * level + img = (img + noise).clip(0,255).astype(np.uint8) + return img + +def jpeg_compressor(img,quality): + if quality > 0: # 0 indicating no lossy compression (i.e losslessly compression) + encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality] + img = cv2.imdecode(cv2.imencode('.jpg', img, encode_param)[1], 1) + return img + +def get_random_degenerate_params(mod='strong'): + ''' + mod : strong | only_downsample | only_4x | weaker_1 | weaker_2 + ''' + params = {} + gaussianBlur_size_list = list(range(3,14,2)) + + if mod == 'strong': + gaussianBlur_sigma_list = [1 + x for x in range(3)] + gaussianBlur_sigma_list += [0] + downsample_scale_list = [1 + x * 0.1 for x in range(0,71)] + awgn_level_list = list(range(1, 8, 1)) + jpeg_quality_list = list(range(10, 41, 1)) + jpeg_quality_list += int(len(jpeg_quality_list) * 0.33) * [0] + + elif mod == 'only_downsample': + gaussianBlur_sigma_list = [0] + downsample_scale_list = [1 + x * 0.1 for x in range(0,71)] + awgn_level_list = [0] + jpeg_quality_list = [0] + + elif mod == 'only_4x': + gaussianBlur_sigma_list = [0] + downsample_scale_list = [4] + awgn_level_list = [0] + jpeg_quality_list = [0] + + elif mod == 'weaker_1': # 0.5 trigger prob + gaussianBlur_sigma_list = [1 + x for x in range(3)] + gaussianBlur_sigma_list += int(len(gaussianBlur_sigma_list)) * [0] # 1/2 trigger this degradation + + downsample_scale_list = [1 + x * 0.1 for x in range(0,71)] + downsample_scale_list += int(len(downsample_scale_list)) * [1] + + awgn_level_list = list(range(1, 8, 1)) + awgn_level_list += int(len(awgn_level_list)) * [0] + + jpeg_quality_list = list(range(10, 41, 1)) + jpeg_quality_list += int(len(jpeg_quality_list)) * [0] + + elif mod == 'weaker_2': # weaker than weaker_1, jpeg [20,40] + gaussianBlur_sigma_list = [1 + x for x in range(3)] + gaussianBlur_sigma_list += int(len(gaussianBlur_sigma_list)) * [0] # 1/2 trigger this degradation + + downsample_scale_list = [1 + x * 0.1 for x in range(0,71)] + downsample_scale_list += int(len(downsample_scale_list)) * [1] + + awgn_level_list = list(range(1, 8, 1)) + awgn_level_list += int(len(awgn_level_list)) * [0] + + jpeg_quality_list = list(range(20, 41, 1)) + jpeg_quality_list += int(len(jpeg_quality_list)) * [0] + + params['blur_sigma'] = random.choice(gaussianBlur_sigma_list) + params['blur_size'] = random.choice(gaussianBlur_size_list) + params['updown_scale'] = random.choice(downsample_scale_list) + params['awgn_level'] = random.choice(awgn_level_list) + params['jpeg_quality'] = random.choice(jpeg_quality_list) + + return params + +def degradate(img,params,jpeg_last = True): + shape = img.shape + if not params: + params = get_random_degenerate_params('original') + + if jpeg_last: + img = gaussian_blur(img,params['blur_sigma'],params['blur_size']) + img = down(img,params['updown_scale'],shape) + img = awgn(img,params['awgn_level']) + img = up(img,params['updown_scale'],shape) + img = jpeg_compressor(img,params['jpeg_quality']) + else: + img = gaussian_blur(img,params['blur_sigma'],params['blur_size']) + img = down(img,params['updown_scale'],shape) + img = awgn(img,params['awgn_level']) + img = jpeg_compressor(img,params['jpeg_quality']) + img = up(img,params['updown_scale'],shape) + + return img \ No newline at end of file diff --git a/hvideotool/core/restore/_deepmosaics/util/ffmpeg.py b/hvideotool/core/restore/_deepmosaics/util/ffmpeg.py new file mode 100644 index 0000000..6efd686 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/ffmpeg.py @@ -0,0 +1,92 @@ +import os,json +import subprocess +# ffmpeg 3.4.6 + +def args2cmd(args): + cmd = '' + for arg in args: + cmd += (arg+' ') + return cmd + +def run(args,mode = 0): + + if mode == 0: + cmd = args2cmd(args) + os.system(cmd) + + elif mode == 1: + ''' + out_string = os.popen(cmd_str).read() + For chinese path in Windows + https://blog.csdn.net/weixin_43903378/article/details/91979025 + ''' + cmd = args2cmd(args) + stream = os.popen(cmd)._stream + sout = stream.buffer.read().decode(encoding='utf-8') + return sout + + elif mode == 2: + cmd = args2cmd(args) + p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + sout = p.stdout.readlines() + return sout + +def video2image(videopath, imagepath, fps=0, start_time='00:00:00', last_time='00:00:00'): + args = ['ffmpeg'] + if last_time != '00:00:00': + args += ['-ss', start_time] + args += ['-t', last_time] + args += ['-i', '"'+videopath+'"'] + if fps != 0: + args += ['-r', str(fps)] + args += ['-f', 'image2','-q:v','-0',imagepath] + run(args) + +def video2voice(videopath, voicepath, start_time='00:00:00', last_time='00:00:00'): + args = ['ffmpeg', '-i', '"'+videopath+'"','-async 1 -f mp3','-b:a 320k'] + if last_time != '00:00:00': + args += ['-ss', start_time] + args += ['-t', last_time] + args += [voicepath] + run(args) + +def image2video(fps,imagepath,voicepath,videopath): + os.system('ffmpeg -y -r '+str(fps)+' -i '+imagepath+' -vcodec libx264 '+os.path.split(voicepath)[0]+'/video_tmp.mp4') + if os.path.exists(voicepath): + os.system('ffmpeg -i '+os.path.split(voicepath)[0]+'/video_tmp.mp4'+' -i "'+voicepath+'" -vcodec copy -acodec aac '+videopath) + else: + os.system('ffmpeg -i '+os.path.split(voicepath)[0]+'/video_tmp.mp4 '+videopath) + +def get_video_infos(videopath): + args = ['ffprobe -v quiet -print_format json -show_format -show_streams', '-i', '"'+videopath+'"'] + out_string = run(args,mode=1) + infos = json.loads(out_string) + try: + fps = eval(infos['streams'][0]['avg_frame_rate']) + endtime = float(infos['format']['duration']) + width = int(infos['streams'][0]['width']) + height = int(infos['streams'][0]['height']) + except Exception as e: + fps = eval(infos['streams'][1]['r_frame_rate']) + endtime = float(infos['format']['duration']) + width = int(infos['streams'][1]['width']) + height = int(infos['streams'][1]['height']) + + return fps,endtime,height,width + +def cut_video(in_path,start_time,last_time,out_path,vcodec='h265'): + if vcodec == 'copy': + os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec copy -acodec copy '+out_path) + elif vcodec == 'h264': + os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec libx264 -b 12M '+out_path) + elif vcodec == 'h265': + os.system('ffmpeg -ss '+start_time+' -t '+last_time+' -i "'+in_path+'" -vcodec libx265 -b 12M '+out_path) + +def continuous_screenshot(videopath,savedir,fps): + ''' + videopath: input video path + savedir: images will save here + fps: save how many images per second + ''' + videoname = os.path.splitext(os.path.basename(videopath))[0] + os.system('ffmpeg -i "'+videopath+'" -vf fps='+str(fps)+' -q:v -0 '+savedir+'/'+videoname+'_%06d.jpg') diff --git a/hvideotool/core/restore/_deepmosaics/util/filt.py b/hvideotool/core/restore/_deepmosaics/util/filt.py new file mode 100644 index 0000000..b99d6e1 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/filt.py @@ -0,0 +1,77 @@ +import numpy as np + +def less_zero(arr,num = 7): + index = np.linspace(0,len(arr)-1,len(arr),dtype='int') + cnt = 0 + for i in range(2,len(arr)-2): + if arr[i] != 0: + arr[i] = arr[i] + if cnt != 0: + if cnt <= num*2: + arr[i-cnt:round(i-cnt/2)] = arr[i-cnt-1-2] + arr[round(i-cnt/2):i] = arr[i+2] + index[i-cnt:round(i-cnt/2)] = i-cnt-1-2 + index[round(i-cnt/2):i] = i+2 + else: + arr[i-cnt:i-cnt+num] = arr[i-cnt-1-2] + arr[i-num:i] = arr[i+2] + index[i-cnt:i-cnt+num] = i-cnt-1-2 + index[i-num:i] = i+2 + cnt = 0 + else: + cnt += 1 + return arr,index + +def medfilt(data,window): + if window%2 == 0 or window < 0: + print('Error: the medfilt window must be even number') + exit(0) + pad = int((window-1)/2) + pad_data = np.zeros(len(data)+window-1, dtype = type(data[0])) + result = np.zeros(len(data),dtype = type(data[0])) + pad_data[pad:pad+len(data)]=data[:] + for i in range(len(data)): + result[i] = np.median(pad_data[i:i+window]) + return result + +def position_medfilt(positions,window): + + x,mask_index = less_zero(positions[:,0],window) + y = less_zero(positions[:,1],window)[0] + area = less_zero(positions[:,2],window)[0] + x_filt = medfilt(x, window) + y_filt = medfilt(y, window) + area_filt = medfilt(area, window) + cnt = 0 + for i in range(1,len(x)): + if 0.8original + ''' + if system_type == 'Linux': + if mod == 'normal': + img = cv2.imread(file_path,1) + elif mod == 'gray': + img = cv2.imread(file_path,0) + elif mod == 'all': + img = cv2.imread(file_path,-1) + + #In windows, for chinese path, use cv2.imdecode insteaded. + #It will loss EXIF, I can't fix it + else: + if mod == 'normal': + img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),1) + elif mod == 'gray': + img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),0) + elif mod == 'all': + img = cv2.imdecode(np.fromfile(file_path,dtype=np.uint8),-1) + + if loadsize != 0: + img = resize(img, loadsize, interpolation=cv2.INTER_CUBIC) + + if rgb and img.ndim==3: + img = img[:,:,::-1] + + return img + +def imwrite(file_path,img,use_thread=False): + ''' + in other to save chinese path images in windows, + this fun just for save final output images + ''' + def subfun(file_path,img): + if system_type == 'Linux': + cv2.imwrite(file_path, img) + else: + cv2.imencode('.jpg', img)[1].tofile(file_path) + if use_thread: + t = Thread(target=subfun,args=(file_path, img,)) + t.daemon() + t.start + else: + subfun(file_path,img) + +def resize(img,size,interpolation=cv2.INTER_LINEAR): + ''' + cv2.INTER_NEAREST      最邻近插值点法 + cv2.INTER_LINEAR        双线性插值法 + cv2.INTER_AREA         邻域像素再取样插补 + cv2.INTER_CUBIC        双立方插补，4*4大小的补点 + cv2.INTER_LANCZOS4 8x8像素邻域的Lanczos插值 + ''' + h, w = img.shape[:2] + if np.min((w,h)) ==size: + return img + if w >= h: + res = cv2.resize(img,(int(size*w/h), size),interpolation=interpolation) + else: + res = cv2.resize(img,(size, int(size*h/w)),interpolation=interpolation) + return res + +def resize_like(img,img_like): + h, w = img_like.shape[:2] + img = cv2.resize(img, (w,h)) + return img + +def ch_one2three(img): + res = cv2.merge([img, img, img]) + return res + +def color_adjust(img,alpha=0,beta=0,b=0,g=0,r=0,ran = False): + ''' + g(x) = (1+α)g(x)+255*β, + g(x) = g(x[:+b*255,:+g*255,:+r*255]) + + Args: + img : input image + alpha : contrast + beta : brightness + b : blue hue + g : green hue + r : red hue + ran : if True, randomly generated color correction parameters + Retuens: + img : output image + ''' + img = img.astype('float') + if ran: + alpha = random.uniform(-0.1,0.1) + beta = random.uniform(-0.1,0.1) + b = random.uniform(-0.05,0.05) + g = random.uniform(-0.05,0.05) + r = random.uniform(-0.05,0.05) + img = (1+alpha)*img+255.0*beta + bgr = [b*255.0,g*255.0,r*255.0] + for i in range(3): img[:,:,i]=img[:,:,i]+bgr[i] + + return (np.clip(img,0,255)).astype('uint8') + +def CAdaIN(src,dst): + ''' + make src has dst's style + ''' + return np.std(dst)*((src-np.mean(src))/np.std(src))+np.mean(dst) + +def makedataset(target_image,orgin_image): + target_image = resize(target_image,256) + orgin_image = resize(orgin_image,256) + img = np.zeros((256,512,3), dtype = "uint8") + w = orgin_image.shape[1] + img[0:256,0:256] = target_image[0:256,int(w/2-256/2):int(w/2+256/2)] + img[0:256,256:512] = orgin_image[0:256,int(w/2-256/2):int(w/2+256/2)] + return img + +def find_mostlikely_ROI(mask): + contours,hierarchy=cv2.findContours(mask, cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE) + if len(contours)>0: + areas = [] + for contour in contours: + areas.append(cv2.contourArea(contour)) + index = areas.index(max(areas)) + mask = np.zeros_like(mask) + mask = cv2.fillPoly(mask,[contours[index]],(255)) + return mask + +def boundingSquare(mask,Ex_mul): + # thresh = mask_threshold(mask,10,threshold) + area = mask_area(mask) + if area == 0 : + return 0,0,0,0 + + x,y,w,h = cv2.boundingRect(mask) + + center = np.array([int(x+w/2),int(y+h/2)]) + size = max(w,h) + point0=np.array([x,y]) + point1=np.array([x+size,y+size]) + + h, w = mask.shape[:2] + if size*Ex_mul > min(h, w): + size = min(h, w) + halfsize = int(min(h, w)/2) + else: + size = Ex_mul*size + halfsize = int(size/2) + size = halfsize*2 + point0 = center - halfsize + point1 = center + halfsize + if point0[0]<0: + point0[0]=0 + point1[0]=size + if point0[1]<0: + point0[1]=0 + point1[1]=size + if point1[0]>w: + point1[0]=w + point0[0]=w-size + if point1[1]>h: + point1[1]=h + point0[1]=h-size + center = ((point0+point1)/2).astype('int') + return center[0],center[1],halfsize,area + +def mask_threshold(mask,ex_mun,threshold): + mask = cv2.threshold(mask,threshold,255,cv2.THRESH_BINARY)[1] + mask = cv2.blur(mask, (ex_mun, ex_mun)) + mask = cv2.threshold(mask,threshold/5,255,cv2.THRESH_BINARY)[1] + return mask + +def mask_area(mask): + mask = cv2.threshold(mask,127,255,0)[1] + # contours= cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[1] #for opencv 3.4 + contours= cv2.findContours(mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)[0]#updata to opencv 4.0 + try: + area = cv2.contourArea(contours[0]) + except: + area = 0 + return area + +def replace_mosaic(img_origin,img_fake,mask,x,y,size,no_feather): + img_fake = cv2.resize(img_fake,(size*2,size*2),interpolation=cv2.INTER_CUBIC) + if no_feather: + img_origin[y-size:y+size,x-size:x+size]=img_fake + return img_origin + else: + # #color correction + # RGB_origin = img_origin[y-size:y+size,x-size:x+size].mean(0).mean(0) + # RGB_fake = img_fake.mean(0).mean(0) + # for i in range(3):img_fake[:,:,i] = np.clip(img_fake[:,:,i]+RGB_origin[i]-RGB_fake[i],0,255) + #eclosion + eclosion_num = int(size/10)+2 + + mask_crop = cv2.resize(mask,(img_origin.shape[1],img_origin.shape[0]))[y-size:y+size,x-size:x+size] + mask_crop = ch_one2three(mask_crop) + + mask_crop = (cv2.blur(mask_crop, (eclosion_num, eclosion_num))) + mask_crop = mask_crop/255.0 + + img_crop = img_origin[y-size:y+size,x-size:x+size] + img_origin[y-size:y+size,x-size:x+size] = np.clip((img_crop*(1-mask_crop)+img_fake*mask_crop),0,255).astype('uint8') + + return img_origin + + +def Q_lapulase(resImg): + ''' + Evaluate image quality + score > 20 normal + score > 50 clear + ''' + img2gray = cv2.cvtColor(resImg, cv2.COLOR_BGR2GRAY) + img2gray = resize(img2gray,512) + res = cv2.Laplacian(img2gray, cv2.CV_64F) + score = res.var() + return score + +def psnr(img1,img2): + mse = np.mean((img1/255.0-img2/255.0)**2) + if mse < 1e-10: + return 100 + psnr_v = 20*np.log10(1/np.sqrt(mse)) + return psnr_v + +def splice(imgs,splice_shape): + '''Stitching multiple images, all imgs must have the same size + imgs : [img1,img2,img3,img4] + splice_shape: (2,2) + ''' + h,w,ch = imgs[0].shape + output = np.zeros((h*splice_shape[0],w*splice_shape[1],ch),np.uint8) + cnt = 0 + for i in range(splice_shape[0]): + for j in range(splice_shape[1]): + if cnt < len(imgs): + output[h*i:h*(i+1),w*j:w*(j+1)] = imgs[cnt] + cnt += 1 + return output + diff --git a/hvideotool/core/restore/_deepmosaics/util/mosaic.py b/hvideotool/core/restore/_deepmosaics/util/mosaic.py new file mode 100644 index 0000000..c76a248 --- /dev/null +++ b/hvideotool/core/restore/_deepmosaics/util/mosaic.py @@ -0,0 +1,164 @@ +import cv2 +import numpy as np +import os +import random +from .image_processing import resize,ch_one2three,mask_area + +def addmosaic(img,mask,opt): + if opt.mosaic_mod == 'random': + img = addmosaic_random(img,mask) + elif opt.mosaic_size == 0: + img = addmosaic_autosize(img, mask, opt.mosaic_mod) + else: + img = addmosaic_base(img,mask,opt.mosaic_size,opt.output_size,model = opt.mosaic_mod) + return img + +def addmosaic_base(img,mask,n,out_size = 0,model = 'squa_avg',rect_rat = 1.6,feather=0,start_point=[0,0]): + ''' + img: input image + mask: input mask + n: mosaic size + out_size: output size 0->original + model : squa_avg squa_mid squa_random squa_avg_circle_edge rect_avg + rect_rat: if model==rect_avg , mosaic w/h=rect_rat + feather : feather size, -1->no 0->auto + start_point : [0,0], please not input this parameter + ''' + n = int(n) + + h_start = np.clip(start_point[0], 0, n) + w_start = np.clip(start_point[1], 0, n) + pix_mid_h = n//2+h_start + pix_mid_w = n//2+w_start + h, w = img.shape[:2] + h_step = (h-h_start)//n + w_step = (w-w_start)//n + if out_size: + img = resize(img,out_size) + if mask.shape[0] != h: + mask = cv2.resize(mask,(w,h)) + img_mosaic = img.copy() + + if model=='squa_avg': + for i in range(h_step): + for j in range(w_step): + if mask[i*n+pix_mid_h,j*n+pix_mid_w]: + img_mosaic[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:]=\ + img[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:].mean(axis=(0,1)) + + elif model=='squa_mid': + for i in range(h_step): + for j in range(w_step): + if mask[i*n+pix_mid_h,j*n+pix_mid_w]: + img_mosaic[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:]=\ + img[i*n+n//2+h_start,j*n+n//2+w_start,:] + + elif model == 'squa_random': + for i in range(h_step): + for j in range(w_step): + if mask[i*n+pix_mid_h,j*n+pix_mid_w]: + img_mosaic[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:]=\ + img[h_start+int(i*n-n/2+n*random.random()),w_start+int(j*n-n/2+n*random.random()),:] + + elif model == 'squa_avg_circle_edge': + for i in range(h_step): + for j in range(w_step): + img_mosaic[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:]=\ + img[i*n+h_start:(i+1)*n+h_start,j*n+w_start:(j+1)*n+w_start,:].mean(axis=(0,1)) + mask = cv2.threshold(mask,127,255,cv2.THRESH_BINARY)[1] + _mask = ch_one2three(mask) + mask_inv = cv2.bitwise_not(_mask) + imgroi1 = cv2.bitwise_and(_mask,img_mosaic) + imgroi2 = cv2.bitwise_and(mask_inv,img) + img_mosaic = cv2.add(imgroi1,imgroi2) + + elif model =='rect_avg': + n_h = n + n_w = int(n*rect_rat) + n_h_half = n_h//2+h_start + n_w_half = n_w//2+w_start + for i in range((h-h_start)//n_h): + for j in range((w-w_start)//n_w): + if mask[i*n_h+n_h_half,j*n_w+n_w_half]: + img_mosaic[i*n_h+h_start:(i+1)*n_h+h_start,j*n_w+w_start:(j+1)*n_w+w_start,:]=\ + img[i*n_h+h_start:(i+1)*n_h+h_start,j*n_w+w_start:(j+1)*n_w+w_start,:].mean(axis=(0,1)) + + if feather != -1: + if feather==0: + mask = (cv2.blur(mask, (n, n))) + else: + mask = (cv2.blur(mask, (feather, feather))) + mask = mask/255.0 + for i in range(3):img_mosaic[:,:,i] = (img[:,:,i]*(1-mask)+img_mosaic[:,:,i]*mask) + img_mosaic = img_mosaic.astype(np.uint8) + + return img_mosaic + +def get_autosize(img,mask,area_type = 'normal'): + h,w = img.shape[:2] + size = np.min([h,w]) + mask = resize(mask,size) + alpha = size/512 + try: + if area_type == 'normal': + area = mask_area(mask) + elif area_type == 'bounding': + w,h = cv2.boundingRect(mask)[2:] + area = w*h + except: + area = 0 + area = area/(alpha*alpha) + if area>50000: + size = alpha*((area-50000)/50000+12) + elif 20000 np.ndarray: - """Return a copy of ``image`` with the detected regions reconstructed.""" + def restore( + self, + image: np.ndarray, + detections: list[Detection], + should_cancel: CancelCheck | None = None, + ) -> np.ndarray: + """Return a copy of ``image`` with the detected regions reconstructed. + + ``should_cancel`` (if given) is polled periodically; when it returns + True the engine should abort and raise :class:`Cancelled`. + """ raise NotImplementedError diff --git a/hvideotool/core/restore/deepmosaics.py b/hvideotool/core/restore/deepmosaics.py index 1442b1d..850e36e 100644 --- a/hvideotool/core/restore/deepmosaics.py +++ b/hvideotool/core/restore/deepmosaics.py @@ -1,107 +1,161 @@ -"""DeepMosaics restorer — real generative mosaic removal. +"""DeepMosaics restorer — real generative mosaic removal, in-process. -Rather than vendoring DeepMosaics' GPL network code (which must match the exact -checkpoint), we drive a **user-installed** DeepMosaics (https://github.com/HypoX64/DeepMosaics) -as a subprocess: write the frame to a temp file, run ``deepmosaic.py --mode clean``, -read the cleaned image back. This reuses their tested pipeline (incl. their own -mosaic locator ``mosaic_position.pth``) and respects the GPL boundary. +The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see +its NOTICE/LICENSE). We load the models **once** and run the per-frame clean path +in-process — far faster than spawning a subprocess per frame (which reloaded the +models every time). Only the model *weights* are user-supplied. -Setup the user must do once (see README → Восстановление): -1. ``git clone https://github.com/HypoX64/DeepMosaics`` and install its deps. -2. Download clean weights (e.g. ``clean_youknow_video.pth``) AND ``mosaic_position.pth`` - into one folder. -3. In the app: Восстановление… → engine "deepmosaics", set the DeepMosaics folder - and the clean-model path (a CUDA GPU is strongly recommended). +Per-frame clean = DeepMosaics' ``cleanmosaic_img_server`` logic, reimplemented +here (so we don't pull in their video/ffmpeg modules): + locate mosaic (BiSeNet ``mosaic_position.pth``) → run the clean generator on the + crop → feather it back. DeepMosaics finds the mosaic itself; our detections are + used for navigation, not passed to it. -NOTE: DeepMosaics finds the mosaic itself; our detections are used for navigation, -not passed to it. +Setup (see README → Восстановление): download the **image** clean weights +``clean_youknow_resnet_9blocks.pth`` + ``mosaic_position.pth`` into one folder and +point the app at the clean-model file. The video model ``clean_youknow_video.pth`` +(BVDNet) needs neighbour frames and does NOT work per-frame. """ from __future__ import annotations -import subprocess import sys -import tempfile from pathlib import Path +from types import SimpleNamespace import numpy as np from ..detection.types import Detection -from ..imageio import imread_unicode, imwrite_unicode -from .base import Restorer +from .base import CancelCheck, Cancelled, Restorer -_IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp"} +_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). + + 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. + """ + dirs = [DEFAULT_WEIGHTS_DIR] + if extra_dir: + dirs.insert(0, Path(extra_dir)) + out: list[tuple[str, str]] = [] + seen: set[str] = set() + for d in dirs: + 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: + continue + seen.add(p.name) + out.append((p.stem, str(p))) + return out + + +def _netg_kind(model_name: str) -> str: + """Pick DeepMosaics' netG type from the weights filename (see their options.py).""" + n = model_name.lower() + if "video" in n: + raise ValueError( + "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен соседний " + "кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth." + ) + if "unet_128" in n: + return "unet_128" + if "hd" in n: + return "HD" + return "resnet_9blocks" class DeepMosaicsRestorer(Restorer): def __init__( self, - deepmosaics_dir: str | None, + deepmosaics_dir: str | None, # kept for factory/config compatibility (weights hint) model_path: str | None, - python_exe: str | None = None, + python_exe: str | None = None, # unused now (in-process) gpu_id: str = "0", ) -> None: - if not deepmosaics_dir or not (Path(deepmosaics_dir) / "deepmosaic.py").is_file(): - raise ValueError( - "Не указана папка DeepMosaics (с deepmosaic.py).\n" - "Установите DeepMosaics и укажите её в «Восстановление…». См. README." - ) if not model_path or not Path(model_path).is_file(): + discovered = discover_models() # fall back to a bundled model + if discovered: + model_path = discovered[0][1] + else: + raise ValueError( + "Не найдены веса DeepMosaics (clean_*.pth).\n" + "Положите clean_youknow_resnet_9blocks.pth + mosaic_position.pth в " + "models/deepmosaics (или выберите в «Восстановление…»). См. README." + ) + model = Path(model_path) + self._netg = _netg_kind(model.name) # raises on a video model + pos = self._find_mosaic_position(model, deepmosaics_dir) + if pos is None: raise ValueError( - "Не найдены веса DeepMosaics (clean_*.pth).\n" - "Скачайте clean_youknow_video.pth + mosaic_position.pth в одну папку. См. README." + "Рядом с clean-моделью не найден mosaic_position.pth.\n" + "Положите mosaic_position.pth в ту же папку, что и clean_*.pth. См. README." ) - self._dir = Path(deepmosaics_dir) - self._model = model_path - self._python = python_exe or sys.executable + self._model = str(model) + self._pos = str(pos) 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})" - def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray: - with tempfile.TemporaryDirectory(prefix="hvt_dm_") as tmp: - tmpd = Path(tmp) - src = tmpd / "frame.jpg" - result_dir = tmpd / "result" - result_dir.mkdir() - imwrite_unicode(str(src), image) + # ------------------------------------------------------------------ engine + def _ensure_loaded(self) -> None: + if self._loaded: + return + if str(_VENDOR) not in sys.path: + sys.path.insert(0, str(_VENDOR)) # so vendored `from models/util import …` resolve + from models import loadmodel, runmodel # type: ignore # noqa: E402 + import util.image_processing as impro # type: ignore # noqa: E402 - cmd = [ - self._python, "deepmosaic.py", - "--media_path", str(src), - "--model_path", str(self._model), - "--mode", "clean", - "--result_dir", str(result_dir), - "--temp_dir", str(tmpd / "dmtmp"), - "--gpu_id", str(self._gpu), - "--no_preview", - ] - proc = subprocess.run( - cmd, cwd=str(self._dir), - stdin=subprocess.DEVNULL, # so DeepMosaics' error input() can't hang - capture_output=True, text=True, - ) - outputs = [p for p in result_dir.iterdir() if p.suffix.lower() in _IMG_EXTS] - if outputs: - newest = max(outputs, key=lambda p: p.stat().st_mtime) - restored = imread_unicode(str(newest)) - if restored is None: - raise RuntimeError("Не удалось прочитать результат DeepMosaics.") - return restored + self._runmodel = runmodel + self._impro = impro + self._opt = SimpleNamespace( + gpu_id=self._gpu, + netG=self._netg, + model_path=self._model, + mosaic_position_model_path=self._pos, + mask_threshold=64, + all_mosaic_area=False, + ex_mult=1.5, + no_feather=False, + traditional=False, + ) + self._netM = loadmodel.bisenet(self._opt, "mosaic") + self._netG = loadmodel.pix2pix(self._opt) + self._loaded = True - # No output file — figure out why. - log = (proc.stderr or "") + (proc.stdout or "") - if "BVDNet.forward()" in log or "argument: 'previous'" in log: - raise RuntimeError( - "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен " - "соседний кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth." - ) - if proc.returncode == 0: - # DeepMosaics ran fine but found no mosaic to clean — keep the frame as is. - return image.copy() - tail = log.strip().splitlines()[-6:] - raise RuntimeError( - f"DeepMosaics не вернул результат (код {proc.returncode}).\n" + "\n".join(tail) - ) + def restore( + self, + image: np.ndarray, + detections: list[Detection], + should_cancel: CancelCheck | None = None, + ) -> np.ndarray: + if should_cancel is not None and should_cancel(): + raise Cancelled("Восстановление отменено") + self._ensure_loaded() + rm, impro, opt = self._runmodel, self._impro, self._opt + + # DeepMosaics' cleanmosaic_img_server, faithfully reproduced. + x, y, size, mask = rm.get_mosaic_position(image, self._netM, opt) + if size <= 100: + return image.copy() # no mosaic located — leave the frame untouched + if should_cancel is not None and should_cancel(): + raise Cancelled("Восстановление отменено") + work = image.copy() + 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) diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py index 464509e..07eba75 100644 --- a/hvideotool/core/restore/factory.py +++ b/hvideotool/core/restore/factory.py @@ -1,8 +1,9 @@ """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 via a user-installed DeepMosaics - (subprocess). Needs the DeepMosaics folder + clean weights + (ideally) a CUDA GPU. +- ``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. """ from __future__ import annotations diff --git a/hvideotool/core/restore/inpaint.py b/hvideotool/core/restore/inpaint.py index 25e992f..f5e2b16 100644 --- a/hvideotool/core/restore/inpaint.py +++ b/hvideotool/core/restore/inpaint.py @@ -13,7 +13,7 @@ import cv2 import numpy as np from ..detection.types import Detection -from .base import Restorer +from .base import CancelCheck, Restorer from .mask import detections_to_mask @@ -27,7 +27,13 @@ class InpaintRestorer(Restorer): def name(self) -> str: return f"InpaintRestorer({self.method})" - def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray: + 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) diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py index 0274bed..f533f27 100644 --- a/hvideotool/settings_store.py +++ b/hvideotool/settings_store.py @@ -1,7 +1,9 @@ """Persist a handful of user choices to ~/HVideoTool/settings.json. -Slimmed down for the image-inspector tool: it remembers the detector, the model -path, the overlay threshold, and the last opened folder. +For the project-based tool this file holds the **defaults for new projects** (the +detector, model path, overlay threshold, restore engine) plus app-level state: the +last opened project, the recent-projects list, and the last directory used in file +dialogs. Per-project settings live in each project's ``project.json``, not here. """ from __future__ import annotations @@ -66,3 +68,30 @@ def set_last_dir(path: str) -> None: data = _read() data["last_dir"] = path _write(data) + + +def last_project() -> str | None: + return _read().get("last_project") + + +def set_last_project(path: str) -> None: + data = _read() + data["last_project"] = path + _write(data) + + +_RECENTS_CAP = 10 + + +def recent_projects() -> list[str]: + recents = _read().get("recent_projects", []) + return [p for p in recents if isinstance(p, str)] + + +def add_recent_project(path: str) -> None: + """Push ``path`` to the front of the recent-projects list (deduped, capped).""" + data = _read() + recents = [p for p in data.get("recent_projects", []) if isinstance(p, str) and p != path] + recents.insert(0, path) + data["recent_projects"] = recents[:_RECENTS_CAP] + _write(data) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index e0ed7ca..b102f4a 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -1,18 +1,23 @@ -"""Main window: open a folder of images and inspect what the detector found. +"""Main window: open a **project** and inspect what the detector found. -Layout: a toolbar (open folder · from-video · detector · model · calc-frame · +A project is a folder (``project.json`` + ``frames/`` + ``detections.json`` + +``collections/``) — see ``core/project.py``. The per-project settings (detector, +model, threshold, restore engine) live in ``project.json``; the global +``settings.json`` only seeds defaults for new projects. + +Layout: a toolbar (new/open project · from-video · detector · model · calc-frame · detect-all · threshold), then a splitter with three panes — left: collection controls + the file list; center: the image with overlays; right: a detail table of every detection. Collection controls sit by the file list (they act on its selection), keeping the toolbar to detection/entry actions only. -Viewing and detecting are decoupled, so browsing a big folder stays instant even +Viewing and detecting are decoupled, so browsing a big project stays instant even with a slow (CPU) detector: - selecting a file just **shows** it (with its cached result, if any); - **double-clicking** a file, or "Рассчитать кадр", runs the detector on it; -- "Детектировать все" runs the whole folder. -Both folder loading and detect-all show a progress bar. Results are cached; -switching detector/model clears the cache. +- "Детектировать все" runs the whole project. +Both project loading and detect-all show a progress bar. Results are cached in the +project; switching detector/model clears the cache. """ from __future__ import annotations @@ -47,9 +52,12 @@ from PySide6.QtWidgets import ( from .. import settings_store from ..config import AppConfig +from ..core.detection import cache as detection_cache from ..core.detection.factory import build_detector from ..core.detection.types import Detection from ..core.imageio import imread_unicode, imwrite_unicode +from ..core.project import PROJECT_FILE, Project +from ..core.restore.base import Cancelled from ..core.restore.factory import build_restorer from ..core.video.extract import extract_frames from ..core.video.frame import Frame @@ -69,16 +77,18 @@ class MainWindow(QMainWindow): self._cfg = config self._detector = None self._detector_key = None - self._folder: Path | None = None + self._project: Project | None = None # the open project (None until one is opened) + self._folder: Path | None = None # == project.frames_dir while a project is open self._files: list[Path] = [] self._results: dict[str, list[Detection]] = {} # path -> detections (cache) self._current: Path | None = None - self._collection: Path | None = None # active destination folder for moves self._restorer = None # un-censor engine, built lazily from config self._restorer_key = None self._restored: dict[str, "object"] = {} # path -> restored image (BGR ndarray) self._showing_restored = False self._nav_sync = False # guard against slider<->list signal loops + self._busy = False # a long operation is running + self._cancel = False # the user asked to stop it self.setWindowTitle("HVideoTool — инспектор детекции цензуры") self.resize(1180, 720) @@ -87,22 +97,25 @@ class MainWindow(QMainWindow): self._build_central() self._build_statusbar() self._build_menu() - self._refresh_collections() - self.statusBar().showMessage("Откройте папку с картинками") + self.statusBar().showMessage("Создайте или откройте проект (Файл)") # ------------------------------------------------------------------ setup def _build_menu(self) -> None: file_menu = self.menuBar().addMenu("Файл") - file_menu.addAction("Открыть папку…", self._choose_folder) + file_menu.addAction("Создать проект…", self._create_project) + file_menu.addAction("Открыть проект…", self._open_project_dialog) + file_menu.addAction("Импортировать папку как проект…", self._import_folder_as_project) file_menu.addAction("Создать из ролика…", self._create_from_video) + self._recent_menu = file_menu.addMenu("Недавние проекты") + self._refresh_recent_menu() file_menu.addSeparator() file_menu.addAction("Рассчитать кадр", self._recompute_current).setShortcut("Space") - file_menu.addAction("Детектировать все", self._detect_all) + file_menu.addAction("Детектировать все (дозапуск)", lambda: self._detect_all(False)) + file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True)) file_menu.addSeparator() file_menu.addAction("Движок восстановления…", self._open_restore_settings) file_menu.addSeparator() - file_menu.addAction("Создать коллекцию…", self._create_collection) - file_menu.addAction("В коллекцию", self._move_to_collection).setShortcut("Ctrl+M") + file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M") file_menu.addSeparator() file_menu.addAction("Выход", self.close) @@ -110,9 +123,10 @@ class MainWindow(QMainWindow): tb = self.addToolBar("Главная") tb.setMovable(False) - tb.addAction(QAction("Открыть папку…", self, triggered=self._choose_folder)) + tb.addAction(QAction("Создать проект…", self, triggered=self._create_project)) + tb.addAction(QAction("Открыть проект…", self, triggered=self._open_project_dialog)) from_video = QAction("Создать из ролика…", self, triggered=self._create_from_video) - from_video.setToolTip("Разложить видео на кадры в папку-коллекцию и открыть её") + from_video.setToolTip("Разложить видео на кадры в новый проект и открыть его") tb.addAction(from_video) tb.addSeparator() @@ -130,7 +144,17 @@ class MainWindow(QMainWindow): calc = QAction("Рассчитать кадр", self, triggered=self._recompute_current) calc.setToolTip("Запустить детектор на выбранном кадре (Space / двойной клик по файлу)") tb.addAction(calc) - tb.addAction(QAction("Детектировать все", self, triggered=self._detect_all)) + detect_all = QAction("Детектировать все", self, triggered=lambda: self._detect_all(False)) + detect_all.setToolTip("Рассчитать все ещё не посчитанные кадры (дозапуск; кэш сохраняется)") + tb.addAction(detect_all) + regen = QAction("Все заново", self, triggered=lambda: self._detect_all(True)) + regen.setToolTip("Очистить кэш детекций и пересчитать всю папку заново") + tb.addAction(regen) + + self.stop_action = QAction("■ Стоп", self, triggered=self._request_cancel) + self.stop_action.setToolTip("Отменить текущую операцию (Esc)") + self.stop_action.setEnabled(False) + tb.addAction(self.stop_action) tb.addSeparator() restore = QAction("Расцензурить кадр", self, triggered=self._restore_current) @@ -158,27 +182,15 @@ class MainWindow(QMainWindow): self.file_list.currentItemChanged.connect(self._on_file_selected) self.file_list.itemDoubleClicked.connect(self._on_file_activated) - # Collection controls live next to the file list — they act on its selection. - self.collection_combo = QComboBox() - self.collection_combo.setToolTip("Активная коллекция, куда перемещаются кадры") - self.collection_combo.activated.connect(self._on_collection_selected) - new_coll = QPushButton("Создать") - new_coll.clicked.connect(self._create_collection) - move_btn = QPushButton("В коллекцию →") - move_btn.setToolTip("Переместить выбранные кадры в активную коллекцию (Ctrl+M)") - move_btn.clicked.connect(self._move_to_collection) - - coll_row = QHBoxLayout() - coll_row.setContentsMargins(0, 0, 0, 0) - coll_row.addWidget(QLabel("Коллекция:")) - coll_row.addWidget(self.collection_combo, 1) - coll_row.addWidget(new_coll) + # One default collection ("Избранное"); the button acts on the list selection. + move_btn = QPushButton("★ В избранное") + move_btn.setToolTip("Переместить выбранные кадры в избранное проекта (Ctrl+M)") + move_btn.clicked.connect(self._move_to_favorites) left = QWidget() left_layout = QVBoxLayout(left) left_layout.setContentsMargins(4, 4, 4, 4) left_layout.setSpacing(4) - left_layout.addLayout(coll_row) left_layout.addWidget(self.file_list, 1) left_layout.addWidget(move_btn) @@ -253,6 +265,7 @@ class MainWindow(QMainWindow): QShortcut(QKeySequence("."), self, lambda: self._step(1)) QShortcut(QKeySequence("["), self, lambda: self._step_hit(-1)) QShortcut(QKeySequence("]"), self, lambda: self._step_hit(1)) + QShortcut(QKeySequence(Qt.Key_Escape), self, self._request_cancel) return bar # -------------------------------------------------------------- navigation @@ -302,6 +315,35 @@ class MainWindow(QMainWindow): self.progress.setVisible(False) self.statusBar().addPermanentWidget(self.progress) + # ------------------------------------------------------------- cancellation + def _begin_busy(self, total: int | None = None) -> None: + """Enter a cancellable long operation. ``total=None`` => busy spinner.""" + self._busy = True + self._cancel = False + self.stop_action.setEnabled(True) + if total is None: + self.progress.setRange(0, 0) # indeterminate + else: + self.progress.setRange(0, total) + self.progress.setValue(0) + self.progress.setVisible(True) + + def _end_busy(self) -> None: + self._busy = False + self.stop_action.setEnabled(False) + self.progress.setVisible(False) + self.progress.setRange(0, 100) # leave it determinate for the next user + + def _request_cancel(self) -> None: + if self._busy: + self._cancel = True + self.statusBar().showMessage("Отмена…") + + def _poll_cancel(self) -> bool: + """Cancel hook for core engines: pump the UI so Стоп registers, then report.""" + QApplication.processEvents() + return self._cancel + # --------------------------------------------------------------- detector def _make_detector(self): d = self._cfg.detection @@ -316,7 +358,7 @@ class MainWindow(QMainWindow): # YOLO/combined need a model — offer to pick one if missing. if name in ("yolo", "combined") and not self._cfg.model_path: self._choose_model() - settings_store.save(self._cfg) + self._persist_settings() self._invalidate_results() def _choose_model(self) -> None: @@ -324,34 +366,210 @@ class MainWindow(QMainWindow): path, _ = QFileDialog.getOpenFileName(self, "Выберите веса (.pt)", start, "Веса YOLO (*.pt);;Все файлы (*.*)") if path: self._cfg.model_path = path - settings_store.save(self._cfg) + self._persist_settings() self.statusBar().showMessage(f"Модель: {path}") self._invalidate_results() def _invalidate_results(self) -> None: - """Detector changed — drop the cache and refresh the current image.""" + """Detector changed — drop the in-memory cache and refresh the current image. + + The on-disk cache is left as-is; it won't be reloaded for the new detector + (key mismatch) and gets overwritten once results for the new detector exist. + """ self._detector_key = None - self._results.clear() - for i in range(self.file_list.count()): - item = self.file_list.item(i) - item.setText(item.data(Qt.UserRole + 1)) - item.setBackground(QBrush()) - self._refresh_marks() + self._clear_results() if self._current is not None: self._show(self._current) - # --------------------------------------------------------------- handlers - def open_path(self, folder: str) -> None: - self._load_folder(Path(folder)) + # ----------------------------------------------------------------- projects + def open_path(self, path: str) -> None: + """Open a project at ``path`` (a project folder or its project.json).""" + p = Path(path) + if Project.is_project(p): + try: + self._open_project(Project.load(p)) + except (OSError, ValueError) as exc: # noqa: BLE001 - surface to the user + QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}") + elif p.is_dir(): + QMessageBox.information( + self, "Не проект", + "Это обычная папка, а не проект. Используйте " + "«Импортировать папку как проект…».", + ) + else: + QMessageBox.warning(self, "Ошибка", f"Путь не найден: {p}") - def _choose_folder(self) -> None: + def _auto_open_last(self) -> None: + """On startup, reopen the last project if it still exists (best-effort).""" + last = settings_store.last_project() + if last and Project.is_project(last): + try: + self._open_project(Project.load(last)) + except (OSError, ValueError): + pass + + def _new_project_root(self, default_name: str = "") -> Path | None: + """Prompt for a parent dir + name; return a fresh (empty) project root or None.""" + start = settings_store.last_dir() or str(Path.home()) + parent = QFileDialog.getExistingDirectory(self, "Где создать проект", start) + if not parent: + return None + name, ok = QInputDialog.getText(self, "Новый проект", "Имя проекта:", text=default_name) + name = name.strip() + if not ok or not name: + return None + root = Path(parent) / name + if root.exists() and any(root.iterdir()): + QMessageBox.warning(self, "Папка занята", f"Папка уже существует и не пуста:\n{root}") + return None + settings_store.set_last_dir(parent) + return root + + def _create_project(self) -> None: + if self._busy: + return + root = self._new_project_root() + if root is None: + return + try: + project = Project.create(root, name=root.name) + except OSError as exc: + QMessageBox.warning(self, "Ошибка", f"Не удалось создать проект:\n{exc}") + return + project.update_from_config(self._cfg) # seed from current global defaults + project.save() + self._open_project(project) + + def _open_project_dialog(self) -> None: + if self._busy: + return + start = settings_store.last_dir() or str(Path.home()) + folder = QFileDialog.getExistingDirectory(self, "Открыть проект (папка проекта)", start) + if not folder: + return + if not Project.is_project(folder): + QMessageBox.warning(self, "Не проект", f"В папке нет {PROJECT_FILE}:\n{folder}") + return + try: + project = Project.load(folder) + except (OSError, ValueError) as exc: + QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}") + return + settings_store.set_last_dir(str(Path(folder).parent)) + self._open_project(project) + + def _import_folder_as_project(self) -> None: + """Create a project and copy a folder of images into its frames/.""" + if self._busy: + return start = settings_store.last_dir() or "" - folder = QFileDialog.getExistingDirectory(self, "Открыть папку с картинками", start) - if folder: - self._load_folder(Path(folder)) + src = QFileDialog.getExistingDirectory(self, "Папка с картинками для импорта", start) + if not src: + return + src = Path(src) + images = sorted(p for p in src.iterdir() if p.suffix.lower() in _IMAGE_EXTS) + if not images: + QMessageBox.warning(self, "Пусто", f"В папке нет картинок:\n{src}") + return + root = self._new_project_root(default_name=src.name) + if root is None: + return + try: + project = Project.create(root, name=root.name, source=str(src)) + except OSError as exc: + QMessageBox.warning(self, "Ошибка", f"Не удалось создать проект:\n{exc}") + return + project.update_from_config(self._cfg) + project.save() + + self._begin_busy(len(images)) + copied = 0 + try: + for i, p in enumerate(images, 1): + self.progress.setValue(i) + self.statusBar().showMessage(f"Импорт {i}/{len(images)}: {p.name}") + QApplication.processEvents() + if self._cancel: + break + dst = self._unique_dest(project.frames_dir, p.name) + try: + shutil.copy2(str(p), str(dst)) + copied += 1 + except OSError: + continue + finally: + self._end_busy() + + # Carry over an old sidecar detection cache (basename-keyed) if present. + old_sidecar = src / ".hvideotool_detections.json" + if old_sidecar.is_file(): + try: + shutil.copy2(str(old_sidecar), str(project.cache_path)) + except OSError: + pass + + self.statusBar().showMessage(f"Импортировано {copied} картинок → {project.name}") + self._open_project(project) + + def _refresh_recent_menu(self) -> None: + self._recent_menu.clear() + recents = settings_store.recent_projects() + if not recents: + empty = self._recent_menu.addAction("(пусто)") + empty.setEnabled(False) + return + for path in recents: + self._recent_menu.addAction(Path(path).name, lambda checked=False, p=path: self._open_recent(p)) + + def _open_recent(self, path: str) -> None: + if self._busy: + return + if not Project.is_project(path): + QMessageBox.warning(self, "Нет проекта", f"Проект не найден:\n{path}") + return + try: + self._open_project(Project.load(path)) + except (OSError, ValueError) as exc: + QMessageBox.warning(self, "Ошибка", f"Не удалось открыть проект:\n{exc}") + + def _open_project(self, project: Project) -> None: + """Core open: set project state, apply its settings, list its frames.""" + if self._busy: + return + self._project = project + project.frames_dir.mkdir(parents=True, exist_ok=True) + project.apply_to_config(self._cfg) # per-project settings -> live config + self._sync_settings_ui() + self._detector_key = None + self._restorer_key = None + self._restored.clear() + settings_store.set_last_project(str(project.root)) + settings_store.add_recent_project(str(project.root)) + self._refresh_recent_menu() + self.setWindowTitle(f"HVideoTool — {project.name}") + self._load_folder(project.frames_dir) + + 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) + self.view.set_threshold(self._cfg.default_threshold) + + def _persist_settings(self) -> None: + """Save settings to the global defaults and (if open) into the project.""" + settings_store.save(self._cfg) # global defaults for new projects + if self._project is not None: + self._project.update_from_config(self._cfg) + self._project.save() def _create_from_video(self) -> None: - """Decode a video into a folder of frames (a collection) and open it.""" + """Decode a video into a new project's frames/ and open the project.""" + if self._busy: + return path, _ = QFileDialog.getOpenFileName( self, "Выберите ролик", settings_store.last_dir() or "", _VIDEO_FILTER ) @@ -362,18 +580,26 @@ class MainWindow(QMainWindow): return keyframes_only, step, max_dim = dialog.options() video = Path(path) - out = video.parent / f"{video.stem}_frames" + root = video.parent / f"{video.stem}_frames" + if Project.is_project(root): + project = Project.load(root) # re-extract into the existing project + elif root.exists() and any(root.iterdir()): + QMessageBox.warning(self, "Папка занята", f"Папка уже существует и не пуста:\n{root}") + return + else: + project = Project.create(root, name=root.name, source=str(video)) + project.update_from_config(self._cfg) + project.save() + out = project.frames_dir - self.progress.setRange(0, 1000) # promille of duration - self.progress.setValue(0) - self.progress.setVisible(True) + self._begin_busy(1000) # promille of duration def cb(done: float, total: float) -> bool: if total > 0: self.progress.setValue(int(1000 * min(done, total) / total)) self.statusBar().showMessage(f"Извлечение кадров: {done:.0f}/{total:.0f} с…") QApplication.processEvents() - return True + return not self._cancel # returning False stops extraction try: saved = extract_frames( @@ -381,19 +607,23 @@ class MainWindow(QMainWindow): max_dim=max_dim, progress=cb, ) except Exception as exc: # noqa: BLE001 - surface decode errors to the user - self.progress.setVisible(False) QMessageBox.warning(self, "Ошибка", f"Не удалось извлечь кадры:\n{exc}") return finally: - self.progress.setVisible(False) + cancelled = self._cancel + self._end_busy() if saved == 0: - QMessageBox.warning(self, "Пусто", "Из ролика не удалось извлечь ни одного кадра.") + msg = "Извлечение отменено — кадров нет." if cancelled \ + else "Из ролика не удалось извлечь ни одного кадра." + QMessageBox.warning(self, "Пусто", msg) return - self.statusBar().showMessage(f"Извлечено {saved} кадров → {out}") - self._load_folder(out) + verb = "Отменено, извлечено" if cancelled else "Извлечено" + self.statusBar().showMessage(f"{verb} {saved} кадров → {out}") + self._open_project(project) def _load_folder(self, folder: Path) -> None: + """List images from ``folder`` (a project's frames/) into the file list.""" if not folder.is_dir(): QMessageBox.warning(self, "Ошибка", f"Папка не найдена: {folder}") return @@ -404,13 +634,11 @@ class MainWindow(QMainWindow): self._files = files self._results.clear() self._current = None - settings_store.set_last_dir(str(folder)) self.file_list.blockSignals(True) self.file_list.setUpdatesEnabled(False) self.file_list.clear() - self.progress.setRange(0, len(files)) - self.progress.setVisible(True) + self._begin_busy(len(files)) for i, p in enumerate(files, 1): item = QListWidgetItem(p.name) item.setData(Qt.UserRole, str(p)) @@ -420,18 +648,27 @@ class MainWindow(QMainWindow): self.progress.setValue(i) self.statusBar().showMessage(f"Загрузка списка: {i}/{len(files)}…") QApplication.processEvents() + if self._cancel: + self._files = files[:i] # keep only what we listed + break self.file_list.setUpdatesEnabled(True) self.file_list.blockSignals(False) - self.progress.setVisible(False) + self._end_busy() + loaded = self._load_cached_results() # reuse a matching on-disk cache self._refresh_marks() - self._refresh_collections() if not files: self.view.set_image(None, []) self._update_nav() - self.statusBar().showMessage(f"В папке нет картинок: {folder}") + self.statusBar().showMessage( + "В проекте пока нет кадров — импортируйте папку или создайте из ролика" + ) return - self.statusBar().showMessage(f"{len(files)} картинок · {folder} · двойной клик / «Рассчитать кадр» для детекции") + cache_note = f" · загружен кэш детекций ({loaded})" if loaded else "" + self.statusBar().showMessage( + f"{len(files)} картинок · {folder} · двойной клик / «Рассчитать кадр» для детекции" + + cache_note + ) self.file_list.setCurrentRow(0) def _on_file_selected(self, current: QListWidgetItem | None, _prev=None) -> None: @@ -441,10 +678,14 @@ class MainWindow(QMainWindow): def _on_file_activated(self, item: QListWidgetItem) -> None: # Double-click: compute if not already cached, then show. - path = Path(item.data(Qt.UserRole)) - if str(path) not in self._results and self._detect(path) is None: + if self._busy: return - self._refresh_marks() + path = Path(item.data(Qt.UserRole)) + if str(path) not in self._results: + if self._detect(path) is None: + return + self._refresh_marks() + self._save_results() # only when a detection actually ran self._show(path) def _detect(self, path: Path) -> list[Detection] | None: @@ -481,40 +722,56 @@ class MainWindow(QMainWindow): def _recompute_current(self) -> None: """Toolbar/Space: (re)run the detector on the selected frame.""" - if self._current is None: + if self._current is None or self._busy: return self._results.pop(str(self._current), None) self._detector_key = None # rebuild the detector so settings changes take effect if self._detect(self._current) is None: return self._refresh_marks() + self._save_results() self._show(self._current) - def _detect_all(self) -> None: - if not self._files: + def _detect_all(self, force: bool = False) -> None: + """Detect the whole folder. ``force`` clears the cache first (full regen); + otherwise already-computed frames are skipped, so it resumes/tops-up.""" + if not self._files or self._busy: return + if force: + self._clear_results() total = len(self._files) - self.progress.setRange(0, total) - self.progress.setVisible(True) + self._begin_busy(total) + done = 0 try: for i, p in enumerate(self._files, 1): self.progress.setValue(i) self.statusBar().showMessage(f"Детекция {i}/{total}: {p.name}") QApplication.processEvents() + if self._cancel: + break if self._detect(p) is None: return # detector unavailable — message already shown + done = i + if i % 50 == 0: + self._refresh_marks() # let marks appear progressively finally: - self.progress.setVisible(False) + self._end_busy() hits = sum(1 for p in self._files if self._results.get(str(p))) self._refresh_marks() - self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок") + self._save_results() # persist progress (works for completed and cancelled runs) + if self._cancel: + self.statusBar().showMessage( + f"Отменено на {done}/{total} · детекции на {hits} картинках" + ) + else: + self.statusBar().showMessage(f"Готово: детекции на {hits} из {total} картинок") if self._current is not None: self._show(self._current) # ------------------------------------------------------------- restoration def _restore_current(self) -> None: """Run the restorer on the current frame's detected regions and show it.""" - if self._current is None: + if self._current is None or self._busy: return key = str(self._current) if key not in self._results and self._detect(self._current) is None: @@ -527,14 +784,20 @@ class MainWindow(QMainWindow): img = imread_unicode(key) if img is None: return + self._begin_busy() # indeterminate — engine drives the duration self.statusBar().showMessage(f"Восстановление: {self._current.name}…") QApplication.processEvents() try: restorer = self._make_restorer() - restored = restorer.restore(img, dets) + restored = restorer.restore(img, dets, should_cancel=self._poll_cancel) + except Cancelled: + self.statusBar().showMessage("Восстановление отменено") + return except Exception as exc: # noqa: BLE001 - surface model/engine errors QMessageBox.warning(self, "Ошибка восстановления", str(exc)) return + finally: + self._end_busy() self._restored[key] = restored self._showing_restored = True self.view.set_image(restored, []) @@ -556,7 +819,7 @@ class MainWindow(QMainWindow): if dlg.exec() != QDialog.Accepted: return dlg.apply_to_config() - settings_store.save(self._cfg) + self._persist_settings() self._restorer_key = None # rebuild on next restore self.statusBar().showMessage(f"Движок восстановления: {self._cfg.restorer}") @@ -582,100 +845,37 @@ class MainWindow(QMainWindow): def _save_restored(self) -> None: if self._current is None or str(self._current) not in self._restored: return - dest_dir = self._collection or self._current.parent - out = self._unique_dest(dest_dir, f"{self._current.stem}_restored.jpg") + out = self._unique_dest(self._current.parent, f"{self._current.stem}_restored.jpg") if imwrite_unicode(str(out), self._restored[str(self._current)]): self.statusBar().showMessage(f"Сохранено: {out}") else: QMessageBox.warning(self, "Ошибка", "Не удалось сохранить файл.") - # ------------------------------------------------------------ collections - def _collections_base(self) -> Path: - """Where new collections are created: next to the opened folder, else home.""" - if self._folder is not None: - return self._folder.parent - return Path.home() / "HVideoTool" / "collections" - - def _refresh_collections(self) -> None: - """Repopulate the collection combo from sibling folders of the opened folder. - - Keeps the active collection selected (and present even if it lives elsewhere). - """ - base = self._collections_base() - subdirs = [] - if base.exists(): - subdirs = sorted( - (p for p in base.iterdir() if p.is_dir() and p != self._folder), - key=lambda p: p.name.lower(), - ) - self.collection_combo.blockSignals(True) - self.collection_combo.clear() - self.collection_combo.addItem("— не выбрана —", None) - for p in subdirs: - self.collection_combo.addItem(p.name, str(p)) - # Make sure the active collection is listed even if it's outside base. - if self._collection is not None and self.collection_combo.findData(str(self._collection)) < 0: - self.collection_combo.addItem(self._collection.name, str(self._collection)) - self.collection_combo.addItem("Выбрать папку…", "__browse__") - self._select_active_in_combo() - self.collection_combo.blockSignals(False) - - def _select_active_in_combo(self) -> None: - idx = self.collection_combo.findData(str(self._collection)) if self._collection else 0 - self.collection_combo.setCurrentIndex(max(0, idx)) - - def _on_collection_selected(self, _index: int) -> None: - data = self.collection_combo.currentData() - if data == "__browse__": - self._browse_collection() + # -------------------------------------------------------------- favorites + def _move_to_favorites(self) -> None: + """Move the selected frames into the project's single default collection.""" + if self._busy: return - self._collection = Path(data) if data else None - if self._collection is not None: - self.statusBar().showMessage(f"Активная коллекция: {self._collection}") - - def _browse_collection(self) -> None: - start = str(self._collections_base()) - folder = QFileDialog.getExistingDirectory(self, "Выбрать коллекцию", start) - if folder: - self._collection = Path(folder) - self._refresh_collections() - self.statusBar().showMessage(f"Активная коллекция: {folder}") - else: - self._select_active_in_combo() # revert the combo to the current collection - - def _create_collection(self) -> None: - name, ok = QInputDialog.getText(self, "Создать коллекцию", "Имя коллекции:") - name = name.strip() - if not ok or not name: - return - path = self._collections_base() / name - try: - path.mkdir(parents=True, exist_ok=True) - except OSError as exc: - QMessageBox.warning(self, "Ошибка", f"Не удалось создать коллекцию:\n{exc}") - return - self._collection = path - self._refresh_collections() - self.statusBar().showMessage(f"Активная коллекция: {path}") - - def _move_to_collection(self) -> None: - if self._collection is None: - QMessageBox.information( - self, "Нет коллекции", - "Сначала выберите коллекцию в списке или создайте новую («Создать…»).", - ) + if self._project is None: + QMessageBox.information(self, "Нет проекта", "Сначала откройте или создайте проект.") return items = self.file_list.selectedItems() if not items: QMessageBox.information(self, "Нет выбора", "Выберите кадры в списке слева.") return + dest = self._project.favorites_dir + try: + dest.mkdir(parents=True, exist_ok=True) + except OSError as exc: + QMessageBox.warning(self, "Ошибка", f"Не удалось создать избранное:\n{exc}") + return moved = 0 for item in items: src = Path(item.data(Qt.UserRole)) if not src.exists(): continue - dst = self._unique_dest(self._collection, src.name) + dst = self._unique_dest(dest, src.name) try: shutil.move(str(src), str(dst)) except OSError as exc: @@ -688,7 +888,9 @@ class MainWindow(QMainWindow): if self._current == src: self._current = None - self.statusBar().showMessage(f"Перемещено {moved} → {self._collection.name}") + if moved: + self._save_results() # cache file should forget the moved frames + self.statusBar().showMessage(f"В избранное перемещено {moved}") cur = self.file_list.currentItem() if cur is not None: self._show(Path(cur.data(Qt.UserRole))) @@ -714,15 +916,60 @@ class MainWindow(QMainWindow): _TINT_HIT = QColor(200, 80, 80, 70) _TINT_CLEAN = QColor(90, 160, 90, 50) + def _set_row_tag(self, item: QListWidgetItem, count: int) -> None: + base = item.data(Qt.UserRole + 1) + item.setText(f"{base} · {count}" if count else f"{base} · —") + item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN) + def _tag_file(self, path: Path, count: int) -> None: for i in range(self.file_list.count()): item = self.file_list.item(i) if item.data(Qt.UserRole) == str(path): - base = item.data(Qt.UserRole + 1) - item.setText(f"{base} · {count}" if count else f"{base} · —") - item.setBackground(self._TINT_HIT if count else self._TINT_CLEAN) + self._set_row_tag(item, count) return + # ------------------------------------------------------------- result cache + def _results_key(self) -> dict: + """Detector identity used to tag/validate the on-disk detection cache.""" + d = self._cfg.detection + return detection_cache.make_key( + self._cfg.detector, self._cfg.model_path, d.yolo_conf, d.yolo_imgsz + ) + + def _save_results(self) -> None: + """Persist the detection cache in the project (skip if nothing to save).""" + if self._project is None or not self._results: + return + detection_cache.save_results( + self._project.cache_path, self._results_key(), self._results + ) + + def _load_cached_results(self) -> int: + """Load a matching on-disk cache into `_results` and tag rows. Returns count.""" + if self._project is None: + return 0 + cached = detection_cache.load_results( + self._project.cache_path, self._results_key(), self._project.frames_dir + ) + if not cached: + return 0 + self._results = cached + for i in range(self.file_list.count()): + item = self.file_list.item(i) + path = item.data(Qt.UserRole) + if path in self._results: # `in`, not truthy: empty list = checked-clean + self._set_row_tag(item, len(self._results[path])) + return len(cached) + + def _clear_results(self) -> None: + """Drop all cached detections and reset row labels/tints (keeps the detector).""" + self._results.clear() + for i in range(self.file_list.count()): + item = self.file_list.item(i) + item.setText(item.data(Qt.UserRole + 1)) + item.setBackground(QBrush()) + self._refresh_marks() + def _refresh_marks(self) -> None: """Project frames-with-detections onto the scrubber as marks.""" marks = { @@ -767,4 +1014,11 @@ class MainWindow(QMainWindow): def _on_threshold_changed(self, value: float) -> None: self._cfg.default_threshold = value self.view.set_threshold(value) - settings_store.save(self._cfg) + self._persist_settings() + + def closeEvent(self, event) -> None: # noqa: N802 - Qt override + self._save_results() # persist the detection cache on exit + if self._project is not None: + self._project.update_from_config(self._cfg) + self._project.save() + super().closeEvent(event) diff --git a/hvideotool/ui/marker_slider.py b/hvideotool/ui/marker_slider.py index 8690af7..61b1409 100644 --- a/hvideotool/ui/marker_slider.py +++ b/hvideotool/ui/marker_slider.py @@ -18,7 +18,8 @@ class MarkerSlider(QSlider): def __init__(self, orientation=Qt.Horizontal, parent=None) -> None: super().__init__(orientation, parent) self._marks: set[int] = set() - self._mark_color = QColor(220, 70, 70) + # Bright cyan stands out against both the dark track and the orange fill. + self._mark_color = QColor(0, 220, 255) def set_marks(self, marks: Iterable[int]) -> None: marks = set(marks) @@ -49,11 +50,15 @@ class MarkerSlider(QSlider): return lo, hi = self.minimum(), self.maximum() half = handle.width() // 2 - top = groove.center().y() - 5 - bottom = groove.center().y() + 5 + # Span (almost) the full widget height so marks read over the fill/handle. + top = self.rect().top() + 1 + bottom = self.rect().bottom() - 1 painter = QPainter(self) - painter.setPen(self._mark_color) + pen = painter.pen() + pen.setColor(self._mark_color) + pen.setWidth(2) + painter.setPen(pen) # Many frames can collapse onto the same pixel column — dedupe to keep # repaint cheap on big folders (tens of thousands of frames). seen_x: set[int] = set() diff --git a/hvideotool/ui/restore_dialog.py b/hvideotool/ui/restore_dialog.py index a1a8870..c046ab8 100644 --- a/hvideotool/ui/restore_dialog.py +++ b/hvideotool/ui/restore_dialog.py @@ -1,7 +1,9 @@ """Configure the restoration ("расцензурить") engine. -inpaint — no setup. deepmosaics — point at a user-installed DeepMosaics folder and -its clean weights; a CUDA GPU is strongly recommended (set GPU id, -1 = CPU/slow). +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). """ from __future__ import annotations @@ -22,6 +24,7 @@ from PySide6.QtWidgets import ( ) from ..config import AppConfig +from ..core.restore.deepmosaics import discover_models class RestoreDialog(QDialog): @@ -29,7 +32,7 @@ class RestoreDialog(QDialog): super().__init__(parent) self._cfg = config self.setWindowTitle("Движок восстановления") - self.setMinimumWidth(520) + self.setMinimumWidth(560) self.engine = QComboBox() self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint") @@ -37,24 +40,21 @@ class RestoreDialog(QDialog): self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0) self.engine.currentIndexChanged.connect(self._sync) - self.dm_dir = QLineEdit(config.dm_dir or "") - self.dm_model = QLineEdit(config.dm_model or "") - self.dm_python = QLineEdit(config.dm_python or "") - self.dm_python.setPlaceholderText("по умолчанию — python текущего venv") + # Model dropdown — bundled clean models, plus the configured one if external. + self.model_combo = QComboBox() + self._populate_models(config.dm_model) + self.dm_gpu = QLineEdit(config.dm_gpu or "0") self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)") form = QFormLayout(self) form.addRow("Движок:", self.engine) - form.addRow("Папка DeepMosaics:", self._with_browse(self.dm_dir, self._browse_dir)) - form.addRow("Веса (clean_*.pth):", self._with_browse(self.dm_model, self._browse_model)) - form.addRow("Python для DeepMosaics:", self._with_browse(self.dm_python, self._browse_python)) + form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model)) form.addRow("GPU id:", self.dm_gpu) hint = QLabel( - "DeepMosaics ставится отдельно (git clone + зависимости). Рядом с весами " - "clean_*.pth должен лежать mosaic_position.pth.\n" - "Для покадрового режима берите clean_youknow_resnet_9blocks.pth — " - "видеомодель clean_youknow_video.pth покадрово не работает. См. README." + "Модели берутся из models/deepmosaics. Нужны clean_youknow_resnet_9blocks.pth " + "и mosaic_position.pth (рядом). Видеомодель clean_*_video.pth покадрово не " + "работает и в списке не показывается. На CPU медленно — лучше GPU. См. README." ) hint.setWordWrap(True) form.addRow(hint) @@ -65,41 +65,41 @@ class RestoreDialog(QDialog): form.addRow(buttons) self._sync() - def _with_browse(self, line: QLineEdit, slot) -> QWidget: + def _populate_models(self, current: str | None) -> None: + self.model_combo.clear() + for name, path in discover_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: + self.model_combo.addItem(Path(current).stem + " (внешняя)", current) + if self.model_combo.count() == 0: + self.model_combo.addItem("(модели не найдены — положите в models/deepmosaics)", None) + idx = self.model_combo.findData(current) if current else 0 + self.model_combo.setCurrentIndex(max(0, idx)) + + def _with_browse(self, widget: QWidget, slot) -> QWidget: w = QWidget() h = QHBoxLayout(w) h.setContentsMargins(0, 0, 0, 0) - h.addWidget(line, 1) - btn = QPushButton("…") - btn.setMaximumWidth(32) + h.addWidget(widget, 1) + btn = QPushButton("Обзор…") btn.clicked.connect(slot) h.addWidget(btn) return w def _sync(self) -> None: is_dm = self.engine.currentData() == "deepmosaics" - for w in (self.dm_dir, self.dm_model, self.dm_python, self.dm_gpu): - w.setEnabled(is_dm) - - def _browse_dir(self) -> None: - d = QFileDialog.getExistingDirectory(self, "Папка DeepMosaics", self.dm_dir.text()) - if d: - self.dm_dir.setText(d) + self.model_combo.setEnabled(is_dm) + self.dm_gpu.setEnabled(is_dm) def _browse_model(self) -> None: - start = self.dm_model.text() or self.dm_dir.text() - p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", start, "Веса (*.pth);;Все файлы (*.*)") + p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)") if p: - self.dm_model.setText(p) - - def _browse_python(self) -> None: - p, _ = QFileDialog.getOpenFileName(self, "Python для DeepMosaics", self.dm_python.text(), "python (*.exe);;Все файлы (*.*)") - if p: - self.dm_python.setText(p) + if self.model_combo.findData(p) < 0: + self.model_combo.addItem(Path(p).stem, p) + self.model_combo.setCurrentIndex(self.model_combo.findData(p)) def apply_to_config(self) -> None: self._cfg.restorer = self.engine.currentData() - self._cfg.dm_dir = self.dm_dir.text().strip() or None - self._cfg.dm_model = self.dm_model.text().strip() or None - self._cfg.dm_python = self.dm_python.text().strip() or None + self._cfg.dm_model = self.model_combo.currentData() self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0" diff --git a/pyproject.toml b/pyproject.toml index 9320adf..11d2463 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,8 @@ version = "0.1.0" description = "Desktop GUI utility that detects and outlines already-applied censorship in images" readme = "README.md" requires-python = ">=3.11" -license = { text = "TBD" } +# GPL-3.0: the project vendors DeepMosaics (GPL-3.0) for mosaic restoration. +license = { text = "GPL-3.0-or-later" } authors = [{ name = "Leonid Pershin" }] dependencies = [ "PySide6>=6.6",