mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-07-07 00:47:50 +02:00
- Add the Jimeng-basic top-left "AI生成" pill as a CAPTURE-LESS mark (pill_engine.py): synthetic-silhouette edge-NCC detect + inpaint-only removal. Gated in remove_auto_marks: kept only when Jimeng is confirmed (TC260 metadata OR the bottom-right "★ 即梦AI" wordmark fired -- the wordmark keeps recall on metadata-STRIPPED uploads) AND Doubao did not fire. - Add an inpaint-fallback removal path + MI-GAN ONNX backend (migan extra, MIT, ~28 MB / ~1 GB peak -- droplet-friendly) alongside big-LaMa. New --method auto|reverse-alpha|inpaint (shared across visible/all/batch) and erase --backend migan; footprint_mask on each engine. - auto is deterministic: reverse-alpha for capture marks (recovers exact pixels, lighter -- measured cleaner than MI-GAN on structured backgrounds) and inpaint only for the capture-less pill. - --mark auto now removes EVERY detected mark in one pass (remove_auto_marks), so a Jimeng-basic image's top-left pill AND bottom-right wordmark both clear. - Bump 0.12.1 -> 0.13.0. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
@@ -21,8 +21,8 @@ Per-command exit-code semantics (the no-signal / GPU-missing skip branches), tes
|
||||
|
||||
- `uv run remove-ai-watermarks all <image.png> -o <output.png>` — full pipeline (visible + invisible + metadata). Same diffusion knobs as `invisible`, plus the visible-pass `--inpaint/--no-inpaint`/`--inpaint-method`. Skips step 2 (invisible/SynthID) when the `[gpu]` extra is absent or no invisible signal is detectable; see the module doc for the distinct exit codes.
|
||||
- `uv run remove-ai-watermarks invisible <image.png> -o <out.png>` — diffusion SynthID removal. **Full knob set** (kept identical across `invisible`/`all`/`batch`): `--strength` (vendor-adaptive default), `--steps`, `--guidance-scale` (CFG, default 7.5), `--pipeline sdxl|controlnet|qwen` (default `controlnet`; `qwen` is a manual opt-in only — see the qwen note in the module map), `--controlnet-scale`, `--model` (HF model id, default SDXL base), `--device`, `--seed`, `--hf-token`, `--max-resolution`/`--min-resolution`, `--upscaler lanczos|esrgan`, `--humanize` (Analog Humanizer grain), `--unsharp` (final sharpen), `--adaptive-polish/--no-adaptive-polish` (**ON by default**), `--tile/--no-tile` + `--tile-size`/`--tile-overlap` (**OFF by default**), `--force/--no-force` (default skip = ON, runs the scrub even with no detected signal). `--auto` is deprecated and a no-op that only warns. Skips the diffusion when no invisible signal is detectable (the no-signal gate); see the module doc.
|
||||
- `uv run remove-ai-watermarks visible <image.png> -o <out.png>` — known-visible-mark removal, CPU, no GPU. Reverse-alpha based. `--mark auto` (default) picks the strongest detected of the Gemini sparkle, Doubao "豆包AI生成", Jimeng "★ 即梦AI", and Samsung Galaxy AI "✦ Contenuti generati dall'AI"; `--mark gemini|doubao|jimeng|samsung` forces one. For arbitrary logos/objects use `erase`. When no known mark is detected the command writes no output and exits with the no-visible-mark code instead of re-serving the input; `--no-detect` forces the gemini fallback and proceeds. See the module doc for the routing/exit detail.
|
||||
- `uv run remove-ai-watermarks erase <image.png> --region x,y,w,h -o <out.png>` — universal region eraser (any logo/object, any position). `--backend cv2` (default, no deps) or `--backend lama` (big-LaMa via onnxruntime, extra `lama`); `--region` is repeatable.
|
||||
- `uv run remove-ai-watermarks visible <image.png> -o <out.png>` — known-visible-mark removal. `--method auto` (default) uses reverse-alpha for the capture marks (recovers exact pixels, CPU, no GPU -- measured cleaner + lighter than inpaint on them) and inpaint only for the capture-less pill; `--method reverse-alpha|inpaint` forces one (inpaint uses MI-GAN with the `migan` extra, else cv2). `--mark auto` (default) removes EVERY detected mark in one pass (a Jimeng-basic image carries the top-left "AI生成" pill AND the bottom-right "★ 即梦AI" wordmark) from: Gemini sparkle, Doubao "豆包AI生成", Jimeng "★ 即梦AI", Samsung Galaxy AI "✦ Contenuti generati dall'AI", and the capture-less Jimeng "AI生成" pill (top-left, metadata-gated); `--mark gemini|doubao|jimeng|samsung|jimeng_pill` forces one. For arbitrary logos/objects use `erase`. When no known mark is detected the command writes no output and exits with the no-visible-mark code instead of re-serving the input; `--no-detect` forces the gemini fallback and proceeds. See the module doc for the routing/exit detail. `--method` is shared across `visible`/`all`/`batch`.
|
||||
- `uv run remove-ai-watermarks erase <image.png> --region x,y,w,h -o <out.png>` — universal region eraser (any logo/object, any position). `--backend cv2` (default, no deps), `--backend migan` (MI-GAN via onnxruntime, extra `migan`; ~28 MB, ~1 GB RAM, near-LaMa), or `--backend lama` (big-LaMa, extra `lama`; best quality but ~4.7 GB RAM); `--region` is repeatable.
|
||||
- `uv run remove-ai-watermarks identify <image>` — provenance verdict (platform + watermark inventory + confidence); `--json` for machine output, `--no-visible` to skip the cv2 sparkle detector
|
||||
- `uv run remove-ai-watermarks metadata <image.png> --check` — inspect AI metadata (C2PA, EXIF, PNG chunks)
|
||||
- `uv run remove-ai-watermarks metadata <image.png> --remove -o <out.png>` — strip all AI metadata
|
||||
@@ -56,11 +56,12 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal
|
||||
- `noai/constants.py` — the single `C2PA_AI_VENDORS` registry (+ `C2PA_SOFT_BINDINGS`) from which `C2PA_ISSUERS` / `SYNTHID_C2PA_ISSUERS` / `identify._ISSUER_PLATFORM` are all derived. Add a new vendor as one registry entry; never edit the derived dicts and never add inline.
|
||||
- `metadata.py` — `scan_head(path)` is the shared (memoized) input for every C2PA/AIGC/IPTC byte scan; use it instead of `open().read(1MB)` for any new marker scan. Also home to `synthid_source`, `xai_signature`, `iptc_ai_system`, `aigc_label`, `huggingface_job`, `samsung_genai`, and `remove_ai_metadata` (fail-safe `strip_c2pa_boxes`). `exif_generator` matches a VALUE against `AI_GENERATOR_TOKENS` across EXIF `Software`/`Make`/`Artist`/`ImageDescription`, XMP `CreatorTool`, AND PNG `tEXt` chunks (`Software`/`Source`/`Title`/`Description` — NovelAI stamps there, not EXIF). **Detection and removal must stay in parity:** a generator that stamps an AI-shaped VALUE under a non-AI KEY (NovelAI's `Title`/`Source`) is dropped on removal by `_is_ai_value` (value-token match, mirrors `exif_generator`), NOT by `_is_ai_key` alone — else the cleaned file still reads as that generator. Add a new no-C2PA generator = one `AI_GENERATOR_TOKENS` entry (use a distinctive token, e.g. `reve.com` not bare `reve`); detection and removal then both follow. Regression: `tests/test_metadata.py::TestExifGenerator::{test_novelai_png_text_chunk_detected,test_novelai_removal_parity}`.
|
||||
- `identify.py` — aggregates every locally-readable signal into one `ProvenanceReport`; `is_ai_generated` is True or None, never asserted False. `ProvenanceReport.ai_source_kind` exposes the C2PA digital-source-type split — `"generated"` (trainedAlgorithmicMedia, fully AI) vs `"enhanced"` (compositeWithTrainedAlgorithmicMedia, a real photo with an AI-composited region), else None — so a caller branches full-frame scrub vs region-targeted clean (see `noai/tiling.feather_region_composite` + `WatermarkRemover.remove_watermark(region=...)`). The sparkle provenance threshold is the SHARED `watermark_registry.GEMINI_SPARKLE_TRUST_CONF` (imported, not a private copy) so the provenance "is there a sparkle" verdict and the removal "take the sparkle" decision can never drift. `import identify` is deliberately light (lazy `noai/__init__`, fits a 512 MB host) — keep heavy imports out (the `watermark_registry` constant import stays light: engines are lazy there). Add capture-camera tokens to `_DEVICE_C2PA_PLATFORM` only when verified against a real C2PA file; editing-app/AI-device signer tokens go to `_SIGNER_C2PA_PLATFORM`; generator/issuer platforms to `C2PA_AI_VENDORS` in `constants.py`. Integrity-clash detection is high-precision by design (only hard generator stamps feed it, source-grouped independence).
|
||||
- `watermark_registry.py` — the single catalog of known visible watermarks (gemini / doubao / jimeng / samsung), reverse-alpha based by policy. Add a new visible text mark = one `_text_mark(...)` row + a `TextMarkConfig` with a captured alpha map; do not re-add per-mark `if` branches. `cli._write_bgr_with_alpha` must NOT zero alpha in the watermark bbox (issue #30 white-box regression). **A new visible mark is HARD-BLOCKED on real controlled flat captures** — the alpha map is solved by `scripts/visible_alpha_solve.py` from solid black + gray (+white) captures of the mark produced by the actual app/device at native resolution, committed under `data/<engine>_capture/captures/`. Do NOT synthesize the alpha by font-rendering the wordmark — the user rejected synthetic reconstruction as below the quality bar (2026-06-22), and per-image alpha falloff/variation needs the real capture. Corpus images are NOT a substitute (data-safety: no corpus-derived committed assets; and reverse-alpha needs the flat capture, not real content). So if no flat capture exists for a mark, the work is parked — do not propose synthetic, do not derive from user uploads. (Meta AI, more Samsung locales, and any Grok visible mark are all parked on this; Grok additionally needs confirming it even HAS a visible mark — its known signal is EXIF-only `xai_signature`.)
|
||||
- `watermark_registry.py` — the single catalog of known visible watermarks (gemini / doubao / jimeng / samsung / jimeng_pill). **`--mark auto` removes EVERY detected mark in one pass** via `remove_auto_marks` (marks coexist -- a Jimeng-basic image has the top-left pill AND the bottom-right wordmark; `best_auto_mark` single-strongest would leave one). **The `jimeng_pill` is CAPTURE-LESS** (`has_capture=False`, `pill_engine.py`): the top-left "AI生成" label has no alpha map, so it is detect-by-synthetic-silhouette + inpaint-only (every method resolves to inpaint). Its weak edge-NCC detector (~7% raw false-fire) is gated in `remove_auto_marks`: the pill is kept only when the image is CONFIRMED Jimeng -- **TC260 metadata present (CLI `_aigc_metadata_present` → `pill_metadata`) OR the bottom-right "★ 即梦AI" wordmark fired** -- AND the **Doubao mark did NOT fire** (a Doubao image is also TC260 but not Jimeng-basic). The wordmark arm is what keeps recall on **metadata-STRIPPED uploads** (screenshots / re-saved files, ~61% of pills carry a detectable wordmark) that a metadata-only gate would miss; do NOT drop it. **Two removal methods, selected by `KnownMark.remove(method=...)` / CLI `--method`:** `reverse-alpha` recovers exact pixels from the captured alpha map (lighter, better on structured backgrounds), `inpaint` erases the mark footprint (`footprint_mask` → MI-GAN, or cv2 without the `migan` extra). `auto` (default) = `resolve_removal_method`: **reverse-alpha for capture marks, inpaint for capture-less** (deterministic, model-independent). Reverse-alpha is the default where a capture exists because it recovers the true pixels and is measured CLEANER than MI-GAN inpaint on the capture marks (doubao/gemini/jimeng — MI-GAN adds a rainbow/smudge artifact on structure) AND lighter (no model/RAM); inpaint is reserved for the capture-less pill. The inpaint backend is chosen by `preferred_inpaint_backend`: **MI-GAN when its extra is installed, else cv2** — big-LaMa is NOT auto-selected (both run on onnxruntime, so availability cannot express intent; big-LaMa stays an explicit `erase --backend lama` opt-in). The inpaint mask MUST be the full NCC-aligned alpha silhouette (`footprint_mask`), NOT the per-image `extract_mask` signature — the signature under-segments and leaves glyph residue (corpus-validated 2026-07). Add a new **reverse-alpha** text mark = one `_text_mark(...)` row + a `TextMarkConfig` with a captured alpha map; do not re-add per-mark `if` branches. A **capture-less inpaint mark** (no alpha map, like `jimeng_pill`) is instead a `has_capture=False` `KnownMark` + a small engine exposing `detect`/`footprint_mask` (geometry + synthetic silhouette) + a metadata/cross-mark gate for its detector -- this path does NOT need a flat capture (removal quality comes from the inpaint backend, so a font-rendered silhouette is data-safe and sufficient). `cli._write_bgr_with_alpha` must NOT zero alpha in the watermark bbox (issue #30 white-box regression). **A new REVERSE-ALPHA visible mark is HARD-BLOCKED on real controlled flat captures** — the alpha map is solved by `scripts/visible_alpha_solve.py` from solid black + gray (+white) captures of the mark produced by the actual app/device at native resolution, committed under `data/<engine>_capture/captures/`. Do NOT synthesize the alpha by font-rendering the wordmark — the user rejected synthetic reconstruction as below the quality bar (2026-06-22), and per-image alpha falloff/variation needs the real capture. Corpus images are NOT a substitute (data-safety: no corpus-derived committed assets; and reverse-alpha needs the flat capture, not real content). So if no flat capture exists for a mark, the work is parked — do not propose synthetic, do not derive from user uploads. (Meta AI, more Samsung locales, and any Grok visible mark are all parked on this; Grok additionally needs confirming it even HAS a visible mark — its known signal is EXIF-only `xai_signature`.)
|
||||
- `gemini_engine.py` — visible Gemini-sparkle remover/detector (cv2/numpy, no GPU): top-K size-weighted fusion candidate selection (`_SELECT_TOPK`), corner-promote, over/under-subtraction guards, false-positive gate, self-verify repair. Detection scores the top-K size-weighted matches by full fusion (spatial+gradient+variance) and keeps the highest — NOT the raw-NCC argmax, which re-admits the tiny-patch FPs the size weight suppresses (the osachub 2026-06-12 sub-0.85 corner-sparkle regression; see `docs/module-internals.md`). Keep the 0.85 corner-promote NCC gate; a margin/chroma-gated lower promote was measured and REJECTED 2026-06-11 (~33% FP on non-Google content). Gate any removal candidate on a physical brightness check, not the detector alone.
|
||||
- `_text_mark_engine.py` — shared base for the three reverse-alpha text-mark engines (extracted 2026-06-09); the per-engine modules are config-only subclasses. New text mark = a `TextMarkConfig` + a thin subclass + one registry row. Gemini stays a separate engine (different model).
|
||||
- `pill_engine.py` — the CAPTURE-LESS Jimeng-basic "AI生成" pill (top-left, issue #54). No alpha map: `detect` is edge-NCC of a synthetic font-rendered silhouette (`assets/jimeng_pill.png`, regenerate via `scripts/render_pill_silhouette.py`; committed, data-safe -- corpus stays out of the repo) in the top-left ROI, calibrated on 61 local real positives to threshold 0.22; `footprint_mask` is a generous FIXED top-left geometry box (NOT the NCC match position -- the synthetic silhouette localizes only approximately, the corner is negative space, so a geometry box removes cleanly while a match box leaves outline residue). Removal is inpaint-only (MI-GAN/cv2). Detector precision is weak (~7% raw false-fire), so it is registry-gated in `remove_auto_marks`: kept only when Jimeng is confirmed (TC260 metadata OR the bottom-right wordmark fired -- the wordmark survives metadata-STRIPPED uploads) AND Doubao did not fire; do NOT loosen those gates.
|
||||
- `doubao_engine.py` / `jimeng_engine.py` / `samsung_engine.py` — thin `TextMarkEngine` subclasses: Doubao "豆包AI生成" (bottom-right), Jimeng "★ 即梦AI" (bottom-right), Samsung Galaxy AI "✦ Contenuti generati dall'AI" (bottom-LEFT, locale-specific — Italian variant calibrated). Removal = reverse-alpha (always-align) + thin residual inpaint, **with an over-subtraction guard ported from `gemini_engine` (2026-06-20)**: `_reverse_alpha_oversubtracts` predicts the reverse-alpha output PER PIXEL over the glyph body from the INPUT, and when the recovered body lands more than `_OVERSUB_DARK_MARGIN` (25) gray levels below the local ring it abandons the reverse-alpha pixels and inpaints the footprint from the original surroundings (`_inpaint_footprint`) — fixing the dark-pit ghost on dark/mid-tone backgrounds (roadmap P0#8). Predicting per-pixel from the input (not the produced output) keeps a clean full-strength mark byte-identical (no false trip). A detector-only removal test is insufficient — assert visual residual (the textured-shift tests + `tests/test_text_mark_oversubtraction.py`).
|
||||
- `region_eraser.py` — universal region eraser (`erase` CLI): cv2 backend default (no deps), optional big-LaMa via onnxruntime (~3.5-4 GB peak RAM, ~5-6 s/call CPU — does not fit a minimal droplet).
|
||||
- `region_eraser.py` — universal region eraser (`erase` CLI) and the inpaint backend for the visible-mark fallback. Three backends: `cv2` (default, no deps), `migan` (MI-GAN ONNX, extra `migan`, MIT, ~28 MB / ~0.95 GB peak / ~0.19 s — the droplet-friendly tier, **the preferred backend for the auto inpaint fallback** via `watermark_registry.preferred_inpaint_backend`), `lama` (big-LaMa ONNX, extra `lama`, ~200 MB / ~4.7 GB peak — best quality, does not fit a minimal droplet, explicit opt-in only). **MI-GAN mask polarity is INVERTED** (0=hole/255=known) vs this package's 255-erase convention; `erase_migan` inverts before feeding the model (feeding 255=hole regenerates the whole frame into stripes — corpus-validated). Both ONNX models download on first use, never bundled.
|
||||
- `invisible_watermark.py` — decodes the OPEN DWT-DCT watermarks (SD / SDXL / FLUX) via `imwatermark` (extra `detect`, pulls torch). Fragile two ways: (1) does not survive JPEG re-encode/resize; (2) **carrier-fragile on a broad class of pristine images** -- a clean encode->decode round-trip recovers 48/48 on chatgpt/firefly/random but FAILS (28-39/48, below the `_MATCH_48`=44 gate) on the FLUX fox, doubao, a flat FLUX generation, AND a clean synthetic flat fill with no watermark. The failure does NOT track texture; it goes with a degenerate **all-ones decode that is a CARRIER ARTIFACT, not a watermark** (synthetic clean image reproduces it). So `detect_invisible_watermark` is **positive-only**: trust a hit; a `None` is inconclusive unless a same-carrier positive-control embed first recovers >=44. Verified 2026-06-19; full caveat in `docs/watermarking-landscape.md`.
|
||||
- `trustmark_detector.py` — Adobe TrustMark open decoder (extra `trustmark`). Do NOT remove the JPEG re-encode false-positive gate — a lone TrustMark hit without it is almost always content noise.
|
||||
- `noai/watermark_remover.py` — `WatermarkRemover` with three diffusion pipelines selected by the explicit `pipeline` ctor arg, never inferred from `model_id`: `sdxl` (plain SDXL img2img), `controlnet` (SDXL + canny ControlNet, **the DEFAULT since 2026-06-09**), and `qwen` (Qwen-Image 20B MMDiT img2img, Apache-2.0, CUDA/cloud-class — best **text** preservation (incl. CJK); `_load_qwen_pipeline`/`_run_qwen`, bf16, no MPS fallback; call shape in the pure `_build_qwen_kwargs` using `true_cfg_scale`). Removal comes from the img2img `strength`; ControlNet only preserves text/face STRUCTURE — SynthID CAN survive controlnet on photoreal content at low strength. Qwen CERTIFIED oracle floors (2026-06-20): OpenAI **0.10** (seed-robust, clean on seeds 0-4), Gemini **0.25** (seed 0 verified, pin a seed — Gemini oracle rate-limits volume; higher than the controlnet Gemini floor 0.15). `resolve_strength(..., pipeline="qwen")` carries the Qwen ladder (`_QWEN_VENDOR_STRENGTH`), so `--pipeline qwen` gets the 0.25 Gemini floor automatically (the old manual `--strength 0.25` workaround is retired). `_build_qwen_kwargs` passes an explicit `height`/`width` from the input (floored to /16 via `_qwen_target_size`) — without it the pipeline defaults to a 1024x1024 SQUARE and silently squishes non-square inputs (fixed 2026-06-20). **`qwen` is a MANUAL opt-in only — there is NO auto-router.** Measured (`scripts/fidelity_metrics.py`, OCR-CER / ArcFace / LPIPS / Laplacian-var, NOT eyeball): qwen beats controlnet on ONE niche only — **clean body text on a plain background, no faces** (openai_1/2 CER 0.241 vs 0.385). controlnet wins FACES (it always has) AND **display/decorative text in a scene** (abba poster: controlnet CER 0.114 vs qwen 0.379 — canny holds letter shapes, qwen re-renders and garbles them). So a content `--pipeline auto` router and a faces+text **mixed dual-pass** were prototyped and **DROPPED** (2026-06-20): on the canonical faces+text case controlnet wins every metric incl. text, so mixed loses; and "text→qwen" can't be auto-decided (it is body-vs-display text that matters, undetectable cheaply). qwen stays for callers who KNOW their content is clean-text-heavy and face-free. No face-restore extra ships, by validated decision (every restore approach looked MORE AI-generated). `remove_watermark(region=(x,y,w,h), region_feather=...)` runs the regeneration but feather-composites only the AI box back over the original (via `noai/tiling.feather_region_composite`), preserving the real photo elsewhere — the **AI-enhanced composite** path (`identify` `ai_source_kind == "enhanced"`); the box is supplied by the caller (a C2PA composite manifest carries no reliable machine-readable region, so we do not fabricate one).
|
||||
|
||||
@@ -165,7 +165,11 @@ Root cause: bad alpha (under-estimated, max ~0.65) + fixed-no-inpaint + tight bo
|
||||
|
||||
## `region_eraser.py`
|
||||
|
||||
`region_eraser.py` — universal region eraser (`erase` CLI). `erase(image, boxes=|mask=, backend=)` accepts grayscale (2D) and RGBA (4-channel) inputs on **both** backends (`erase_cv2` and `erase_lama` each split off any alpha plane and re-attach it unchanged, and promote grayscale to BGR for processing — LaMa would otherwise crash on grayscale and drop alpha on BGRA): `boxes_to_mask` → `cv2.inpaint` (`cv2` backend, default, no deps) or big-LaMa via onnxruntime (`lama` backend, extra `lama`, `Carve/LaMa-ONNX` Apache-2.0 model downloaded on first use, never bundled). `erase_lama` crops a padded region around the mask, runs LaMa at its fixed 512² input, pastes only masked pixels back (untouched areas stay pixel-exact). Lazy `_get_lama_session` singleton; `lama_available()` guards the optional import.
|
||||
`region_eraser.py` — universal region eraser (`erase` CLI) AND the inpaint backend for the visible-mark fallback (`watermark_registry._inpaint_remove`). `erase(image, boxes=|mask=, backend=)` accepts grayscale (2D) and RGBA (4-channel) inputs on **all** backends (each splits off any alpha plane and re-attaches it unchanged, and promotes grayscale to BGR): `boxes_to_mask` → one of three backends.
|
||||
- `cv2` (default, no deps): `cv2.inpaint`.
|
||||
- `migan` (extra `migan`, `andraniksargsyan/migan` ONNX, MIT, ~28 MB): `erase_migan`. The MI-GAN ONNX crops around the mask bbox and re-composites internally, so the FULL image is fed at native resolution; only masked pixels are pasted back. **Mask polarity is INVERTED** vs this package's 255-erase convention — the shipped ONNX wants 0=hole / 255=known, so `erase_migan` feeds `(mask<=127)*255`; feeding 255=hole regenerates the whole frame into stripes (corpus-validated 2026-07, cost hours to find). ~0.95 GB peak / ~0.19 s. This is the **preferred** inpaint-fallback backend.
|
||||
- `lama` (extra `lama`, `Carve/LaMa-ONNX` Apache-2.0, ~200 MB): `erase_lama` crops a padded region around the mask, runs at LaMa's fixed 512² input, pastes only masked pixels back. Best quality but ~4.7 GB peak — explicit opt-in only, NOT auto-selected.
|
||||
Lazy `_get_{lama,migan}_session` singletons; `{lama,migan}_available()` guard the optional imports (both == onnxruntime present). Note both extras install the same onnxruntime, so the two `*_available()` checks are identical — the registry's `preferred_inpaint_backend` therefore prefers MI-GAN whenever onnxruntime is present, and big-LaMa is reachable only by explicit `--backend lama`.
|
||||
|
||||
**LaMa-ONNX costs ~3.5-4 GB peak RAM and ~5-6 s/call on CPU** (FFC working set, not arena — `enable_cpu_mem_arena=False` does not help), so it does NOT fit a minimal droplet; the cv2 backend (tens of MB, ~30 ms) does. LaMa quality at low RAM = serverless/GPU, mirroring how raiw.cc offloads SDXL to fal.
|
||||
|
||||
@@ -254,7 +258,7 @@ Diffusion SynthID removal. The `--tile/--no-tile` knob is the *lossless* alterna
|
||||
|
||||
### `visible`
|
||||
|
||||
Known-visible-mark removal, CPU, no GPU. Reverse-alpha based: each mark is removed by inverting its captured alpha map. `--mark auto` (default) picks the strongest detected of the Gemini sparkle, the Doubao "豆包AI生成" text strip, the Jimeng "★ 即梦AI" wordmark, and the Samsung Galaxy AI "✦ Contenuti generati dall'AI" strip (bottom-LEFT, locale-specific — Italian variant calibrated); `--mark gemini` / `--mark doubao` / `--mark jimeng` / `--mark samsung` force one (choices come from the registry). Gemini/Doubao recover pixels exactly with no inpaint at native; **Jimeng and Samsung add an always-on thin residual inpaint over the glyph footprint** (their marks re-rasterize per image, so reverse-alpha alone leaves a faint outline). For arbitrary logos/objects use `erase`. **When `--mark auto` finds no known mark (the common case — ~74% of real uploads carry no registered visible mark), the command does NOT silently re-serve the input as a finished result.** It runs a cheap metadata-only `identify`, prints actionable guidance (if the image carries an invisible/metadata mark, e.g. an OpenAI/Gemini C2PA image, it points to `all`; otherwise it does NOT imply the image is clean -- it warns that an invisible pixel watermark like SynthID cannot be detected once the metadata proxy is gone and routes to both `all` and `erase --region`), writes NO output file, and exits **`EXIT_NO_VISIBLE_MARK` (2)** — distinct from success (0) and a hard error (1) so a wrapping service (raiw.cc) can surface the message instead of treating the unchanged image as done (the production "it didn't work" / score-0 trap). Same handling for an explicit `--mark <name>` that is not detected. Helper `cli._no_visible_mark_exit`; regression-guarded by `tests/test_cli.py::TestVisibleCommand::test_visible_auto_no_mark_exits_two_with_eraser_hint` and `test_visible_auto_no_mark_routes_to_all_when_metadata`. `--no-detect` still forces the gemini fallback and proceeds (exit 0).
|
||||
Known-visible-mark removal. Two methods via `--method` (default `auto`): `reverse-alpha` inverts the captured alpha map to recover exact pixels (CPU, no GPU, better on structured backgrounds); `inpaint` erases the mark footprint (`footprint_mask` → MI-GAN when the `migan` extra is installed, else cv2; see `preferred_inpaint_backend` — big-LaMa is NOT auto-selected). `auto` = `registry.resolve_removal_method`: reverse-alpha for capture marks, inpaint for capture-less (deterministic, model-independent -- reverse-alpha is measured cleaner + lighter than MI-GAN on the capture marks, so inpaint is reserved for the pill); the resolved method is echoed in the status line, and `--method` is shared across `visible`/`all`/`batch`. The inpaint mask is the full NCC-aligned alpha silhouette (`footprint_mask`), NOT the per-image `extract_mask` signature — the signature under-segments and leaves glyph residue (corpus-validated 2026-07; the mask is dilated to absorb alpha-alignment slop). `--mark auto` (default) removes EVERY detected mark in one pass (`registry.remove_auto_marks`, not the single strongest -- a Jimeng-basic image carries both the top-left pill and the bottom-right wordmark) from: the Gemini sparkle, the Doubao "豆包AI生成" text strip, the Jimeng "★ 即梦AI" wordmark, the Samsung Galaxy AI "✦ Contenuti generati dall'AI" strip (bottom-LEFT, Italian locale), and the capture-less Jimeng "AI生成" pill (top-left, `pill_engine`, inpaint-only). The pill's weak edge-NCC detector is gated in `remove_auto_marks`: kept only when the image is confirmed Jimeng -- TC260 metadata present (CLI `_aigc_metadata_present` → `pill_metadata`) OR the bottom-right wordmark fired -- AND Doubao did not fire. The wordmark arm keeps recall on metadata-STRIPPED uploads (screenshots / re-saves) that a metadata-only gate would miss. `--mark gemini|doubao|jimeng|samsung|jimeng_pill` forces one (choices come from the registry). Under `reverse-alpha`, Gemini/Doubao recover pixels exactly with no inpaint at native; **Jimeng and Samsung add an always-on thin residual inpaint over the glyph footprint** (their marks re-rasterize per image, so reverse-alpha alone leaves a faint outline). For arbitrary logos/objects use `erase`. **When `--mark auto` finds no known mark (the common case — ~74% of real uploads carry no registered visible mark), the command does NOT silently re-serve the input as a finished result.** It runs a cheap metadata-only `identify`, prints actionable guidance (if the image carries an invisible/metadata mark, e.g. an OpenAI/Gemini C2PA image, it points to `all`; otherwise it does NOT imply the image is clean -- it warns that an invisible pixel watermark like SynthID cannot be detected once the metadata proxy is gone and routes to both `all` and `erase --region`), writes NO output file, and exits **`EXIT_NO_VISIBLE_MARK` (2)** — distinct from success (0) and a hard error (1) so a wrapping service (raiw.cc) can surface the message instead of treating the unchanged image as done (the production "it didn't work" / score-0 trap). Same handling for an explicit `--mark <name>` that is not detected. Helper `cli._no_visible_mark_exit`; regression-guarded by `tests/test_cli.py::TestVisibleCommand::test_visible_auto_no_mark_exits_two_with_eraser_hint` and `test_visible_auto_no_mark_routes_to_all_when_metadata`. `--no-detect` still forces the gemini fallback and proceeds (exit 0).
|
||||
|
||||
### `batch`
|
||||
|
||||
|
||||
+11
-2
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "remove-ai-watermarks"
|
||||
version = "0.12.1"
|
||||
version = "0.13.0"
|
||||
description = "AI watermark remover: strip visible and invisible AI watermarks (Gemini / Nano Banana sparkle, SynthID) and provenance metadata (C2PA, EXIF) from images"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -112,6 +112,15 @@ lama = [
|
||||
"onnxruntime>=1.16.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
# Lightweight inpaint backend -- MI-GAN via onnxruntime (andraniksargsyan/migan,
|
||||
# MIT). CPU, no torch. Model (~28 MB) downloaded on first use and cached by
|
||||
# huggingface_hub; never bundled. ~700-950 MB peak RAM / ~0.19 s/call -- the
|
||||
# droplet-friendly tier (vs big-LaMa's ~4.7 GB) and the preferred inpaint backend
|
||||
# for the visible-mark fallback when installed. Same runtime as `lama`.
|
||||
migan = [
|
||||
"onnxruntime>=1.16.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
# Optional pre-diffusion super-resolution for small inputs (Real-ESRGAN). Loaded via
|
||||
# spandrel (MIT) -- a pure model-loader with NO basicsr dependency (it pulls only
|
||||
# torch / torchvision / safetensors / numpy / einops).
|
||||
@@ -136,7 +145,7 @@ dev = [
|
||||
"uv-outdated>=0.1.0; python_version >= '3.12'",
|
||||
"uv-secure>=0.12.0; python_version >= '3.12'",
|
||||
]
|
||||
all = ["remove-ai-watermarks[gpu,detect,trustmark,lama,dev]"]
|
||||
all = ["remove-ai-watermarks[gpu,detect,trustmark,lama,migan,dev]"]
|
||||
|
||||
# PyTorch Intel-GPU (XPU) wheel index. ``explicit = true`` keeps it inert for
|
||||
# the default CPU/CUDA install: uv consults it only when a torch install
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Render the SYNTHETIC 'AI生成' pill silhouette asset (data-safe, font-rendered).
|
||||
|
||||
The Jimeng-basic TC260 visible label is a rounded pill with 'AI生成' in the TOP-LEFT
|
||||
corner (issue #54). Unlike the reverse-alpha marks, this is a CAPTURE-LESS mark:
|
||||
the committed asset is a font-rendered binary SILHOUETTE (mark shape only, zero photo
|
||||
content), used ONLY to (a) detect the pill by edge-NCC in the top-left corner and
|
||||
(b) build the inpaint mask. It is NOT an alpha map -- removal quality comes from the
|
||||
inpaint backend (MI-GAN/cv2), so the silhouette need not be pixel-accurate, and the
|
||||
synthetic render keeps corpus/user content out of the tracked repo (data-safety).
|
||||
|
||||
Detection was calibrated on the retained local corpus (61 real positives + jimeng
|
||||
negatives): edge-NCC threshold ~0.22 in the top-left ROI. Re-run to regenerate the
|
||||
asset: uv run python scripts/render_pill_silhouette.py
|
||||
|
||||
Requires a CJK font (macOS STHeiti by default); the asset itself is committed, so this
|
||||
script only runs when regenerating it (never in CI).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
_ASSET = Path(__file__).resolve().parents[1] / "src" / "remove_ai_watermarks" / "assets" / "jimeng_pill.png"
|
||||
_FONT = "/System/Library/Fonts/STHeiti Medium.ttc"
|
||||
|
||||
|
||||
def render_silhouette(w: int = 320) -> np.ndarray:
|
||||
"""Rounded-pill outline + 'AI生成' text as a binary silhouette (255 = mark)."""
|
||||
h = int(w * 0.5)
|
||||
im = Image.new("L", (w, h), 0)
|
||||
d = ImageDraw.Draw(im)
|
||||
pad = int(w * 0.03)
|
||||
r = (h - 2 * pad) // 3
|
||||
d.rounded_rectangle([pad, pad, w - pad, h - pad], radius=r, outline=255, width=max(2, w // 90))
|
||||
fsz = int(h * 0.42)
|
||||
font = ImageFont.truetype(_FONT, fsz)
|
||||
txt = "AI生成"
|
||||
tb = d.textbbox((0, 0), txt, font=font)
|
||||
tw, th = tb[2] - tb[0], tb[3] - tb[1]
|
||||
d.text(((w - tw) // 2 - tb[0], (h - th) // 2 - tb[1]), txt, font=font, fill=255)
|
||||
return np.array(im)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
try:
|
||||
sil = render_silhouette()
|
||||
except OSError as e:
|
||||
print(f"Font not found ({e}); install a CJK font or edit _FONT.", file=sys.stderr)
|
||||
raise SystemExit(1) from e
|
||||
_ASSET.parent.mkdir(parents=True, exist_ok=True)
|
||||
Image.fromarray(sil).save(_ASSET)
|
||||
print(f"wrote {_ASSET} ({sil.shape[1]}x{sil.shape[0]})")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -11,4 +11,4 @@ _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
|
||||
_warnings.filterwarnings("ignore", message=r".*ImageProcessorFast.*")
|
||||
|
||||
|
||||
__version__ = "0.12.1"
|
||||
__version__ = "0.13.0"
|
||||
|
||||
@@ -504,3 +504,44 @@ class TextMarkEngine:
|
||||
rm = cv2.dilate(rm, kernel)
|
||||
best_out = cv2.inpaint(best_out, rm, c.residual_inpaint_radius, cv2.INPAINT_NS)
|
||||
return best_out
|
||||
|
||||
# ── Inpaint footprint (for the inpaint-fallback removal path) ────────
|
||||
|
||||
def footprint_mask(
|
||||
self, image: NDArray[Any], *, force: bool = False, dilate: int | None = None
|
||||
) -> NDArray[Any] | None:
|
||||
"""Full-frame uint8 mask (255 = mark) of the mark footprint, for the
|
||||
inpaint-fallback removal path (LaMa / cv2), or None if no placement fits.
|
||||
|
||||
``force`` is accepted for a uniform engine signature (the caller passes it to
|
||||
every engine) but ignored here -- the text-mark footprint is always the
|
||||
geometry-placed captured silhouette, present with or without a detection.
|
||||
|
||||
Uses the NCC-ALIGNED captured silhouette, NOT the per-image
|
||||
:meth:`extract_mask` signature: the signature under-segments the glyphs, so
|
||||
inpainting it leaves a residual ghost (corpus-validated 2026-07 -- Doubao
|
||||
left a "三包" remnant). The mask is dilated to absorb alpha-alignment slop
|
||||
(a scale/position mismatch at low detect confidence otherwise leaves a thin
|
||||
residual ring); ``dilate`` defaults to a mark-relative margin.
|
||||
|
||||
The caller gates on detection -- this returns the geometric footprint
|
||||
regardless, so a clean corner would be masked too.
|
||||
"""
|
||||
image = image_io.to_bgr(image)
|
||||
h, w = image.shape[:2]
|
||||
if h < 32 or w < 64:
|
||||
return None
|
||||
placed = self._aligned_alpha_map(image) or self._fixed_alpha_map(image)
|
||||
if placed is None:
|
||||
return None
|
||||
block, (ax, ay, gw, gh) = placed
|
||||
sil = (block > self.config.residual_alpha_floor).astype(np.uint8) * 255
|
||||
if int((sil > 0).sum()) == 0:
|
||||
return None
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
ch, cw = min(gh, h - ay), min(gw, w - ax)
|
||||
mask[ay : ay + ch, ax : ax + cw] = sil[:ch, :cw]
|
||||
d = dilate if dilate is not None else max(9, int(0.05 * gw))
|
||||
if d > 0:
|
||||
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1)))
|
||||
return mask
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 3.1 KiB |
+108
-51
@@ -293,6 +293,18 @@ _force_option = click.option(
|
||||
)
|
||||
|
||||
|
||||
_visible_method_option = click.option(
|
||||
"--method",
|
||||
"removal_method",
|
||||
type=click.Choice(["auto", "reverse-alpha", "inpaint"]),
|
||||
default="auto",
|
||||
help="Visible-mark removal method. auto: reverse-alpha for capture marks (exact "
|
||||
"pixels, lighter), inpaint for the capture-less pill. reverse-alpha recovers "
|
||||
"pixels from a captured alpha map; inpaint erases the footprint (MI-GAN with the "
|
||||
"'migan' extra, else cv2).",
|
||||
)
|
||||
|
||||
|
||||
def _resolve_auto_polish(auto: bool, adaptive_polish: bool) -> bool:
|
||||
"""Warn on the retired ``--auto`` flag, returning ``adaptive_polish`` unchanged.
|
||||
|
||||
@@ -327,9 +339,23 @@ def _warn_if_esrgan_unavailable(upscaler: str) -> None:
|
||||
console.print(" Note: --upscaler esrgan needs the 'esrgan' extra; falling back to Lanczos.")
|
||||
|
||||
|
||||
def _aigc_metadata_present(path: Path) -> bool:
|
||||
"""True when the file carries a China-AIGC (TC260) metadata label. Used to gate
|
||||
the weak-detector 'AI生成' pill: metadata confirms Jimeng-class provenance. NB
|
||||
this is only ONE of two confirmations -- ``remove_auto_marks`` also accepts the
|
||||
bottom-right wordmark, so a metadata-STRIPPED upload can still be handled."""
|
||||
with contextlib.suppress(Exception):
|
||||
from remove_ai_watermarks import metadata
|
||||
|
||||
return bool(metadata.aigc_label(path))
|
||||
return False
|
||||
|
||||
|
||||
def _remove_visible_auto(
|
||||
image: NDArray[Any],
|
||||
*,
|
||||
source_path: Path | None = None,
|
||||
removal_method: str = "auto",
|
||||
inpaint: bool = True,
|
||||
inpaint_method: str = "ns",
|
||||
inpaint_strength: float = 0.85,
|
||||
@@ -340,19 +366,26 @@ def _remove_visible_auto(
|
||||
standalone ``visible`` command uses, so EVERY registered mark is handled (the
|
||||
Gemini sparkle AND the Doubao/Jimeng/Samsung text marks), not just the sparkle.
|
||||
Returns ``(result, label-or-None)``; when no ``in_auto`` mark fires the image is
|
||||
returned unchanged with ``None``. ``inpaint*`` tune the Gemini edge-residual
|
||||
cleanup only (the text engines ignore them).
|
||||
returned unchanged with ``None``. ``removal_method`` selects reverse-alpha vs the
|
||||
inpaint fallback (see ``KnownMark.remove``); ``inpaint*`` tune the Gemini
|
||||
edge-residual cleanup only (the text engines ignore them).
|
||||
"""
|
||||
from remove_ai_watermarks import watermark_registry
|
||||
|
||||
best = watermark_registry.best_auto_mark(image)
|
||||
if best is None:
|
||||
return image, None
|
||||
rmethod: watermark_registry.RemovalMethod = removal_method # type: ignore[assignment]
|
||||
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
|
||||
result, _ = watermark_registry.get_mark(best.key).remove(
|
||||
image, inpaint_method=method, inpaint=inpaint, inpaint_strength=inpaint_strength, force=False
|
||||
pill_md = _aigc_metadata_present(source_path) if source_path is not None else False
|
||||
result, removed = watermark_registry.remove_auto_marks(
|
||||
image,
|
||||
pill_metadata=pill_md,
|
||||
method=rmethod,
|
||||
inpaint_method=method,
|
||||
inpaint=inpaint,
|
||||
inpaint_strength=inpaint_strength,
|
||||
)
|
||||
return result, best.label
|
||||
if not removed:
|
||||
return image, None
|
||||
return result, ", ".join(removed)
|
||||
|
||||
|
||||
# Exit code for the standalone ``visible`` command when no visible mark was
|
||||
@@ -535,8 +568,9 @@ def main(ctx: click.Context, verbose: bool) -> None:
|
||||
type=click.Choice(["auto", *watermark_registry.mark_keys()]),
|
||||
default="auto",
|
||||
help="Which known visible mark to target (auto picks the strongest detected). "
|
||||
"All marks are removed by exact reverse-alpha against a captured alpha map.",
|
||||
"Removal method is chosen by --method (default auto).",
|
||||
)
|
||||
@_visible_method_option
|
||||
@click.option("--strip-metadata/--keep-metadata", default=True, help="Strip AI metadata from output.")
|
||||
@click.pass_context
|
||||
def cmd_visible(
|
||||
@@ -548,14 +582,17 @@ def cmd_visible(
|
||||
inpaint_strength: float,
|
||||
detect: bool,
|
||||
mark: str,
|
||||
removal_method: str,
|
||||
strip_metadata: bool,
|
||||
) -> None:
|
||||
"""Remove a known visible AI watermark from an image.
|
||||
|
||||
Finds a known mark in its usual place (Gemini sparkle / Doubao text) via the
|
||||
watermark registry and removes it by exact reverse-alpha against a captured
|
||||
alpha map -- recovering the true pixels, not an inpaint guess. ``--mark auto``
|
||||
picks the strongest detected mark. For arbitrary logos/objects, use ``erase``.
|
||||
watermark registry and removes it. Default ``--method auto`` recovers the true
|
||||
pixels by exact reverse-alpha for the capture marks, and inpaints only the
|
||||
capture-less "AI生成" pill (MI-GAN with the ``migan`` extra, else cv2).
|
||||
``--mark auto`` picks the strongest detected mark. For arbitrary logos/objects,
|
||||
use ``erase``.
|
||||
"""
|
||||
from remove_ai_watermarks import watermark_registry as registry
|
||||
|
||||
@@ -574,41 +611,52 @@ def cmd_visible(
|
||||
h, w = image.shape[:2]
|
||||
console.print(f" Input: {source.name} ({w}x{h})")
|
||||
|
||||
# Resolve the target mark from the known-watermark registry. ``auto`` scans
|
||||
# every in-auto mark in its usual place and picks the strongest; an explicit
|
||||
# ``--mark <key>`` targets that one (the user asserts its presence).
|
||||
if mark == "auto":
|
||||
best = registry.best_auto_mark(image)
|
||||
if best is None:
|
||||
console.print(" No known visible mark detected (gemini / doubao / jimeng / samsung).")
|
||||
if detect:
|
||||
_no_visible_mark_exit(source)
|
||||
target = "gemini" # forced (no-detect): fall back to the default mark
|
||||
else:
|
||||
target = best.key
|
||||
console.print(f" Mark auto: {best.label} ({best.location}, conf {best.confidence:.2f})")
|
||||
else:
|
||||
target = mark
|
||||
|
||||
chosen = registry.get_mark(target)
|
||||
det = chosen.detect(image)
|
||||
if detect and not det.detected:
|
||||
console.print(f" {chosen.label} not detected (conf {det.confidence:.2f}). Use --no-detect to force.")
|
||||
_no_visible_mark_exit(source)
|
||||
if det.detected:
|
||||
console.print(f" {chosen.label} detected ({chosen.location}, conf {det.confidence:.2f})")
|
||||
|
||||
method: Literal["telea", "ns"] = "ns" if inpaint_method == "ns" else "telea"
|
||||
t0 = time.monotonic()
|
||||
with console.status(f"Removing {chosen.label}... ({chosen.recovery})"):
|
||||
result, _ = chosen.remove(
|
||||
image,
|
||||
inpaint_method=method,
|
||||
inpaint=inpaint,
|
||||
inpaint_strength=inpaint_strength,
|
||||
force=not detect,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
# ``auto`` removes EVERY detected in_auto mark in one pass (a Jimeng-basic image
|
||||
# carries the top-left pill AND the bottom-right wordmark); an explicit
|
||||
# ``--mark <key>`` targets that one (the user asserts its presence).
|
||||
if mark == "auto" and detect:
|
||||
t0 = time.monotonic()
|
||||
with console.status("Detecting & removing visible marks..."):
|
||||
result, removed = registry.remove_auto_marks(
|
||||
image,
|
||||
pill_metadata=_aigc_metadata_present(source),
|
||||
method=removal_method, # type: ignore[arg-type]
|
||||
inpaint_method=method,
|
||||
inpaint=inpaint,
|
||||
inpaint_strength=inpaint_strength,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
if not removed:
|
||||
console.print(" No known visible mark detected (gemini / doubao / jimeng / jimeng-pill / samsung).")
|
||||
_no_visible_mark_exit(source)
|
||||
console.print(f" Removed: {', '.join(removed)}")
|
||||
else:
|
||||
target = "gemini" if mark == "auto" else mark # --no-detect auto: gemini fallback
|
||||
chosen = registry.get_mark(target)
|
||||
det = chosen.detect(image)
|
||||
if detect and not det.detected:
|
||||
console.print(f" {chosen.label} not detected (conf {det.confidence:.2f}). Use --no-detect to force.")
|
||||
_no_visible_mark_exit(source)
|
||||
if det.detected:
|
||||
console.print(f" {chosen.label} detected ({chosen.location}, conf {det.confidence:.2f})")
|
||||
resolved = registry.resolve_removal_method(removal_method, chosen.has_capture) # type: ignore[arg-type]
|
||||
if resolved == "inpaint" and not registry.inpaint_model_available():
|
||||
console.print(
|
||||
" Note: --method inpaint using cv2 (install the 'migan' extra for a lightweight ONNX model)."
|
||||
)
|
||||
t0 = time.monotonic()
|
||||
with console.status(f"Removing {chosen.label}... ({resolved})"):
|
||||
result, _ = chosen.remove(
|
||||
image,
|
||||
method=removal_method, # type: ignore[arg-type]
|
||||
inpaint_method=method,
|
||||
inpaint=inpaint,
|
||||
inpaint_strength=inpaint_strength,
|
||||
force=not detect,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
# Save (rejoins the original alpha plane unchanged)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -653,9 +701,10 @@ def _parse_region(spec: str) -> tuple[int, int, int, int]:
|
||||
)
|
||||
@click.option(
|
||||
"--backend",
|
||||
type=click.Choice(["cv2", "lama"]),
|
||||
type=click.Choice(["cv2", "migan", "lama"]),
|
||||
default="cv2",
|
||||
help="Inpaint backend. cv2: instant, no deps. lama: onnxruntime big-LaMa, better quality (extra 'lama').",
|
||||
help="Inpaint backend. cv2: instant, no deps. migan: light ONNX MI-GAN, ~1 GB RAM, "
|
||||
"near-LaMa quality (extra 'migan'). lama: big-LaMa, best quality but ~4.7 GB RAM (extra 'lama').",
|
||||
)
|
||||
@click.option("--inpaint-method", type=click.Choice(["telea", "ns"]), default="telea", help="cv2 inpaint method.")
|
||||
@click.option("--dilate", type=int, default=3, help="Grow the box by this many px before inpainting.")
|
||||
@@ -666,7 +715,7 @@ def cmd_erase(
|
||||
source: Path,
|
||||
regions: tuple[str, ...],
|
||||
output: Path | None,
|
||||
backend: Literal["cv2", "lama"],
|
||||
backend: Literal["cv2", "migan", "lama"],
|
||||
inpaint_method: str,
|
||||
dilate: int,
|
||||
strip_metadata: bool,
|
||||
@@ -999,6 +1048,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
|
||||
@click.option(
|
||||
"--inpaint-method", type=click.Choice(["ns", "telea", "gaussian"]), default="ns", help="Inpainting method."
|
||||
)
|
||||
@_visible_method_option
|
||||
@_strength_option
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps for invisible removal.")
|
||||
@_pipeline_option
|
||||
@@ -1036,6 +1086,7 @@ def cmd_all(
|
||||
output: Path | None,
|
||||
inpaint: bool,
|
||||
inpaint_method: Literal["ns", "telea", "gaussian"],
|
||||
removal_method: str,
|
||||
strength: float | None,
|
||||
steps: int,
|
||||
pipeline: str,
|
||||
@@ -1105,7 +1156,9 @@ def cmd_all(
|
||||
console.print(f" Input: {source.name} ({w}x{h})")
|
||||
|
||||
with console.status("Removing visible watermark..."):
|
||||
result, removed_label = _remove_visible_auto(image, inpaint=inpaint, inpaint_method=inpaint_method)
|
||||
result, removed_label = _remove_visible_auto(
|
||||
image, source_path=source, removal_method=removal_method, inpaint=inpaint, inpaint_method=inpaint_method
|
||||
)
|
||||
if removed_label is not None:
|
||||
console.print(f" Visible watermark removed ({removed_label})")
|
||||
else:
|
||||
@@ -1244,6 +1297,7 @@ def _process_batch_image(
|
||||
seed: int | None,
|
||||
hf_token: str | None,
|
||||
humanize: float,
|
||||
removal_method: str = "auto",
|
||||
unsharp: float = 0.0,
|
||||
max_resolution: int = 0,
|
||||
min_resolution: int = 1024,
|
||||
@@ -1276,7 +1330,7 @@ def _process_batch_image(
|
||||
if image is None:
|
||||
raise ValueError("Failed to read image")
|
||||
|
||||
result, _ = _remove_visible_auto(image, inpaint=inpaint)
|
||||
result, _ = _remove_visible_auto(image, source_path=img_path, removal_method=removal_method, inpaint=inpaint)
|
||||
|
||||
_write_bgr_with_alpha(out_path, result, alpha)
|
||||
saved_alpha = alpha
|
||||
@@ -1361,6 +1415,7 @@ def _process_batch_image(
|
||||
@_strength_option
|
||||
@click.option("--steps", type=int, default=50, help="Number of denoising steps (invisible mode).")
|
||||
@click.option("--inpaint/--no-inpaint", default=True, help="Apply inpainting (visible mode).")
|
||||
@_visible_method_option
|
||||
@click.option(
|
||||
"--humanize", type=float, default=0.0, help="Analog Humanizer film grain intensity (0 = off, typical: 2.0-6.0)."
|
||||
)
|
||||
@@ -1402,6 +1457,7 @@ def cmd_batch(
|
||||
seed: int | None,
|
||||
hf_token: str | None,
|
||||
inpaint: bool,
|
||||
removal_method: str,
|
||||
humanize: float,
|
||||
unsharp: float,
|
||||
max_resolution: int,
|
||||
@@ -1468,6 +1524,7 @@ def cmd_batch(
|
||||
seed=seed,
|
||||
hf_token=hf_token,
|
||||
humanize=humanize,
|
||||
removal_method=removal_method,
|
||||
unsharp=unsharp,
|
||||
max_resolution=max_resolution,
|
||||
min_resolution=min_resolution,
|
||||
|
||||
@@ -603,6 +603,40 @@ class GeminiEngine:
|
||||
self._reverse_alpha_blend(result, alpha_map, pos)
|
||||
return self._verify_and_repair(result, alpha_map, pos, size)
|
||||
|
||||
def footprint_mask(self, image: NDArray[Any], *, force: bool = False, dilate: int = 13) -> NDArray[Any] | None:
|
||||
"""Full-frame uint8 mask (255 = sparkle) of the sparkle footprint, for the
|
||||
inpaint-fallback removal path (LaMa / cv2), or None.
|
||||
|
||||
The footprint is the interpolated captured alpha at the detected scale --
|
||||
the same region reverse-alpha operates on. When ``force`` and nothing is
|
||||
detected, falls back to the default sparkle slot for the image size (the
|
||||
``--no-detect`` path). The caller gates on the trust-confidence detection.
|
||||
"""
|
||||
image = image_io.to_bgr(image)
|
||||
h, w = image.shape[:2]
|
||||
det = self.detect_watermark(image)
|
||||
if det.detected:
|
||||
x, y, scale = det.region[0], det.region[1], det.region[2]
|
||||
elif force:
|
||||
cfg = get_watermark_config(w, h)
|
||||
x, y = cfg.get_position(w, h)
|
||||
scale = cfg.logo_size
|
||||
else:
|
||||
return None
|
||||
alpha = self.get_interpolated_alpha(scale)
|
||||
fp = self._footprint_indices(alpha, (x, y), image.shape)
|
||||
if fp is None:
|
||||
return None
|
||||
aroi, (y1, y2, x1, x2) = fp
|
||||
sil = (aroi > 0.10).astype(np.uint8) * 255
|
||||
if int((sil > 0).sum()) == 0:
|
||||
return None
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
mask[y1:y2, x1:x2] = sil
|
||||
if dilate > 0:
|
||||
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * dilate + 1, 2 * dilate + 1)))
|
||||
return mask
|
||||
|
||||
def remove_watermark_custom(
|
||||
self,
|
||||
image: NDArray[Any],
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Jimeng-basic 'AI生成' pill: a CAPTURE-LESS visible mark (issue #54).
|
||||
|
||||
The Jimeng free-tier TC260 label is a rounded pill with 'AI生成' in the TOP-LEFT
|
||||
corner -- distinct from the reverse-alpha ``jimeng`` "★ 即梦AI" mark (bottom-right).
|
||||
No flat capture / alpha map exists for it, so it is removed by INPAINT, not
|
||||
reverse-alpha:
|
||||
|
||||
* Detect: edge-NCC of a font-rendered SILHOUETTE (``assets/jimeng_pill.png``,
|
||||
synthetic, data-safe -- see ``scripts/render_pill_silhouette.py``) against the
|
||||
top-left ROI, at the pill's known width fraction. Corpus-calibrated threshold
|
||||
(61 real positives + jimeng negatives): ``_DETECT_THRESHOLD`` 0.22.
|
||||
* Remove: place the pill footprint at the matched location and inpaint it
|
||||
(MI-GAN / cv2 via the registry). Quality comes from the inpaint backend, so the
|
||||
silhouette need not be pixel-accurate -- which is why a synthetic render is
|
||||
sufficient and no corpus-derived asset is committed.
|
||||
|
||||
Geometry measured on 51 real examples (8 resolutions, all 3:4): width ~0.161*W,
|
||||
height ~0.091*W, top-left, margins ~0.02-0.05.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, NamedTuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from remove_ai_watermarks import image_io
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from numpy.typing import NDArray
|
||||
|
||||
# cv2/numpy boundary: cv2 ships no usable type info, so strict pyright cannot know
|
||||
# its array element types. Relax the unknown-type rules for this file only; the
|
||||
# public signatures are still annotated with NDArray[Any].
|
||||
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportUnknownParameterType=false, reportMissingTypeArgument=false, reportMissingTypeStubs=false, reportMissingImports=false, reportArgumentType=false, reportAssignmentType=false, reportReturnType=false, reportCallIssue=false, reportIndexIssue=false, reportOperatorIssue=false, reportOptionalMemberAccess=false, reportOptionalSubscript=false, reportAttributeAccessIssue=false, reportUnnecessaryComparison=false
|
||||
|
||||
_ASSET = Path(__file__).parent / "assets" / "jimeng_pill.png"
|
||||
|
||||
# Geometry (fractions of image WIDTH unless noted); top-left corner.
|
||||
_WIDTH_FRAC = 0.161
|
||||
_ROI_W_FRAC = 0.34 # search window width (of W)
|
||||
_ROI_H_FRAC = 0.14 # search window height (of H)
|
||||
_DETECT_THRESHOLD = 0.22 # edge-NCC gate, corpus-calibrated
|
||||
# Inpaint mask GEOMETRY (fractions of W unless noted): a generous fixed top-left box
|
||||
# covering the pill (measured ~0.167*W wide, ~0.09*W tall, margin ~0.02-0.05) plus
|
||||
# margin. The mask uses stable geometry, NOT the NCC match position -- the synthetic
|
||||
# silhouette localizes only approximately, and the corner is negative space, so
|
||||
# over-covering is harmless while a match-positioned box leaves outline residue.
|
||||
_MASK_X0, _MASK_Y0 = 0.012, 0.006 # x0 of W, y0 of H
|
||||
_MASK_W, _MASK_H = 0.205, 0.115 # width of W, height of W
|
||||
|
||||
_silhouette: NDArray[Any] | None = None
|
||||
|
||||
|
||||
class PillDetection(NamedTuple):
|
||||
detected: bool
|
||||
confidence: float
|
||||
region: tuple[int, int, int, int] # x, y, w, h of the matched pill
|
||||
|
||||
|
||||
def _load_silhouette() -> NDArray[Any] | None:
|
||||
global _silhouette
|
||||
if _silhouette is None:
|
||||
if not _ASSET.exists():
|
||||
return None
|
||||
_silhouette = image_io.imread(str(_ASSET), cv2.IMREAD_GRAYSCALE)
|
||||
return _silhouette
|
||||
|
||||
|
||||
def _grad(gray: NDArray[Any]) -> NDArray[Any]:
|
||||
gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
|
||||
gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
|
||||
return cv2.normalize(cv2.magnitude(gx, gy), None, 0, 255, cv2.NORM_MINMAX)
|
||||
|
||||
|
||||
class PillEngine:
|
||||
"""Detect + inpaint-mask the top-left 'AI生成' pill (edge-NCC, no reverse-alpha)."""
|
||||
|
||||
def _match(self, image: NDArray[Any]) -> tuple[float, tuple[int, int, int, int]] | None:
|
||||
sil = _load_silhouette()
|
||||
if sil is None or image is None or image.size == 0:
|
||||
return None
|
||||
h, w = image.shape[:2]
|
||||
if h < 64 or w < 64:
|
||||
return None
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image
|
||||
rh, rw = int(h * _ROI_H_FRAC), int(w * _ROI_W_FRAC)
|
||||
roi = gray[0:rh, 0:rw]
|
||||
tw = max(24, int(_WIDTH_FRAC * w))
|
||||
th = max(12, int(tw * sil.shape[0] / sil.shape[1]))
|
||||
if th >= rh or tw >= rw:
|
||||
return None
|
||||
tmpl = cv2.resize(sil, (tw, th))
|
||||
res = cv2.matchTemplate(_grad(roi.astype(np.float32)), _grad(tmpl.astype(np.float32)), cv2.TM_CCOEFF_NORMED)
|
||||
_, score, _, loc = cv2.minMaxLoc(res)
|
||||
return float(score), (int(loc[0]), int(loc[1]), tw, th)
|
||||
|
||||
def detect(self, image: NDArray[Any]) -> PillDetection:
|
||||
m = self._match(image)
|
||||
if m is None:
|
||||
return PillDetection(False, 0.0, (0, 0, 0, 0))
|
||||
score, box = m
|
||||
return PillDetection(score >= _DETECT_THRESHOLD, score, box)
|
||||
|
||||
def footprint_mask(self, image: NDArray[Any], *, force: bool = False) -> NDArray[Any] | None:
|
||||
"""Full-frame uint8 mask (255 = pill) over the pill's known top-left region.
|
||||
|
||||
Uses stable GEOMETRY (a generous fixed box), not the NCC match position: the
|
||||
synthetic silhouette localizes only approximately, so a match-positioned mask
|
||||
leaves outline residue, while the top-left corner is negative space, so a
|
||||
generous geometric box removes the pill cleanly and harmlessly. The caller
|
||||
gates on :meth:`detect`, so a clean corner is never masked. ``force`` is
|
||||
accepted for a uniform engine signature but ignored (the geometry box is
|
||||
fixed regardless)."""
|
||||
if image is None or image.size == 0:
|
||||
return None
|
||||
h, w = image.shape[:2]
|
||||
x0, y0 = int(_MASK_X0 * w), int(_MASK_Y0 * h)
|
||||
x1, y1 = min(w, x0 + int(_MASK_W * w)), min(h, y0 + int(_MASK_H * w))
|
||||
if x1 <= x0 or y1 <= y0:
|
||||
return None
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
mask[y0:y1, x0:x1] = 255
|
||||
return mask
|
||||
@@ -9,10 +9,15 @@ deterministic per-generator engines (Gemini sparkle, Doubao) do not cover.
|
||||
Backends:
|
||||
- ``cv2`` (default): ``cv2.inpaint`` (Telea / Navier-Stokes). Instant, no extra
|
||||
dependencies, lower quality on large or textured regions.
|
||||
- ``migan`` (optional, extra ``migan``): MI-GAN via onnxruntime
|
||||
(``andraniksargsyan/migan``, MIT). CPU, ~28 MB model, ~700-950 MB peak RAM,
|
||||
~0.19 s/call -- the droplet-friendly tier: near-big-LaMa quality on small
|
||||
marks at ~5x less RAM and ~8x faster. Model downloaded on first use.
|
||||
- ``lama`` (optional, extra ``lama``): big-LaMa via onnxruntime
|
||||
(``Carve/LaMa-ONNX``, Apache-2.0). CPU, resolution-robust, much better on
|
||||
texture. The model (~200 MB) is downloaded on first use and cached by
|
||||
huggingface_hub; it is never bundled in this repo.
|
||||
(``Carve/LaMa-ONNX``, Apache-2.0). CPU, resolution-robust, best quality on
|
||||
texture but ~200 MB model and ~4.7 GB peak RAM (too heavy for a small host).
|
||||
The model is downloaded on first use and cached by huggingface_hub; it is
|
||||
never bundled in this repo.
|
||||
"""
|
||||
|
||||
# cv2/numpy boundary: cv2 ships no usable type info, so strict pyright cannot know
|
||||
@@ -32,13 +37,17 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
Backend = Literal["cv2", "lama"]
|
||||
Backend = Literal["cv2", "lama", "migan"]
|
||||
|
||||
_LAMA_REPO = "Carve/LaMa-ONNX"
|
||||
_LAMA_FILE = "lama_fp32.onnx"
|
||||
|
||||
# Cached onnxruntime session (loading is expensive; reuse across calls).
|
||||
_MIGAN_REPO = "andraniksargsyan/migan"
|
||||
_MIGAN_FILE = "migan.onnx"
|
||||
|
||||
# Cached onnxruntime sessions (loading is expensive; reuse across calls).
|
||||
_lama_session: object | None = None
|
||||
_migan_session: object | None = None
|
||||
|
||||
|
||||
def boxes_to_mask(
|
||||
@@ -167,6 +176,72 @@ def erase_lama(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
|
||||
return result
|
||||
|
||||
|
||||
def migan_available() -> bool:
|
||||
"""True when the optional MI-GAN backend can run (onnxruntime installed)."""
|
||||
from .optional_deps import module_available
|
||||
|
||||
return module_available("onnxruntime")
|
||||
|
||||
|
||||
def _get_migan_session() -> object:
|
||||
"""Load (once) the MI-GAN ONNX session, downloading the model on first use."""
|
||||
global _migan_session
|
||||
if _migan_session is not None:
|
||||
return _migan_session
|
||||
|
||||
import onnxruntime as ort
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
model_path = hf_hub_download(repo_id=_MIGAN_REPO, filename=_MIGAN_FILE)
|
||||
logger.info("Loading MI-GAN ONNX model: %s", model_path)
|
||||
_migan_session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
|
||||
return _migan_session
|
||||
|
||||
|
||||
def erase_migan(image_bgr: NDArray[Any], mask: NDArray[Any]) -> NDArray[Any]:
|
||||
"""Inpaint ``mask`` (255 = erase) with MI-GAN via onnxruntime (CPU).
|
||||
|
||||
The MI-GAN ONNX pipeline crops around the mask bbox internally and re-composites,
|
||||
so the full image is fed at native resolution. Only the masked pixels are pasted
|
||||
back, so untouched areas stay pixel-exact.
|
||||
|
||||
Mask polarity: the shipped ``andraniksargsyan/migan`` ONNX expects 0 = hole
|
||||
(inpaint) / 255 = known (keep) -- the INVERSE of this package's 255-erase
|
||||
convention -- so the mask is inverted before feeding the model (corpus-validated
|
||||
2026-07; feeding 255=hole regenerates the whole frame into stripes).
|
||||
|
||||
Like ``erase_lama``, accepts 1-channel (grayscale) and 4-channel (BGRA) input.
|
||||
"""
|
||||
if image_bgr.ndim == 2:
|
||||
bgr = erase_migan(cv2.cvtColor(image_bgr, cv2.COLOR_GRAY2BGR), mask)
|
||||
return cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
|
||||
if image_bgr.ndim == 3 and image_bgr.shape[2] == 4:
|
||||
bgr = erase_migan(np.ascontiguousarray(image_bgr[:, :, :3]), mask)
|
||||
return np.dstack([bgr, image_bgr[:, :, 3]])
|
||||
|
||||
session = _get_migan_session()
|
||||
inp = session.get_inputs() # type: ignore[attr-defined]
|
||||
img_name, mask_name = inp[0].name, inp[1].name
|
||||
h, w = image_bgr.shape[:2]
|
||||
|
||||
rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
|
||||
img_in = np.transpose(rgb, (2, 0, 1))[None].astype(np.uint8) # (1,3,H,W)
|
||||
# invert to MI-GAN polarity: 255 where KNOWN (keep), 0 where hole (erase)
|
||||
known = (mask <= 127).astype(np.uint8) * 255
|
||||
mask_in = known[None, None] # (1,1,H,W)
|
||||
|
||||
out = session.run(None, {img_name: img_in, mask_name: mask_in})[0] # type: ignore[attr-defined]
|
||||
res = np.transpose(np.asarray(out)[0], (1, 2, 0)).astype(np.uint8) # (H',W',3) RGB
|
||||
if res.shape[:2] != (h, w):
|
||||
res = cv2.resize(res, (w, h), interpolation=cv2.INTER_LINEAR)
|
||||
out_bgr = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
|
||||
|
||||
result = image_bgr.copy()
|
||||
hole = mask > 127
|
||||
result[hole] = out_bgr[hole]
|
||||
return result
|
||||
|
||||
|
||||
def erase(
|
||||
image_bgr: NDArray[Any],
|
||||
*,
|
||||
@@ -191,6 +266,12 @@ def erase(
|
||||
if not mask.any():
|
||||
return image_bgr.copy()
|
||||
|
||||
if backend == "migan":
|
||||
if not migan_available():
|
||||
raise RuntimeError(
|
||||
"MI-GAN backend requires onnxruntime. Install the extra: pip install 'remove-ai-watermarks[migan]'"
|
||||
)
|
||||
return erase_migan(image_bgr, mask)
|
||||
if backend == "lama":
|
||||
if not lama_available():
|
||||
raise RuntimeError(
|
||||
|
||||
@@ -41,6 +41,11 @@ if TYPE_CHECKING:
|
||||
InpaintMethod = Literal["telea", "ns"]
|
||||
Region = tuple[int, int, int, int]
|
||||
|
||||
# Removal method selection for the visible pass. ``auto`` prefers the inpaint
|
||||
# fallback (LaMa footprint) when the ``lama`` extra is installed, else reverse-alpha
|
||||
# for marks with a captured alpha map (cv2 inpaint for capture-less marks).
|
||||
RemovalMethod = Literal["auto", "reverse-alpha", "inpaint"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MarkDetection:
|
||||
@@ -62,9 +67,10 @@ class KnownMark:
|
||||
label: str
|
||||
location: str # usual place, human-readable ("bottom-right")
|
||||
in_auto: bool # participate in `--mark auto` scanning
|
||||
recovery: str # removal strategy (all reverse-alpha today)
|
||||
recovery: str # default removal strategy label ("reverse-alpha")
|
||||
_detect: Callable[[NDArray[Any]], MarkDetection]
|
||||
_remove: Callable[..., tuple[NDArray[Any], Region | None]]
|
||||
has_capture: bool = True # a captured alpha map exists (reverse-alpha possible)
|
||||
|
||||
def detect(self, image: NDArray[Any]) -> MarkDetection:
|
||||
return self._detect(image)
|
||||
@@ -73,20 +79,29 @@ class KnownMark:
|
||||
self,
|
||||
image: NDArray[Any],
|
||||
*,
|
||||
method: RemovalMethod = "auto",
|
||||
inpaint_method: InpaintMethod = "ns",
|
||||
inpaint: bool = True,
|
||||
inpaint_strength: float = 0.85,
|
||||
force: bool = False,
|
||||
) -> tuple[NDArray[Any], Region | None]:
|
||||
"""Remove this mark by reverse-alpha; returns ``(result, region)`` where
|
||||
``region`` is the removed mark's bbox (for residual-inpaint positioning),
|
||||
or None if nothing was removed. NB: the CLI does NOT use ``region`` to
|
||||
clear alpha on save -- that zeroing caused the issue-#30 white box.
|
||||
"""Remove this mark; returns ``(result, region)`` where ``region`` is the
|
||||
removed mark's bbox (for residual-inpaint positioning), or None if nothing
|
||||
was removed. NB: the CLI does NOT use ``region`` to clear alpha on save --
|
||||
that zeroing caused the issue-#30 white box.
|
||||
|
||||
``inpaint`` / ``inpaint_strength`` / ``inpaint_method`` tune the Gemini
|
||||
reverse-alpha edge-residual cleanup only. ``force`` removes at the mark's
|
||||
usual location even without a positive detection (the ``--no-detect`` path).
|
||||
``method`` selects the removal path: ``reverse-alpha`` recovers the true
|
||||
pixels from the captured alpha map (lighter, exact, better on structured
|
||||
backgrounds); ``inpaint`` erases the footprint with LaMa (or cv2 when the
|
||||
``lama`` extra is absent), needing no capture; ``auto`` (default) prefers
|
||||
inpaint when LaMa is installed and falls back to reverse-alpha otherwise
|
||||
(cv2 inpaint for capture-less marks). ``inpaint``/``inpaint_strength``/
|
||||
``inpaint_method`` tune the Gemini reverse-alpha edge-residual cleanup only.
|
||||
``force`` removes at the mark's usual location even without a positive
|
||||
detection (the ``--no-detect`` path).
|
||||
"""
|
||||
if resolve_removal_method(method, self.has_capture) == "inpaint":
|
||||
return _inpaint_remove(self, image, force)
|
||||
return self._remove(image, inpaint_method, inpaint, inpaint_strength, force)
|
||||
|
||||
|
||||
@@ -135,11 +150,74 @@ def _engine(key: str) -> Any:
|
||||
from remove_ai_watermarks.samsung_engine import SamsungEngine
|
||||
|
||||
_engines[key] = SamsungEngine()
|
||||
elif key == "jimeng_pill":
|
||||
from remove_ai_watermarks.pill_engine import PillEngine
|
||||
|
||||
_engines[key] = PillEngine()
|
||||
else: # pragma: no cover - guarded by the registry keys
|
||||
raise KeyError(key)
|
||||
return _engines[key]
|
||||
|
||||
|
||||
def inpaint_model_available() -> bool:
|
||||
"""True when any ONNX inpaint-model backend (MI-GAN or big-LaMa) can run."""
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
return region_eraser.migan_available() or region_eraser.lama_available()
|
||||
|
||||
|
||||
def preferred_inpaint_backend() -> str:
|
||||
"""Backend used by the inpaint fallback: MI-GAN (light, droplet-friendly, the
|
||||
default) when its ONNX runtime is available, else cv2. big-LaMa is NOT auto-
|
||||
selected -- it is a heavier explicit opt-in via ``erase --backend lama`` (both
|
||||
models run on onnxruntime, so availability alone cannot express the user's
|
||||
intent; the light model is the safe default)."""
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
return "migan" if region_eraser.migan_available() else "cv2"
|
||||
|
||||
|
||||
def resolve_removal_method(method: RemovalMethod, has_capture: bool) -> Literal["reverse-alpha", "inpaint"]:
|
||||
"""Resolve the requested method to a concrete one. A capture-less mark has no
|
||||
alpha map, so it can only be inpainted -- even explicit ``reverse-alpha`` falls
|
||||
back to inpaint there.
|
||||
|
||||
``auto`` uses **reverse-alpha for capture marks** and **inpaint for capture-less**.
|
||||
Reverse-alpha recovers the true pixels under the mark (measured cleaner than
|
||||
MI-GAN inpaint on the capture marks -- doubao/gemini/jimeng -- especially on
|
||||
structured backgrounds, and it needs no model / RAM), so it stays the default
|
||||
where a capture exists; inpaint is reserved for marks that have no alpha map
|
||||
(the Jimeng pill). ``--method inpaint`` still forces inpaint for anyone who
|
||||
wants it."""
|
||||
# A capture-less mark can only be inpainted; a capture mark inpaints only when
|
||||
# explicitly asked (auto + reverse-alpha both recover the true pixels).
|
||||
if method == "inpaint" or not has_capture:
|
||||
return "inpaint"
|
||||
return "reverse-alpha"
|
||||
|
||||
|
||||
def _inpaint_remove(mark: KnownMark, image: NDArray[Any], force: bool) -> tuple[NDArray[Any], Region | None]:
|
||||
"""Remove ``mark`` by inpainting its footprint: the NCC-aligned captured
|
||||
silhouette (:meth:`footprint_mask`), erased with MI-GAN when its extra is
|
||||
installed, else cv2 (see :func:`preferred_inpaint_backend`). No-op (returns a
|
||||
copy) on a clean corner unless ``force``. Gated on the mark's trust-confidence
|
||||
detection, so a clean image is never touched."""
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
det = mark.detect(image)
|
||||
if not (det.detected or force):
|
||||
return image.copy(), None
|
||||
engine = _engine(mark.key)
|
||||
fm = engine.footprint_mask(image, force=force) # uniform signature; text/pill ignore force
|
||||
if fm is None or not fm.any():
|
||||
return image.copy(), None
|
||||
if preferred_inpaint_backend() == "migan":
|
||||
result = region_eraser.erase_migan(image, fm)
|
||||
else:
|
||||
result = region_eraser.erase_cv2(image, fm, radius=6)
|
||||
return result, (det.region if det.detected else None)
|
||||
|
||||
|
||||
def _gemini_detect(image: NDArray[Any]) -> MarkDetection:
|
||||
d = _engine("gemini").detect_watermark(image)
|
||||
detected = bool(d.detected) and d.confidence >= _GEMINI_AUTO_MIN_CONF
|
||||
@@ -207,11 +285,29 @@ def _text_mark(key: str, label: str, location: str) -> KnownMark:
|
||||
)
|
||||
|
||||
|
||||
# ── Capture-less mark: the Jimeng-basic "AI生成" pill (top-left, inpaint-only) ──
|
||||
# No alpha map exists, so removal is inpaint only (has_capture=False routes every
|
||||
# method to inpaint). Detection is edge-NCC of a synthetic silhouette; see pill_engine.
|
||||
def _pill_detect(image: NDArray[Any]) -> MarkDetection:
|
||||
d = _engine("jimeng_pill").detect(image)
|
||||
return MarkDetection("jimeng_pill", "Jimeng AI生成 pill", "top-left", d.detected, d.confidence, d.region)
|
||||
|
||||
|
||||
def _pill_noop_remove(
|
||||
image: NDArray[Any], _im: InpaintMethod, _ip: bool, _st: float, _force: bool
|
||||
) -> tuple[NDArray[Any], Region | None]:
|
||||
# Capture-less: reverse-alpha is impossible. resolve_removal_method routes every
|
||||
# method to inpaint (_inpaint_remove), so this reverse-alpha slot is never reached;
|
||||
# return an untouched copy defensively.
|
||||
return image.copy(), None
|
||||
|
||||
|
||||
_REGISTRY: tuple[KnownMark, ...] = (
|
||||
KnownMark("gemini", "Google Gemini sparkle", "bottom-right", True, "reverse-alpha", _gemini_detect, _gemini_remove),
|
||||
_text_mark("doubao", "Doubao 豆包AI生成 text", "bottom-right"),
|
||||
_text_mark("jimeng", "Jimeng 即梦AI wordmark", "bottom-right"),
|
||||
_text_mark("samsung", "Samsung Galaxy AI text", "bottom-left"),
|
||||
KnownMark("jimeng_pill", "Jimeng AI生成 pill", "top-left", True, "inpaint", _pill_detect, _pill_noop_remove, False),
|
||||
)
|
||||
|
||||
|
||||
@@ -238,8 +334,7 @@ def detect_marks(image: NDArray[Any], *, include_explicit: bool = True) -> list[
|
||||
|
||||
Returns one MarkDetection per scanned mark (``detected`` flags which fired).
|
||||
``include_explicit=False`` scans only the ``in_auto`` marks -- the set used
|
||||
by ``--mark auto``.
|
||||
"""
|
||||
by ``--mark auto``."""
|
||||
return [m.detect(image) for m in _REGISTRY if include_explicit or m.in_auto]
|
||||
|
||||
|
||||
@@ -247,3 +342,45 @@ def best_auto_mark(image: NDArray[Any]) -> MarkDetection | None:
|
||||
"""The highest-confidence detected ``in_auto`` mark, or None if none fired."""
|
||||
fired = [d for d in detect_marks(image, include_explicit=False) if d.detected]
|
||||
return max(fired, key=lambda d: d.confidence) if fired else None
|
||||
|
||||
|
||||
def remove_auto_marks(
|
||||
image: NDArray[Any],
|
||||
*,
|
||||
pill_metadata: bool = False,
|
||||
method: RemovalMethod = "auto",
|
||||
inpaint_method: InpaintMethod = "ns",
|
||||
inpaint: bool = True,
|
||||
inpaint_strength: float = 0.85,
|
||||
) -> tuple[NDArray[Any], list[str]]:
|
||||
"""Remove EVERY detected ``in_auto`` mark in one pass, chaining the result.
|
||||
|
||||
Marks coexist in different corners -- a Jimeng-basic image carries BOTH the
|
||||
top-left "AI生成" pill AND the bottom-right "★ 即梦AI" wordmark -- and their
|
||||
confidences are on different scales, so ``best_auto_mark`` (single strongest)
|
||||
would clean only one and leave the other (issue #54). Each mark re-detects at
|
||||
its own corner on the progressively-cleaned image, so order does not matter.
|
||||
|
||||
The capture-less ``jimeng_pill`` has a weak edge-NCC detector (~7% raw
|
||||
false-fire), so it is kept only when the image is CONFIRMED Jimeng and NOT
|
||||
Doubao. Confirmation is ``pill_metadata`` (the China-AIGC / TC260 metadata
|
||||
signal, supplied by the caller) OR the reliable bottom-right wordmark firing --
|
||||
the wordmark keeps recall on metadata-STRIPPED uploads (screenshots / re-saved
|
||||
files) that the metadata gate alone would miss. Returns ``(result, [labels
|
||||
removed])``; an empty list means nothing fired."""
|
||||
fired = [d for d in detect_marks(image, include_explicit=False) if d.detected]
|
||||
keys = {d.key for d in fired}
|
||||
jimeng_confirmed = pill_metadata or "jimeng" in keys
|
||||
if "jimeng_pill" in keys and (not jimeng_confirmed or "doubao" in keys):
|
||||
fired = [d for d in fired if d.key != "jimeng_pill"]
|
||||
result = image
|
||||
for det in fired:
|
||||
result, _ = get_mark(det.key).remove(
|
||||
result,
|
||||
method=method,
|
||||
inpaint_method=inpaint_method,
|
||||
inpaint=inpaint,
|
||||
inpaint_strength=inpaint_strength,
|
||||
force=False,
|
||||
)
|
||||
return result, [d.label for d in fired]
|
||||
|
||||
+10
-4
@@ -459,19 +459,25 @@ class TestAllCommand:
|
||||
|
||||
def test_all_visible_step_uses_registry(self, runner, sample_png, tmp_path):
|
||||
"""Regression (#1): the `all` visible step must route through the registry
|
||||
(best_auto_mark), so Doubao/Jimeng/Samsung text marks are handled -- not just
|
||||
the Gemini sparkle via a hardcoded GeminiEngine."""
|
||||
(remove_auto_marks), so Doubao/Jimeng/Samsung/pill marks are handled -- not
|
||||
just the Gemini sparkle via a hardcoded GeminiEngine."""
|
||||
mock_cls, _mock_engine = _mock_invisible_engine()
|
||||
output = tmp_path / "clean.png"
|
||||
|
||||
def _fake_remove_auto(image, **kwargs):
|
||||
return image, []
|
||||
|
||||
with (
|
||||
patch("remove_ai_watermarks.cli.InvisibleEngine", mock_cls, create=True),
|
||||
patch("remove_ai_watermarks.invisible_engine.InvisibleEngine", mock_cls),
|
||||
patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True),
|
||||
patch("remove_ai_watermarks.watermark_registry.best_auto_mark", return_value=None) as mock_best,
|
||||
patch(
|
||||
"remove_ai_watermarks.watermark_registry.remove_auto_marks", side_effect=_fake_remove_auto
|
||||
) as mock_auto,
|
||||
):
|
||||
result = runner.invoke(main, ["all", str(sample_png), "-o", str(output), "--force"])
|
||||
assert result.exit_code == 0, result.output
|
||||
mock_best.assert_called() # the registry auto-detector drove the visible pass
|
||||
mock_auto.assert_called() # the registry auto-detector drove the visible pass
|
||||
|
||||
def test_all_skips_invisible_on_no_signal_but_succeeds(self, runner, sample_png, tmp_path):
|
||||
"""P0#5: with no detectable invisible watermark and no --force, `all` skips
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
"""Inpaint-fallback visible removal: method resolution, footprint masks, dispatch.
|
||||
|
||||
The inpaint path erases the mark footprint (MI-GAN when the ``migan`` extra is
|
||||
installed, else cv2) instead of reverse-alpha, so a mark needs no captured alpha
|
||||
map for removal (only for the footprint silhouette). ``auto`` is deterministic:
|
||||
reverse-alpha for capture marks, inpaint for capture-less. These tests avoid any
|
||||
ONNX model download by pinning the backend to cv2 via ``preferred_inpaint_backend``;
|
||||
only pure cv2/numpy paths run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks import watermark_registry as registry
|
||||
from remove_ai_watermarks.doubao_engine import DoubaoEngine
|
||||
from remove_ai_watermarks.gemini_engine import GeminiEngine
|
||||
|
||||
|
||||
def _compose_textmark(engine, bg: float = 120.0, w: int = 1024, h: int = 1024):
|
||||
"""Composite the engine's captured mark onto a flat ``bg`` at full opacity so
|
||||
the mark is detectable. Returns ``(watermarked_uint8, (ax, ay, gw, gh))``."""
|
||||
img = np.full((h, w, 3), float(bg), np.float32)
|
||||
block, (ax, ay, gw, gh) = engine._fixed_alpha_map(img)
|
||||
a = np.clip(block, 0.0, 0.99)[:, :, None]
|
||||
logo = np.array(engine.config.alpha_logo_bgr, np.float32)
|
||||
img[ay : ay + gh, ax : ax + gw] = img[ay : ay + gh, ax : ax + gw] * (1 - a) + logo * a
|
||||
return np.clip(img, 0, 255).astype(np.uint8), (ax, ay, gw, gh)
|
||||
|
||||
|
||||
class TestResolveMethod:
|
||||
@pytest.mark.parametrize("explicit", ["reverse-alpha", "inpaint"])
|
||||
def test_explicit_passthrough(self, explicit: str) -> None:
|
||||
# capture mark: the explicit method passes through unchanged
|
||||
assert registry.resolve_removal_method(explicit, True) == explicit # type: ignore[arg-type]
|
||||
# capture-less: reverse-alpha is impossible, so it collapses to inpaint
|
||||
expected = "inpaint" if explicit == "reverse-alpha" else explicit
|
||||
assert registry.resolve_removal_method(explicit, False) == expected # type: ignore[arg-type]
|
||||
|
||||
def test_auto_uses_reverse_alpha_for_capture_marks(self) -> None:
|
||||
# auto is deterministic and model-independent: reverse-alpha where a capture
|
||||
# exists (cleaner + lighter than inpaint), inpaint only where it does not.
|
||||
assert registry.resolve_removal_method("auto", True) == "reverse-alpha"
|
||||
|
||||
def test_auto_uses_inpaint_for_capture_less(self) -> None:
|
||||
assert registry.resolve_removal_method("auto", False) == "inpaint"
|
||||
|
||||
|
||||
class TestFootprintMask:
|
||||
def test_textmark_footprint_geometry(self) -> None:
|
||||
mask = DoubaoEngine().footprint_mask(np.full((1024, 1024, 3), 120, np.uint8))
|
||||
assert mask is not None
|
||||
assert mask.shape == (1024, 1024)
|
||||
assert mask.dtype == np.uint8
|
||||
assert mask.any()
|
||||
# Doubao sits bottom-right: the mask mass is in the bottom-right quadrant.
|
||||
ys, xs = np.where(mask > 0)
|
||||
assert ys.mean() > 512
|
||||
assert xs.mean() > 512
|
||||
|
||||
def test_textmark_small_image_returns_none(self) -> None:
|
||||
assert DoubaoEngine().footprint_mask(np.full((20, 20, 3), 120, np.uint8)) is None
|
||||
|
||||
def test_gemini_footprint_needs_detection_or_force(self) -> None:
|
||||
eng = GeminiEngine()
|
||||
clean = np.full((1024, 1024, 3), 128, np.uint8)
|
||||
assert eng.footprint_mask(clean) is None # nothing detected -> no mask
|
||||
forced = eng.footprint_mask(clean, force=True) # default sparkle slot
|
||||
assert forced is not None
|
||||
assert forced.any()
|
||||
|
||||
|
||||
class TestInpaintDispatch:
|
||||
"""Force the cv2 backend (patch preferred_inpaint_backend) so no ONNX model
|
||||
downloads; the dispatch/gating logic is backend-agnostic."""
|
||||
|
||||
def test_clean_image_is_untouched(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registry, "preferred_inpaint_backend", lambda: "cv2")
|
||||
img = np.full((1024, 1024, 3), 120, np.uint8)
|
||||
out, region = registry.get_mark("doubao").remove(img, method="inpaint")
|
||||
assert region is None
|
||||
assert np.array_equal(out, img) # not detected, not forced -> no-op
|
||||
|
||||
def test_forced_inpaint_edits_only_footprint(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registry, "preferred_inpaint_backend", lambda: "cv2")
|
||||
img, (ax, ay, gw, gh) = _compose_textmark(DoubaoEngine())
|
||||
out, _ = registry.get_mark("doubao").remove(img, method="inpaint", force=True)
|
||||
assert not np.array_equal(out[ay : ay + gh, ax : ax + gw], img[ay : ay + gh, ax : ax + gw])
|
||||
assert np.array_equal(out[:200, :200], img[:200, :200]) # far corner untouched
|
||||
|
||||
def test_detected_inpaint_lowers_confidence(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registry, "preferred_inpaint_backend", lambda: "cv2")
|
||||
mark = registry.get_mark("doubao")
|
||||
img, _ = _compose_textmark(DoubaoEngine())
|
||||
before = mark.detect(img)
|
||||
assert before.detected # the composed mark is detectable
|
||||
out, region = mark.remove(img, method="inpaint")
|
||||
assert region is not None
|
||||
assert mark.detect(out).confidence < before.confidence
|
||||
|
||||
def test_reverse_alpha_method_still_selectable(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(registry, "inpaint_model_available", lambda: True) # would be inpaint on auto
|
||||
img, _ = _compose_textmark(DoubaoEngine())
|
||||
# explicit reverse-alpha bypasses the inpaint fallback even with a model present
|
||||
out, region = registry.get_mark("doubao").remove(img, method="reverse-alpha")
|
||||
assert region is not None
|
||||
assert not np.array_equal(out, img)
|
||||
|
||||
|
||||
class TestBackendSelection:
|
||||
"""MI-GAN is the preferred inpaint backend (light, droplet-friendly); big-LaMa
|
||||
is NOT auto-selected. cv2 is the floor when no ONNX model is present."""
|
||||
|
||||
def test_prefers_migan_when_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
monkeypatch.setattr(region_eraser, "migan_available", lambda: True)
|
||||
assert registry.preferred_inpaint_backend() == "migan"
|
||||
|
||||
def test_cv2_when_no_model(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
monkeypatch.setattr(region_eraser, "migan_available", lambda: False)
|
||||
assert registry.preferred_inpaint_backend() == "cv2"
|
||||
|
||||
def test_inpaint_model_available_reflects_either(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
monkeypatch.setattr(region_eraser, "migan_available", lambda: False)
|
||||
monkeypatch.setattr(region_eraser, "lama_available", lambda: False)
|
||||
assert not registry.inpaint_model_available()
|
||||
monkeypatch.setattr(region_eraser, "lama_available", lambda: True)
|
||||
assert registry.inpaint_model_available()
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Jimeng-basic 'AI生成' pill: capture-less mark (detect via synthetic silhouette
|
||||
edge-NCC, remove via inpaint). No model download -- cv2 fallback / pure logic only."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
from remove_ai_watermarks import watermark_registry as registry
|
||||
from remove_ai_watermarks.pill_engine import _DETECT_THRESHOLD, PillEngine
|
||||
|
||||
_FONT = "/System/Library/Fonts/STHeiti Medium.ttc"
|
||||
|
||||
|
||||
def _font_ok() -> bool:
|
||||
try:
|
||||
ImageFont.truetype(_FONT, 20)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
_HAS_FONT = _font_ok()
|
||||
_needs_font = pytest.mark.skipif(
|
||||
not _HAS_FONT, reason="CJK font unavailable (compose helper needs it; asset is committed)"
|
||||
)
|
||||
|
||||
|
||||
def _compose_pill(w: int = 1200, h: int = 1600, bg: int = 150) -> np.ndarray:
|
||||
"""Composite a semi-transparent 'AI生成' pill top-left onto a flat BGR frame."""
|
||||
img = Image.new("RGB", (w, h), (bg, bg, bg))
|
||||
ov = Image.new("RGBA", (w, h), (0, 0, 0, 0))
|
||||
d = ImageDraw.Draw(ov)
|
||||
mw, mh = int(0.167 * w), int(0.09 * w)
|
||||
mx, my = int(0.03 * w), int(0.02 * w)
|
||||
d.rounded_rectangle([mx, my, mx + mw, my + mh], radius=mh // 3, outline=(255, 255, 255, 150), width=3)
|
||||
font = ImageFont.truetype(_FONT, int(mh * 0.5))
|
||||
d.text((mx + mw // 6, my + mh // 5), "AI生成", font=font, fill=(255, 255, 255, 170))
|
||||
out = Image.alpha_composite(img.convert("RGBA"), ov).convert("RGB")
|
||||
return np.asarray(out)[:, :, ::-1].copy() # RGB->BGR
|
||||
|
||||
|
||||
class TestPillDetect:
|
||||
@_needs_font
|
||||
def test_detects_composited_pill(self) -> None:
|
||||
det = PillEngine().detect(_compose_pill())
|
||||
assert det.detected
|
||||
assert det.confidence >= _DETECT_THRESHOLD
|
||||
|
||||
def test_clean_frame_does_not_fire(self) -> None:
|
||||
clean = np.full((1600, 1200, 3), 150, np.uint8)
|
||||
assert not PillEngine().detect(clean).detected
|
||||
|
||||
def test_small_image_no_fire(self) -> None:
|
||||
assert not PillEngine().detect(np.full((40, 40, 3), 150, np.uint8)).detected
|
||||
|
||||
|
||||
class TestPillMask:
|
||||
def test_footprint_mask_top_left_geometry(self) -> None:
|
||||
mask = PillEngine().footprint_mask(np.full((1600, 1200, 3), 150, np.uint8))
|
||||
assert mask is not None
|
||||
assert mask.shape == (1600, 1200)
|
||||
assert mask.any()
|
||||
ys, xs = np.where(mask > 0)
|
||||
# pill sits top-left: mask mass in the top-left quadrant
|
||||
assert ys.mean() < 800
|
||||
assert xs.mean() < 600
|
||||
|
||||
|
||||
class TestPillRegistry:
|
||||
def test_pill_is_capture_less(self) -> None:
|
||||
m = registry.get_mark("jimeng_pill")
|
||||
assert m.has_capture is False
|
||||
|
||||
def test_capture_less_routes_every_method_to_inpaint(self) -> None:
|
||||
# a capture-less mark cannot reverse-alpha; even explicit reverse-alpha -> inpaint
|
||||
assert registry.resolve_removal_method("reverse-alpha", False) == "inpaint"
|
||||
assert registry.resolve_removal_method("auto", False) == "inpaint"
|
||||
assert registry.resolve_removal_method("inpaint", False) == "inpaint"
|
||||
|
||||
|
||||
class TestPillGate:
|
||||
"""The pill is kept only when the image is CONFIRMED Jimeng (TC260 metadata OR
|
||||
the bottom-right wordmark fired) and NOT Doubao. Fakes detect_marks so no image
|
||||
content is needed; cv2 backend so nothing downloads."""
|
||||
|
||||
@staticmethod
|
||||
def _fakes(monkeypatch: pytest.MonkeyPatch, keys: set[str]) -> None:
|
||||
from remove_ai_watermarks.watermark_registry import MarkDetection
|
||||
|
||||
labels = {
|
||||
"doubao": "Doubao 豆包AI生成 text",
|
||||
"jimeng": "Jimeng 即梦AI wordmark",
|
||||
"jimeng_pill": "Jimeng AI生成 pill",
|
||||
}
|
||||
monkeypatch.setattr(registry, "preferred_inpaint_backend", lambda: "cv2")
|
||||
monkeypatch.setattr(
|
||||
registry,
|
||||
"detect_marks",
|
||||
lambda image, *, include_explicit=True: [
|
||||
MarkDetection(k, labels[k], "loc", True, 0.6, (10, 10, 40, 40)) for k in keys
|
||||
],
|
||||
)
|
||||
|
||||
def test_pill_kept_with_metadata(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
self._fakes(monkeypatch, {"jimeng_pill"})
|
||||
_, removed = registry.remove_auto_marks(np.full((400, 300, 3), 150, np.uint8), pill_metadata=True)
|
||||
assert "Jimeng AI生成 pill" in removed
|
||||
|
||||
def test_pill_kept_via_wordmark_without_metadata(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# metadata-stripped upload: the bottom-right wordmark confirms Jimeng
|
||||
self._fakes(monkeypatch, {"jimeng", "jimeng_pill"})
|
||||
_, removed = registry.remove_auto_marks(np.full((400, 300, 3), 150, np.uint8), pill_metadata=False)
|
||||
assert "Jimeng AI生成 pill" in removed
|
||||
|
||||
def test_pill_dropped_without_metadata_or_wordmark(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
self._fakes(monkeypatch, {"jimeng_pill"})
|
||||
_, removed = registry.remove_auto_marks(np.full((400, 300, 3), 150, np.uint8), pill_metadata=False)
|
||||
assert "Jimeng AI生成 pill" not in removed
|
||||
|
||||
def test_pill_dropped_on_doubao_even_with_metadata(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
self._fakes(monkeypatch, {"doubao", "jimeng_pill"})
|
||||
_, removed = registry.remove_auto_marks(np.full((400, 300, 3), 150, np.uint8), pill_metadata=True)
|
||||
assert "Doubao 豆包AI生成 text" in removed
|
||||
assert "Jimeng AI生成 pill" not in removed
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from remove_ai_watermarks.region_eraser import boxes_to_mask, erase, lama_available
|
||||
from remove_ai_watermarks.region_eraser import boxes_to_mask, erase, lama_available, migan_available
|
||||
|
||||
|
||||
class TestBoxesToMask:
|
||||
@@ -134,3 +134,74 @@ class TestLamaChannelHandling:
|
||||
out = erase(bgra, boxes=[(40, 40, 20, 20)], backend="lama")
|
||||
assert out.shape == bgra.shape
|
||||
assert np.array_equal(out[..., 3], bgra[..., 3]) # alpha carried through unchanged
|
||||
|
||||
|
||||
class TestMiganBackend:
|
||||
def test_migan_raises_when_unavailable(self):
|
||||
img = np.full((100, 100, 3), 50, np.uint8)
|
||||
if migan_available():
|
||||
pytest.skip("onnxruntime installed; cannot test the unavailable path")
|
||||
with pytest.raises(RuntimeError, match="onnxruntime"):
|
||||
erase(img, boxes=[(10, 10, 20, 20)], backend="migan")
|
||||
|
||||
|
||||
class TestMiganWrapper:
|
||||
"""erase_migan without the real model: fake session returns a solid-red field
|
||||
and captures the fed mask. Exercises the mask-polarity inversion, masked-only
|
||||
compositing, and grayscale/BGRA channel handling."""
|
||||
|
||||
captured: dict
|
||||
|
||||
@pytest.fixture
|
||||
def _fake_migan(self, monkeypatch: pytest.MonkeyPatch):
|
||||
from remove_ai_watermarks import region_eraser
|
||||
|
||||
self.captured = {}
|
||||
|
||||
class _In:
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, outer):
|
||||
self.outer = outer
|
||||
|
||||
def get_inputs(self):
|
||||
return [_In("image"), _In("mask")]
|
||||
|
||||
def run(self, _outputs, feeds):
|
||||
self.outer.captured["mask"] = feeds["mask"]
|
||||
img = feeds["image"] # (1,3,H,W) RGB
|
||||
red = np.zeros_like(img)
|
||||
red[:, 0] = 255 # pure red in RGB
|
||||
return [red]
|
||||
|
||||
monkeypatch.setattr(region_eraser, "migan_available", lambda: True)
|
||||
monkeypatch.setattr(region_eraser, "_get_migan_session", lambda: _FakeSession(self))
|
||||
|
||||
@pytest.mark.usefixtures("_fake_migan")
|
||||
def test_composites_only_masked_region_and_inverts_mask(self):
|
||||
img = np.full((100, 100, 3), 120, np.uint8) # BGR
|
||||
out = erase(img, boxes=[(40, 40, 20, 20)], backend="migan", dilate=0)
|
||||
# inside the box -> red (BGR (0,0,255)); outside -> untouched
|
||||
assert tuple(int(v) for v in out[50, 50]) == (0, 0, 255)
|
||||
assert np.array_equal(out[:30, :30], img[:30, :30])
|
||||
# mask fed to MI-GAN is inverted: 0 (hole) inside the box, 255 (known) outside
|
||||
m = self.captured["mask"][0, 0]
|
||||
assert m[50, 50] == 0
|
||||
assert m[10, 10] == 255
|
||||
|
||||
@pytest.mark.usefixtures("_fake_migan")
|
||||
def test_grayscale_2d_does_not_raise(self):
|
||||
gray = np.full((100, 100), 120, np.uint8)
|
||||
out = erase(gray, boxes=[(40, 40, 20, 20)], backend="migan", dilate=0)
|
||||
assert out.ndim == 2
|
||||
assert out.shape == gray.shape
|
||||
|
||||
@pytest.mark.usefixtures("_fake_migan")
|
||||
def test_bgra_preserves_alpha(self):
|
||||
bgra = np.full((100, 100, 4), 120, np.uint8)
|
||||
bgra[..., 3] = 200
|
||||
out = erase(bgra, boxes=[(40, 40, 20, 20)], backend="migan", dilate=0)
|
||||
assert out.shape == bgra.shape
|
||||
assert np.array_equal(out[..., 3], bgra[..., 3])
|
||||
|
||||
@@ -14,14 +14,17 @@ DOUBAO_SAMPLE = Path(__file__).resolve().parents[1] / "data" / "samples" / "doub
|
||||
|
||||
class TestCatalog:
|
||||
def test_keys(self):
|
||||
assert reg.mark_keys() == ["gemini", "doubao", "jimeng", "samsung"]
|
||||
assert reg.mark_keys() == ["gemini", "doubao", "jimeng", "samsung", "jimeng_pill"]
|
||||
|
||||
def test_all_in_auto(self):
|
||||
assert all(m.in_auto for m in reg.known_marks())
|
||||
|
||||
def test_recovery_is_reverse_alpha(self):
|
||||
# Every catalogued mark is removed by exact reverse-alpha (no inpaint).
|
||||
assert all(m.recovery == "reverse-alpha" for m in reg.known_marks())
|
||||
def test_recovery(self):
|
||||
# Capture marks recover by reverse-alpha; the capture-less pill is inpaint-only.
|
||||
by_key = {m.key: m for m in reg.known_marks()}
|
||||
assert all(by_key[k].recovery == "reverse-alpha" for k in ("gemini", "doubao", "jimeng", "samsung"))
|
||||
assert by_key["jimeng_pill"].recovery == "inpaint"
|
||||
assert by_key["jimeng_pill"].has_capture is False
|
||||
|
||||
def test_locations(self):
|
||||
by_key = {m.key: m for m in reg.known_marks()}
|
||||
@@ -29,6 +32,7 @@ class TestCatalog:
|
||||
assert by_key["doubao"].location == "bottom-right"
|
||||
assert by_key["jimeng"].location == "bottom-right"
|
||||
assert by_key["samsung"].location == "bottom-left"
|
||||
assert by_key["jimeng_pill"].location == "top-left"
|
||||
|
||||
def test_get_mark_unknown_raises(self):
|
||||
with pytest.raises(KeyError):
|
||||
@@ -39,7 +43,7 @@ class TestScan:
|
||||
def test_detect_marks_scans_all(self):
|
||||
img = np.zeros((256, 256, 3), np.uint8)
|
||||
keys = {d.key for d in reg.detect_marks(img)}
|
||||
assert keys == {"gemini", "doubao", "jimeng", "samsung"}
|
||||
assert keys == {"gemini", "doubao", "jimeng", "samsung", "jimeng_pill"}
|
||||
|
||||
def test_blank_image_no_auto_mark(self):
|
||||
assert reg.best_auto_mark(np.zeros((256, 256, 3), np.uint8)) is None
|
||||
|
||||
@@ -651,7 +651,7 @@ name = "cuda-bindings"
|
||||
version = "13.3.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cuda-pathfinder", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "cuda-pathfinder", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
|
||||
@@ -1491,7 +1491,7 @@ name = "nvidia-cublas"
|
||||
version = "13.1.1.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cuda-nvrtc", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
|
||||
@@ -1530,7 +1530,7 @@ name = "nvidia-cudnn-cu13"
|
||||
version = "9.20.0.48"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cublas", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
|
||||
@@ -1542,7 +1542,7 @@ name = "nvidia-cufft"
|
||||
version = "12.0.0.61"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
|
||||
@@ -1572,9 +1572,9 @@ name = "nvidia-cusolver"
|
||||
version = "12.0.4.66"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cusparse", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-nvjitlink", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cublas", marker = "sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-cusparse", marker = "sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
|
||||
@@ -1586,7 +1586,7 @@ name = "nvidia-cusparse"
|
||||
version = "12.6.3.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "python_full_version >= '3.12' or sys_platform != 'darwin'" },
|
||||
{ name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
|
||||
@@ -2697,7 +2697,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "remove-ai-watermarks"
|
||||
version = "0.12.1"
|
||||
version = "0.13.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "c2pa-python" },
|
||||
@@ -2759,6 +2759,11 @@ lama = [
|
||||
{ name = "onnxruntime", version = "1.24.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "onnxruntime", version = "1.27.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
]
|
||||
migan = [
|
||||
{ name = "huggingface-hub" },
|
||||
{ name = "onnxruntime", version = "1.24.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "onnxruntime", version = "1.27.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
|
||||
]
|
||||
trustmark = [
|
||||
{ name = "trustmark" },
|
||||
]
|
||||
@@ -2770,10 +2775,12 @@ requires-dist = [
|
||||
{ name = "click", specifier = ">=8.0.0" },
|
||||
{ name = "diffusers", marker = "extra == 'gpu'", specifier = ">=0.38.0" },
|
||||
{ name = "huggingface-hub", marker = "extra == 'lama'", specifier = ">=0.20.0" },
|
||||
{ name = "huggingface-hub", marker = "extra == 'migan'", specifier = ">=0.20.0" },
|
||||
{ name = "invisible-watermark", marker = "extra == 'detect'", specifier = ">=0.2.0" },
|
||||
{ name = "invisible-watermark", marker = "extra == 'dev'", specifier = ">=0.2.0" },
|
||||
{ name = "numpy", specifier = ">=1.24.0" },
|
||||
{ name = "onnxruntime", marker = "extra == 'lama'", specifier = ">=1.16.0" },
|
||||
{ name = "onnxruntime", marker = "extra == 'migan'", specifier = ">=1.16.0" },
|
||||
{ name = "opencv-python-headless", specifier = ">=4.8.0" },
|
||||
{ name = "piexif", specifier = ">=1.1.3" },
|
||||
{ name = "pillow", specifier = ">=10.0.0" },
|
||||
@@ -2782,7 +2789,7 @@ requires-dist = [
|
||||
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.1.0" },
|
||||
{ name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.5.0" },
|
||||
{ name = "python-dotenv", specifier = ">=1.0.0" },
|
||||
{ name = "remove-ai-watermarks", extras = ["gpu", "detect", "trustmark", "lama", "dev"], marker = "extra == 'all'" },
|
||||
{ name = "remove-ai-watermarks", extras = ["gpu", "detect", "trustmark", "lama", "migan", "dev"], marker = "extra == 'all'" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.4.0" },
|
||||
{ name = "safetensors", marker = "extra == 'gpu'" },
|
||||
{ name = "spandrel", marker = "extra == 'esrgan'", specifier = ">=0.3.0" },
|
||||
@@ -2793,7 +2800,7 @@ requires-dist = [
|
||||
{ name = "uv-outdated", marker = "python_full_version >= '3.12' and extra == 'dev'", specifier = ">=0.1.0" },
|
||||
{ name = "uv-secure", marker = "python_full_version >= '3.12' and extra == 'dev'", specifier = ">=0.12.0" },
|
||||
]
|
||||
provides-extras = ["gpu", "detect", "trustmark", "lama", "esrgan", "dev", "all"]
|
||||
provides-extras = ["gpu", "detect", "trustmark", "lama", "migan", "esrgan", "dev", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "requests"
|
||||
|
||||
Reference in New Issue
Block a user