CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-06 22:18:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

Report failed writes instead of crashing, and reject directories at parse time

Browse Source 
The Tier E adversarial sweep (new, scripts/robustness_suite.py) drove the real CLI
over truncated, corrupt, zero-byte, absurdly-shaped and bomb inputs, unicode and
RTL paths, hostile output directories and concurrent runs. It found two crashes;
the /simplify review then reproduced a third and worse one.

1. A FAILED WRITE CRASHED ON THE SIZE REPORT. image_io.imwrite is contractually
   non-raising and returns False, but write_bgr_with_alpha discarded that bool and
   returned None, so no caller could tell a failed write from a successful one.
   Every write site then ran output.stat() to print the size, so a read-only
   destination died with a bare FileNotFoundError pointing at the stat rather than
   the write. The fix is deliberately NOT uniform: single-image commands exit via
   the new cli._write_output_or_exit; api._write_visible_result RAISES so a library
   caller gets an accurate error instead of a confusing FileNotFoundError from the
   downstream metadata strip; and the batch sites raise but never SystemExit,
   because the batch loop counts per-image exceptions and aborting would kill the
   whole run.

2. BATCH LOST DATA SILENTLY. Into a read-only output directory it wrote ZERO files
   for 2 inputs and exited 0 -- no traceback, no error, an empty output directory a
   wrapping service would read as a completed run. The robustness harness could not
   see this class at all, since it scored exit codes and traceback markers and this
   failure has neither; it now asserts on the artifacts written.

3. A DIRECTORY PASSED AS THE IMAGE crashed the metadata scanner with
   IsADirectoryError, because click.Path(exists=True) accepts directories. Fixed
   with dir_okay=False on all six source arguments, so argument parsing refuses it.

Also adds Tier B4 (scripts/resource_ceilings.py): peak RSS per fill backend from
1 MP to 25 MP, one fresh process per cell. migan 603->775 MB and lama 4679->4779 MB
are flat in input size, confirming the crop-around-the-mask design and both
documented figures; cv2 is the only backend that grows (74->440 MB, 5.9x). The
harness's own no-op check originally allocated a full-frame temp before reading
peak RSS and inflated the numbers with input size -- it now compares only the mask
box, and the conclusion survived re-measurement.

And scripts/real_examples_e2e.py, which drives every command over real corpus
examples and checks the outcome rather than the exit code: 6/6 provenance classes
identified, 10/10 metadata strips re-scan clean, all three fill backends write,
diffusion on MPS writes genuinely changed images. It records samsung as a real
partial (the faintest mark, 0.431 -> 0.404 against a 0.40 gate on the weakest of
its 3 corpus positives) and treats the gated pill's refusal to act as correct.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-07-20 22:45:02 -07:00
co-authored by Claude Opus 4.8
parent c2afb92832
commit 633fc3aa52
10 changed files with 1150 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+3 -3
View File
@@ -17,7 +17,7 @@ Consequences for contributors (do not drift back into the stock niche just becau
## How to run
Per-command exit-code semantics (the no-signal / GPU-missing skip branches), test traps, and regression-guard paths live in `docs/module-internals.md` (section "CLI commands (`cli.py`)") — read it before changing any command's skip/exit behavior.
Per-command exit-code semantics (the no-signal / GPU-missing skip branches), test traps, and regression-guard paths live in `docs/module-internals.md` (section "CLI commands (`cli.py`)") — read it before changing any command's skip/exit behavior. Every single-image command's `source` argument declares `dir_okay=False`: without it `click.Path(exists=True)` accepts a directory, which then reached `open()` and raised `IsADirectoryError` (Tier E, 2026-07-20; `batch`'s `directory` already declared `file_okay=False`). Regression: `tests/test_cli_robustness.py::TestDirectoryInputIsRejected`.
- `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 `--backend auto|cv2|migan|lama` (default `auto`) and `--sensitivity auto|strict` (default `auto`) for the localize -> fill visible removal (see the `visible` bullet). 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` (**interacts with `--strength`**: diffusers derives its timesteps as `int(steps * strength)`, so a low `--steps` used to crash inside torch with `cannot reshape tensor of 0 elements` -- at the default 0.15 that was every value below 7. `noai/watermark_profiles.viable_steps` now raises the count to the minimum that denoises and logs the adjustment; keep the guard where it is, above `_generate_one`, so all three pipelines and the tiled path inherit it), `--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.
@@ -61,14 +61,14 @@ Compact map. The full per-module detail (design decisions, tuned thresholds, cal
- `_text_mark_engine.py` — shared base for the three text-mark engines (extracted 2026-06-09); the per-engine modules are config-only subclasses. Detection still matches the glyph silhouette (NCC, keys on glyph shape). The removal mask is TEMPLATE-FREE: it is the bounding box of the top-hat glyph blob (`extract_mask`), filled solid + dilated, so the shared fill inpaints the whole wordmark rectangle. This drops the fixed alpha-template placement, so a re-rendered or differently-placed mark is still masked; the captured alpha maps are now used only for the detection silhouette, not for removal. 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 fills cleanly while a match box leaves outline residue). `footprint_texture`/`footprint_is_flat` (median-Sobel over that box, `_FLAT_TEXTURE_MAX`) back the metadata-only safe-fill gate. Removal is the shared localize -> fill (MI-GAN/cv2). Detector precision is weak (~7% raw false-fire), so it is registry-gated in `remove_auto_marks` via `_keep_pill`: never on Doubao; the bottom-right wordmark removes it unrestricted (~94% precise, survives metadata-STRIPPED uploads); TC260-metadata-only removes it ONLY on a flat footprint (its textured false fires -- ceilings/walls -- are what the fill smears). 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). Detection matches the glyph silhouette (NCC); removal localizes the glyph blob to a solid dilated box (`extract_mask`) and hands it to the shared fill. Corpus validation: doubao and jimeng localize + remove at ~100% with clean footprints (the filled region blends into its surroundings within a few LAB levels, no color shift, no dark pit); clean images with no vendor signature had 0% false removal. **Samsung detection is calibrated only for the Italian "Contenuti generati dall'AI" string** (a pre-existing limit, unchanged by the localize -> fill refactor but now surfaced because detection gates removal): non-Italian Samsung locales are not detected, and thus not removed, even though the fill mask itself is locale-independent; other locales need their own detection silhouette (the locale string font-rendered + calibrated on real positives), NOT an app capture.
- `region_eraser.py` — universal region eraser (`erase` CLI) and the shared fill backend behind `watermark_registry.fill` for the visible localize -> fill removal. Three backends: `cv2` (default, no deps, the floor), `migan` (MI-GAN ONNX, extra `migan`, MIT, ~28 MB / ~0.19 s — the droplet-friendly tier, **the preferred default fill** when the extra is installed), `lama` (big-LaMa ONNX, extra `lama`, ~200 MB / ~4.7 GB peak — best quality, does not fit a minimal droplet, explicit opt-in only). Both `migan` and `lama` **crop a padded region around the mask** before inference and paste only masked pixels back, so peak RAM is bounded by the MARK size, not the image (`migan` ~0.6-0.9 GB regardless of upload size — feeding the whole frame scaled it to ~2.4 GB at 25 MP; `migan` feeds the crop at native resolution, `lama` resizes to its fixed 512²). **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. The `erase` command keeps its own `--backend`/`--inpaint-method` (unchanged).
- `region_eraser.py` — universal region eraser (`erase` CLI) and the shared fill backend behind `watermark_registry.fill` for the visible localize -> fill removal. Three backends: `cv2` (default, no deps, the floor), `migan` (MI-GAN ONNX, extra `migan`, MIT, ~28 MB / ~0.19 s — the droplet-friendly tier, **the preferred default fill** when the extra is installed), `lama` (big-LaMa ONNX, extra `lama`, ~200 MB / ~4.7 GB peak — best quality, does not fit a minimal droplet, explicit opt-in only). Both `migan` and `lama` **crop a padded region around the mask** before inference and paste only masked pixels back, so peak RAM is bounded by the MARK size, not the image (`migan` ~0.6-0.9 GB regardless of upload size — feeding the whole frame scaled it to ~2.4 GB at 25 MP; `migan` feeds the crop at native resolution, `lama` resizes to its fixed 512²). **Measured end to end 2026-07-20** (`scripts/resource_ceilings.py`, fresh process per cell, 1 MP → 25 MP): `migan` 603 → 775 MB and `lama` 4679 → 4779 MB, both **flat in input size** — the crop-around-the-mask design holds and both documented figures reproduce. **`cv2` is the only backend that GROWS with the input** (74 → 440 MB, 5.9x) because it inpaints the full frame rather than a crop; still the cheapest tier, but size it for the largest upload accepted. Cold wall time 0.02-0.12 s (cv2) / ~0.6 s (migan) / ~3.8 s (lama), model load included. (The harness's own no-op check originally allocated a full-frame temp before reading peak RSS and inflated these by up to 17% at 25 MP; it now compares only the mask box. The conclusion survived re-measurement, the digits moved.) **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. The `erase` command keeps its own `--backend`/`--inpaint-method` (unchanged).
- `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. **Both SDXL loaders (`_load_pipeline`, `_load_controlnet_pipeline`) pass `add_watermarker=False`** — diffusers otherwise embeds an open "Stable Diffusion XL" DWT-DCT invisible watermark on EVERY SDXL output whenever `invisible-watermark` is installed (the `detect` extra), so a watermark REMOVER would re-stamp a detectable AI watermark and the cleaned output re-reads as AI (`identify` → "Open invisible watermark: Stable Diffusion XL"; corpus-observed on the SynthID sample before the fix). Only the pipeline accepts the kwarg, not the `ControlNetModel` sub-model; qwen is not SDXL and has no such watermarker. Regression: `tests/test_platform.py::TestNoReembeddedWatermark`. 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).
- `noai/tiling.py` — sliding-window tiled diffusion for large inputs (CLI `--tile`). `WatermarkRemover.remove_watermark` branches to `run_tiled` when `tile` is set AND the long side exceeds `tile_size`, refactoring the single-pass `_generate` into a per-tile `_generate_one` (the ControlNet edge map is rebuilt per tile inside it). Pure helpers `plan_tiles` (uniform-size tiles, last one flush to the edge) and `feather_weights` (strictly-positive separable taper -> partition-of-unity blend) are unit-tested without the model. Also home to `feather_region_composite(base, regenerated, box, *, feather)` — the pure region-targeted compositor for **AI-enhanced composites** (`ai_source_kind == "enhanced"`): blends the regenerated AI box back over the original with a feathered seam, leaving the real photo OUTSIDE the box pixel-exact. It backs `WatermarkRemover.remove_watermark(region=...)` (regenerate ONLY the AI region, not the whole frame); the no-model lossless region path stays `region_eraser.erase`. New tile/region-blend tuning goes in these pure helpers; do not inline blend math into the runner.
- `auto_config.py` + the content-detection layer were REMOVED 2026-06-09; `--auto` is a deprecated no-op (controlnet is the default pipeline and the adaptive polish is ON by default and self-gates to a no-op where there is no detail deficit).
- `upscaler.py` — optional Real-ESRGAN pre-diffusion super-resolution for small inputs (extra `esrgan`, spandrel only). Manual opt-in; the default `--upscaler` stays `lanczos` and the engine always falls back to Lanczos on absence/error. ESRGAN can degrade faces and thin text.
- `image_io.py` — Unicode-safe cv2 IO (issue #17). Every cv2 file read/write in the package routes through `imread`/`imwrite`; do not call `cv2.imread`/`cv2.imwrite` directly. `to_bgr(image)` is the shared channel normalizer — use it instead of inlining `cvtColor` branches. `read_bgr_and_alpha`/`write_bgr_with_alpha` (+ `ALPHA_FORMATS`) are the alpha-preserving IO helpers shared by the CLI and the library `api` (moved here from cli so both use ONE implementation; the write MUST NOT zero alpha in the mark bbox — issue #30 white box). cv2/numpy import lazily, so importing `image_io` is cheap. **`imread` has a Pillow fallback (`_pil_read`) for HEIC/AVIF**: cv2 can't decode those containers, so when its decode returns None it opens via Pillow (AVIF native; HEIC via the core `pillow-heif` dep, whose libheif also covers AVIF) and converts to the same BGR/BGRA layout the flags imply — so the pixel/removal path reads iPhone HEIC and AVIF, not just the metadata path. Normal PNG/JPEG/WebP never reach the fallback. Corpus-verified: 54/55 HEIC+AVIF now decode (the 1 miss is a truncated upload). **`imwrite` PRESERVES the input format at max quality** ("work with originals"): the removal only touches the mark footprint (cv2 AND MI-GAN fills composite over the original — untouched pixels are bit-exact), so the container re-encode must not degrade the rest. JPEG is written at quality 100 / 4:4:4 (no chroma subsampling) — PSNR ~55 dB vs the old default-95's ~48; HEIC/AVIF write via Pillow (`_pil_write`) since cv2 has NO encoder for them (writing `.heic` via cv2 RAISES — a HEIC input used to crash on save). `imwrite` never raises (catches `cv2.error`). **`api.remove_visible` copies the original bytes verbatim on a no-op** (nothing removed + same output format) rather than a lossy re-encode, so a clean image round-trips byte-identical. `noai/constants.SUPPORTED_FORMATS` now includes `.heic`/`.heif`/`.avif` alongside png/jpg/jpeg/webp (pillow-heif is core, so read+write both work), so `batch` discovers them and the CLI no longer warns on an iPhone HEIC; JPEG-XL stays OUT (metadata/strip-only, no pixel decoder without pillow-jxl). **The invisible/SynthID path is inherently a full-frame diffusion regeneration (every pixel changes by design — you cannot "work with originals" there), but it no longer piles gratuitous re-encodes on top:** `watermark_remover` saves the regenerated output through `image_io.imwrite` (not raw `PIL.save`, which defaults to JPEG q75), `invisible_engine` writes its pre-diffusion temp as lossless PNG (not a re-compressed copy of a JPEG input), and the output metadata strip goes through the byte-level `metadata.remove_ai_metadata` (see its bullet) which for JPEG does NOT re-encode the DCT at all — pixels stay bit-identical.
- `image_io.py` — Unicode-safe cv2 IO (issue #17). Every cv2 file read/write in the package routes through `imread`/`imwrite`; do not call `cv2.imread`/`cv2.imwrite` directly. `to_bgr(image)` is the shared channel normalizer — use it instead of inlining `cvtColor` branches. `read_bgr_and_alpha`/`write_bgr_with_alpha` (+ `ALPHA_FORMATS`) are the alpha-preserving IO helpers shared by the CLI and the library `api` (moved here from cli so both use ONE implementation; the write MUST NOT zero alpha in the mark bbox — issue #30 white box). **`write_bgr_with_alpha` RETURNS `imwrite`'s success flag and every caller must check it** — `imwrite` is contractually non-raising, so that bool is the ONLY signal the file was not created. The wrapper used to return `None` and swallow it, so a write that a read-only directory silently prevented ran on to `output.stat()` and died with a bare `FileNotFoundError` traceback pointing at the stat rather than the write (Tier E, 2026-07-20). **Three different layers each needed their own handling, and the right one is NOT the same everywhere:** the single-image commands (`visible --mark <name>`, `erase`, `all`) write through the shared `cli._write_output_or_exit`, which exits; `api._write_visible_result` (behind `visible --mark auto`) RAISES `OSError` so a library caller gets an accurate error instead of a confusing `FileNotFoundError` from the downstream metadata strip; and the **batch** sites raise too, never `SystemExit` — the batch loop catches per-image exceptions, counts them and exits non-zero, so aborting there would kill the whole run. Discarding the flag in batch made a read-only output directory produce **zero files and still exit 0** — silent data loss contradicting the documented batch contract. Regression: `tests/test_cli_robustness.py`. cv2/numpy import lazily, so importing `image_io` is cheap. **`imread` has a Pillow fallback (`_pil_read`) for HEIC/AVIF**: cv2 can't decode those containers, so when its decode returns None it opens via Pillow (AVIF native; HEIC via the core `pillow-heif` dep, whose libheif also covers AVIF) and converts to the same BGR/BGRA layout the flags imply — so the pixel/removal path reads iPhone HEIC and AVIF, not just the metadata path. Normal PNG/JPEG/WebP never reach the fallback. Corpus-verified: 54/55 HEIC+AVIF now decode (the 1 miss is a truncated upload). **`imwrite` PRESERVES the input format at max quality** ("work with originals"): the removal only touches the mark footprint (cv2 AND MI-GAN fills composite over the original — untouched pixels are bit-exact), so the container re-encode must not degrade the rest. JPEG is written at quality 100 / 4:4:4 (no chroma subsampling) — PSNR ~55 dB vs the old default-95's ~48; HEIC/AVIF write via Pillow (`_pil_write`) since cv2 has NO encoder for them (writing `.heic` via cv2 RAISES — a HEIC input used to crash on save). `imwrite` never raises (catches `cv2.error`). **`api.remove_visible` copies the original bytes verbatim on a no-op** (nothing removed + same output format) rather than a lossy re-encode, so a clean image round-trips byte-identical. `noai/constants.SUPPORTED_FORMATS` now includes `.heic`/`.heif`/`.avif` alongside png/jpg/jpeg/webp (pillow-heif is core, so read+write both work), so `batch` discovers them and the CLI no longer warns on an iPhone HEIC; JPEG-XL stays OUT (metadata/strip-only, no pixel decoder without pillow-jxl). **The invisible/SynthID path is inherently a full-frame diffusion regeneration (every pixel changes by design — you cannot "work with originals" there), but it no longer piles gratuitous re-encodes on top:** `watermark_remover` saves the regenerated output through `image_io.imwrite` (not raw `PIL.save`, which defaults to JPEG q75), `invisible_engine` writes its pre-diffusion temp as lossless PNG (not a re-compressed copy of a JPEG input), and the output metadata strip goes through the byte-level `metadata.remove_ai_metadata` (see its bullet) which for JPEG does NOT re-encode the DCT at all — pixels stay bit-identical.
- `api.py` — the high-level convenience API, re-exported lazily at the package top level via `__init__.__getattr__` (PEP 562, so `import remove_ai_watermarks` stays cheap): `remove_visible(source, output=None, *, sensitivity="auto", backend="auto", strip_metadata=True, write_noop=True) -> (result_bgr, [labels])` (source = path OR BGR ndarray; a PATH auto-reads metadata provenance and preserves alpha, an ARRAY does neither; `write_noop=True` writes a clean passthrough copy when nothing is removed, `False` leaves `output` untouched so a "no mark = produce nothing" caller like the CLI `visible` command does not clobber a pre-existing file there) and `visible_provenance(path) -> frozenset[str]` (the single metadata→vendor-keys mapper; `cli._visible_provenance` is a thin None-guarded wrapper over it). **`remove_visible` is the ONE path the CLI and library share** — `cli.cmd_visible`'s `--mark auto` branch delegates entirely to it (read → provenance → `remove_auto_marks` → write → `strip_metadata`), so there is no CLI-vs-library drift; `strip_metadata` defaults True to match `visible --strip-metadata`. This is where a library caller should start — NOT the engines directly (`GeminiEngine`/`TextMarkEngine` have no `remove_watermark` any more; removal is registry `remove_auto_marks`/`KnownMark.remove`; the old single-strongest `best_auto_mark` is gone — removal takes EVERY mark). `identify` is NOT top-level re-exported (it collides with the `identify` submodule); use `from remove_ai_watermarks.identify import identify`.
For the Doubao alpha-distillation history (why content-image reverse-alpha distillation fails by physics and controlled captures were required), see `docs/research-doubao-distillation.md`.
docs/module-internals.md
+3 -1
View File
@@ -189,7 +189,7 @@ The factor is now a per-mark `TextMarkConfig.provenance_ncc_factor`. Doubao stay
**`cli._remove_visible_auto` is the shared visible-removal helper used by `cmd_all`/`cmd_batch` too** (they no longer hardcode `GeminiEngine`), so `all`/`batch` remove Doubao/Jimeng/Samsung text marks, not just the Gemini sparkle (regression-guarded by `test_all_visible_step_uses_registry`). The three text-mark adapters were consolidated 2026-06-09: a single `_text_mark(key, label, location)` builds the registry row from one parameterized `_text_mark_detect`/`_text_mark_remove` pair (the remove adapter localizes the glyph footprint and hands it to the shared `fill` only when detected/forced, else skipped); the gemini adapters stay bespoke. Add a new visible mark = one `_text_mark(...)` row + its `TextMarkConfig` (with a captured alpha map for the detection silhouette); do not re-add per-mark `if` branches or copy-paste adapters.
**Alpha-on-save policy (issue #30):** `cli._write_bgr_with_alpha` rejoins the input's alpha plane **unchanged** — it must NOT zero alpha in the watermark bbox. The fill reconstructs real pixels there, so zeroing alpha punched a transparent hole that renders as a solid **white box** on any non-transparent viewer (Gemini app exports are opaque RGBA, so every user hit it; regression-guarded by `test_visible_keeps_alpha_opaque_in_watermark_region`). The registry `remove()` still returns its region, but the CLI no longer uses it to clear alpha.
**Alpha-on-save policy (issue #30):** `image_io.write_bgr_with_alpha` (it lives in `image_io`, not `cli` — moved so the CLI and the library `api` share ONE implementation) rejoins the input's alpha plane **unchanged** — it must NOT zero alpha in the watermark bbox. The fill reconstructs real pixels there, so zeroing alpha punched a transparent hole that renders as a solid **white box** on any non-transparent viewer (Gemini app exports are opaque RGBA, so every user hit it; regression-guarded by `test_visible_keeps_alpha_opaque_in_watermark_region`). The registry `remove()` still returns its region, but the CLI no longer uses it to clear alpha. **It returns `imwrite`'s success flag and callers must check it** (2026-07-20): `imwrite` is contractually non-raising, so that bool is the only signal the file was not created. The wrapper previously returned `None` and swallowed it, so every CLI write site ran `output.stat()` to report the size and a read-only destination died with a bare `FileNotFoundError` traceback pointing at the stat instead of the write. The CLI now writes through the shared `cli._write_output_or_exit`. Regression: `tests/test_cli_robustness.py::TestFailedWriteIsReported`.
## `gemini_engine.py`
@@ -338,6 +338,8 @@ KEPT from that work (independently valid for the manual `--pipeline qwen`): the
Full per-command behavior for the skip/exit branches summarized in `CLAUDE.md`'s "How to run". The CLI distinguishes three exit codes: success (0), hard error (1), and a "nothing to do" code (2, `EXIT_NO_VISIBLE_MARK` / `EXIT_NO_INVISIBLE_SIGNAL`) so a wrapping service (raiw.cc) can surface guidance instead of treating an unchanged image as done (the production "it didn't work" / score-0 trap).
**Every single-image command's `source` argument declares `dir_okay=False`** (2026-07-20). `click.Path(exists=True)` accepts a directory unless told otherwise, so `identify <dir>` sailed past argument parsing and raised `IsADirectoryError` out of `metadata.scan_head`'s `open()` — a traceback, not a usage error. Refusing it at the argument layer is the right place: every command gets it, and none needs its own check. (`batch`'s `directory` argument was already correct with `file_okay=False`.) Found by the Tier E adversarial sweep; regression: `tests/test_cli_robustness.py::TestDirectoryInputIsRejected`.
### `all`
Full pipeline (visible + invisible + metadata). Same diffusion knobs as `invisible`, plus the visible-pass `--backend auto|cv2|migan|lama` (default `auto`) that picks the fill for the localize -> fill visible removal. **When the `[gpu]` extra is absent, step 2 (invisible/SynthID) is skipped** — `all` still writes an output (visible mark + metadata stripped) but prints a prominent end-of-run banner ("the invisible (SynthID) watermark was NOT removed") AND exits **non-zero** (1), so a skipped SynthID pass is not mistaken for a clean result (the recurring #14/#47 trap, where the old quiet inline warning was missed). `invisible` already hard-errors without the extra; only `all` continued, hence the loud end-banner. Regression-guarded by `tests/test_cli.py::TestAllCommand::test_all_loud_warning_and_nonzero_exit_when_gpu_missing`. **No-signal skip (P0#5):** step 2 also runs the same `has_invisible_target` gate (see `invisible` below) — when no invisible watermark is detectable and `--force` is not set, step 2 is skipped and the pixels are left intact, but unlike the GPU-missing skip this is a **SUCCESS (exit 0)**: the visible pass + metadata strip still ran and a file is written (the message says so without claiming the image is clean). Distinct exit semantics by design: GPU-missing = couldn't do the work (non-zero); no-signal = nothing to do (zero). Regression-guarded by `test_all_skips_invisible_on_no_signal_but_succeeds`. **Test trap:** any `all` test that exercises the full pipeline MUST `patch("remove_ai_watermarks.invisible_engine.is_available", return_value=True)` — CI installs core+dev only (no `[gpu]`), so an unpatched `all` test takes the skip branch and now hits the non-zero exit. This passed locally (gpu present → `is_available()` True) but red-failed every matrix cell on the v0.11.0 commit (`test_all_basic`/`test_all_visible_step_uses_registry` asserted exit 0); both now patch `is_available` True.
docs/verification-plan.md
+146 -4
View File
@@ -579,6 +579,145 @@ bug (`scale_basis`) directly.
This section records what the measurements imply technically. Prioritization is tracked
separately, outside this repo.
## Real-example end-to-end run (2026-07-20)
The corpus sweeps all call the library IN-PROCESS. This run drives the actual
`remove-ai-watermarks` entry point as a user would, over real corpus images and the
committed real fixtures, and checks the OUTCOME rather than the exit code
(`scripts/real_examples_e2e.py`). **26 of 28 behaviours passed.**
| Command | Real examples | Result |
|---|---|---|
| `identify --json` | 6 provenance classes (OpenAI, Adobe, Midjourney, Doubao/TC260, Grok, FLUX) | 6/6 correct verdict + platform |
| `metadata --check` / `--remove` | the same 5 real AI-metadata files | 10/10 detected, and the output re-scans clean |
| `visible --mark auto` | a live positive per mark | doubao / jimeng / gemini removed and re-detect clean; pill correctly DECLINED by the gate; samsung partial (below) |
| `erase --region` | one real image x 3 backends | cv2 / MI-GAN / big-LaMa all wrote output |
| `batch --mode visible` | a real 5-image directory | 5/5 outputs |
| `invisible` (MPS, `--max-resolution 512`) | a real Gemini and a real OpenAI carrier | both wrote a genuinely CHANGED image |
| `all` (MPS) | a real Gemini carrier | one transient exit 1, not reproducible (see below) |
**Samsung is a real, reproducible partial.** It is the faintest registered mark (peak alpha
~0.38) sitting on a 0.40 gate, so the margin between "detected" and "removed" is razor thin.
Over the entire corpus population (n=3, all of it) the CLI clears 2 of 3 outright
(0.446 and 0.440 -> below gate) and on the weakest one reduces 0.431 -> **0.404**, which is
still fractionally over the gate. The glyph IS filled and the confidence IS reduced; the
residual re-detects. This is the faint-mark residual class on the `binary` front-end, which
has no equivalent of the `tophat` faint-mask fallback. **Not fixed:** with n=3 corpus-wide
there is no way to tell an improvement from noise, which is the same
improving-what-3-samples-measure trap recorded under Open items.
**The pill "failure" was the harness, not the product.** `detect_marks` fires `jimeng_pill`
but `remove_auto_marks` returns no label -- `_keep_pill` correctly declines an uncorroborated
low-confidence pill, so `visible` writes nothing and exits 2. The harness now asks the
product what it DECIDED and treats a correct decline as a pass.
**The `all` exit 1 did not reproduce** -- the same invocation ran clean standalone twice and
again in the exact three-run sequence that produced it (all exit 0, step 2 executed, no skip
banner). It stays recorded as a transient because it could not be diagnosed: the harness
discarded the command output, and `cmd_all` has three distinct `SystemExit(1)` paths
(unreadable input, unreadable intermediate, and the deliberate synthid-skipped banner) that
an exit code alone cannot distinguish. The harness now retains the output tail on any
failure, so a recurrence is identifiable.
## Tier E: robustness (2026-07-20) -- RUN, and it found two real crashes
`scripts/robustness_suite.py` drives the real CLI over adversarial and degenerate inputs and
scores GRACEFULNESS, not success: a non-zero exit with a readable message is a pass, an
unhandled traceback or a hang is a fail. **33/33 graceful after two fixes; 31/33 before.**
Covered: truncated and corrupt files, zero-byte, a text file named `.jpg`, 1x1 and
1x4000 slivers, a 729 KB decompression bomb declaring 16000x16000, Unicode and RTL
filenames (read AND write), a mismatched extension, a nonexistent nested output dir, a
read-only output dir, a directory passed as a file, a missing path, 4x concurrent runs
against one input, and batch over both an empty and an undecodable directory.
**Both defects were invisible to the 849-test suite**, because unit tests feed well-formed
fixtures into writable directories. Both are now regression-guarded by
`tests/test_cli_robustness.py`.
1. **A failed write crashed on the size report.** `image_io.imwrite` is contractually
non-raising -- it returns `False` when the path cannot be written. But
`write_bgr_with_alpha` discarded that bool and returned `None`, so **no caller could
distinguish a failed write from a successful one**, and all five write sites then ran
`output.stat()` to print the size. A read-only output directory produced a bare
`FileNotFoundError` traceback pointing at the stat rather than at the write. The signal
existed the whole way down and was thrown away by one wrapper. Fixed at that wrapper
(propagate the flag) plus one shared `cli._write_output_or_exit`, so all sites are
covered rather than the one that happened to be caught.
2. **A directory passed as the image crashed the scanner.** `click.Path(exists=True)`
accepts directories unless told otherwise, so `identify <dir>` reached `open()` and
raised `IsADirectoryError`. Fixed by `dir_okay=False` on all six `source` arguments --
argument parsing now refuses it, which is where it belongs. (`batch` was already correct:
it declares `file_okay=False`.)
3. **The worst instance was found by the /simplify review, not by the sweep: `batch` into a
read-only directory wrote ZERO files for 2 inputs and exited 0.** No traceback, no
error, a success code and an empty output directory -- a wrapping service would treat
that as a completed run. `graceful()` structurally cannot see this class, because it
scores exit code and traceback markers and this failure has neither. The suite now
carries a `batch_readonly_outdir` case that asserts on the ARTIFACTS (how many files
exist) rather than on the status. **Any check for a silent no-op has to assert on the
output, not the exit code.**
The fix is NOT uniform, and that matters: the single-image commands exit via
`cli._write_output_or_exit`; `api._write_visible_result` RAISES so a library caller gets an
accurate error rather than a confusing `FileNotFoundError` from the downstream metadata
strip; and the batch sites raise too but must never `SystemExit`, because the batch loop
counts per-image exceptions and aborting would kill the whole run instead of failing one
image.
The lesson worth keeping: **a non-raising IO contract needs its return value checked at
every call site, and a wrapper that swallows it silently disables the contract for
everyone downstream.** Grep for other `-> None` wrappers over non-raising primitives --
`invisible_engine.py:346` still discards `imwrite`'s flag and is the same shape.
## Tier B4: resource ceilings (2026-07-20) -- RUN
`scripts/resource_ceilings.py`, one FRESH process per cell (peak RSS inside a long-lived
process is contaminated by whatever ran before it, and the question is what a per-request
worker needs). The mask is a fixed, small corner region at every input size, so the thing
under test is whether the learned backends really crop around the MARK rather than
processing the frame.
| backend | peak RSS 1MP -> 25MP | wall | scaling |
|---|---|---|---|
| cv2 | 74 -> 440 MB | 0.02-0.12 s | **grows 5.9x with input size** |
| migan | 603 -> 775 MB | ~0.6 s | flat |
| lama | 4679 -> 4779 MB | ~3.8 s | flat |
**The documented claims hold.** CLAUDE.md's "migan ~0.6-0.9 GB regardless of upload size"
and "lama ~4.7 GB peak" both reproduce to the digit, and the crop-around-the-mask design is
confirmed: both learned backends are flat in input size. This is also the measurement
behind the deployment split (free tier `migan` fits a small droplet under 1 GB; paid tier
`lama` needs ~5 GB).
**New information: cv2 is the only backend that scales with the input.** It inpaints the
full frame rather than a crop, so at 25 MP it costs 440 MB -- still the cheapest tier, but
6x its small-image footprint, which is worth knowing when sizing a worker that accepts
phone-camera uploads. Its wall time stays trivial throughout.
Caveat on the wall times: each is a COLD process including model load, so migan's ~0.6 s
here is not comparable to the ~0.19 s warm figure quoted elsewhere. Cold is the honest
number for a per-request worker; warm is the honest number for a long-lived one.
**Two method notes, both cases of the harness corrupting its own measurement.**
1. The first version called `erase()` with a positional `boxes` argument, which the
keyword-only signature rejects. All 12 cells failed identically and still reported
plausible-looking RSS (60-129 MB) -- the process's numpy footprint with no backend work
at all. Twelve identical failures were one bad call, and only the printed error made
that visible. The child now asserts the output actually differs, so a no-op cannot be
reported as a measurement.
2. That assertion was then itself the contaminant: `(out != img).any()` allocates a boolean
temp the size of the IMAGE before `getrusage` is read -- ~75 MB at 25 MP, and it grows
with the input, so the harness partly manufactured the very "cv2 scales with input size"
conclusion it was measuring. Now it compares only the mask box. **Re-measured: the
conclusion survived, the digits moved.** cv2 6.1x -> 5.9x, migan max 847 -> 775 MB, lama
max 4849 -> 4779 MB; the table above is the corrected run. Worth keeping because the
contaminated numbers were already committed to CLAUDE.md before the check was
questioned -- a verification step is part of the instrument and needs the same scrutiny
as the thing it verifies.
## Open items (as of 2026-07-20)
Everything below is known, measured, and deliberately not done yet. Each line says what it
@@ -592,6 +731,7 @@ would take, so none of it has to be rediscovered.
| 2 | Exit code 2 means three different things (no visible mark / no invisible signal / Click usage error) | any wrapper must parse stderr to tell them apart | split the codes; **breaking for existing wrappers**, so it needs a deliberate call |
| 3a | ~~the full visible-parity sweep has NOT been re-run since the doubao front-end fix~~ **DONE 2026-07-20** | doubao parity moved **91.8% -> 99.3%** (2562/2580) as predicted; gemini/jimeng/samsung unchanged, `_visible_parity_cv2_v2.csv` | closed |
| 3 | `visible_removal_audit.py` measures the UNGATED per-mark path | reports the pill at 32% where the product runs at 100% precision | teach it the product path (`remove_auto_marks`) for gated marks, or at minimum say so loudly in its docstring |
| 5 | samsung leaves a residual just over its gate on the weakest of its 3 corpus positives | 0.431 -> 0.404 against a 0.40 gate; the other two clear outright | a faint-mask fallback for the `binary` front-end (samsung has none), but **do not tune on n=3** -- it needs more positives first, same blocker as everything else in Tier C |
| 4 | the faint-mask fallback fills the WHOLE corner box, not the glyph | 12 of 12 real faint-path frames covered 100% of the corner ROI (120.9% median with padding); the path fires on ~8% of doubao detections | fixed 2026-07-20 -- see below |
**Defect 4 in full, because it is instructive.** The fallback I added on 2026-07-19 reads
@@ -675,10 +815,12 @@ Measured this session, in the order the evidence supports:
### Verification tiers not run
- **B4 resource ceilings** -- peak RSS and wall time per backend x input size to 25 MP.
- **E robustness** -- truncated, corrupt, absurd dimensions, decompression bombs, unicode
and RTL filenames, read-only output dirs, concurrent runs on one file.
- **C recall expansion** -- gated by labelling appetite.
- **C recall expansion** -- gated by labelling appetite. This is now the ONLY unrun tier,
and it is blocked on data rather than effort: every remaining detector question needs
labelled positives (30+ per uncovered vendor; jimeng still rests on n=14, the pill on
n=6, samsung on n=3).
(Tiers E and B4 are now RUN -- see the sections above.)
### Recommended next step
scripts/real_examples_e2e.py
+328
View File
@@ -0,0 +1,328 @@
"""End-to-end confidence run: drive the ACTUAL CLI over REAL corpus examples.
WHY THIS EXISTS AND WHAT IT IS NOT
The 849-test suite and `smoke_matrix.py` prove the code paths behave on fixtures and
synthetic inputs. This is the other half: run the real `remove-ai-watermarks` entry point,
as a user would, over real corpus images spanning every command and every provenance
class, and CHECK THE OUTPUT -- not that it exited 0, but that it did the right thing (the
mark is actually gone on re-detect, the metadata actually strips, the diffusion actually
writes a changed image). A green exit is not evidence the work happened.
WHAT IT COVERS
identify one real image per provenance class -> the verdict is right
metadata real AI-metadata files -> --check detects, --remove strip-and-verifies clean
visible real marked images per mark -> the mark is gone on re-detect, output written
erase a real image, each fill backend (cv2 / migan / lama) -> output written
invisible a real SynthID image on MPS at reduced resolution -> a CHANGED image is written
all a real marked image through the full pipeline -> output written
batch a real directory -> every input produces an output
invisible/all run the diffusion model, so they are gated behind --diffusion and run at a
small --max-resolution on MPS (the user's "reduced size on MPS" path). Everything else is
cv2/numpy and fast.
DATA SAFETY
Corpus images are user uploads: read-only, local analysis, outputs to a gitignored temp
dir. Records example uids and pass/fail, never image content.
uv run python scripts/real_examples_e2e.py # fast surface (no diffusion)
uv run python scripts/real_examples_e2e.py --diffusion # + invisible/all on MPS
"""
from __future__ import annotations
import argparse
import glob
import json
import random
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
REPO = Path(__file__).resolve().parents[1]
CORPUS = REPO / "data" / "spaces" / "originals"
DATASETS = REPO / "data" / "spaces" / "_visible_datasets"
SAMPLES = REPO / "data" / "samples"
_UV = shutil.which("uv") or "uv" # full path avoids the partial-executable lint
def run(args: list[str], timeout: int = 300) -> tuple[int, str]:
"""Invoke the real installed CLI. Returns (exit_code, combined output)."""
proc = subprocess.run( # noqa: S603 -- fixed argv, the whole point is the real entry point
[_UV, "run", "remove-ai-watermarks", *args],
cwd=REPO,
capture_output=True,
text=True,
timeout=timeout,
check=False,
)
return proc.returncode, (proc.stdout + proc.stderr)
def find_visible_positive(mark: str) -> Path | None:
"""A real corpus image the parity run bucketed as carrying this mark, that the current
detector STILL fires on (the bucket was built by an older run; re-confirm live)."""
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.watermark_registry import detect_marks
pool = sorted(glob.glob(str(DATASETS / mark / "*")))
random.Random(3).shuffle(pool) # noqa: S311 -- deterministic sampling, not cryptography
for p in pool[:60]:
img = imread(p)
if img is None:
continue
if any(d.detected and d.key == mark for d in detect_marks(img)):
return Path(p)
return None
class Results:
"""Collects one row per checked behaviour.
A FAIL keeps the command's OUTPUT. That is not cosmetic: the first run of this harness
discarded it, an `all` invocation failed once with exit 1, and because the output was
gone there was no way to tell which of that command's three distinct `SystemExit(1)`
paths had fired -- the failure was undiagnosable and did not reproduce. Keep the tail
so a transient is at least identifiable after the fact.
"""
def __init__(self) -> None:
self.rows: list[tuple[str, str, bool, str, str]] = []
def add(self, cmd: str, example: str, ok: bool, detail: str, output: str = "") -> None:
self.rows.append((cmd, example, ok, detail, "" if ok else output[-800:]))
mark = "PASS" if ok else "FAIL"
print(f" [{mark}] {cmd:28s} {example:26s} {detail}", flush=True)
def report(self) -> int:
n = len(self.rows)
bad = [r for r in self.rows if not r[2]]
print(f"\n{'=' * 78}\nREAL-EXAMPLE E2E {n - len(bad)}/{n} passed")
print(f"{'=' * 78}")
for cmd, ex, ok, detail, out in self.rows:
if not ok:
print(f" FAIL {cmd} {ex} {detail}")
if out.strip():
print(" --- command output (tail) ---")
for line in out.strip().splitlines()[-12:]:
print(f" {line}")
if not bad:
print(" every command produced the right result on real corpus examples")
return 1 if bad else 0
def check_identify(res: Results) -> None:
"""One real image per provenance class -> the verdict is right (--json must parse)."""
print("\nidentify -- real image per provenance class")
cases = [
("chatgpt-1.png", True),
("firefly-1.png", True),
("mj-1.png", True),
("doubao-1.png", True),
("grok-1.jpg", True),
("flux-1.jpg", True),
]
for name, want_ai in cases:
p = SAMPLES / name
if not p.exists():
res.add("identify", name, False, "sample missing")
continue
code, out = run(["identify", str(p), "--json"])
try:
data = json.loads(out[out.index("{") : out.rindex("}") + 1])
except (ValueError, json.JSONDecodeError):
res.add("identify", name, False, f"json did not parse (exit {code})")
continue
is_ai = data.get("is_ai_generated")
plat = (data.get("platform") or "").lower()
ok = bool(is_ai) is want_ai
res.add("identify", name, ok, f"is_ai={is_ai} platform={plat or '-'}")
def check_metadata(res: Results, tmp: Path) -> None:
"""Real AI-metadata files -> --check detects, --remove strip-and-verifies clean."""
print("\nmetadata -- real AI-metadata files, detect then strip-and-verify")
for name in ("chatgpt-1.png", "doubao-1.png", "mj-1.png", "grok-1.jpg", "flux-1.jpg"):
p = SAMPLES / name
if not p.exists():
res.add("metadata --check", name, False, "sample missing")
continue
code, out = run(["metadata", str(p), "--check"])
detected = "no ai" not in out.lower() and code == 0
res.add("metadata --check", name, detected, "AI metadata detected" if detected else "nothing detected")
outp = tmp / f"stripped_{name}"
code, out = run(["metadata", str(p), "--remove", "-o", str(outp)])
# strip-and-verify: the command re-scans the OUTPUT and fails loudly on leftovers.
clean = outp.exists() and code == 0 and "still" not in out.lower()
res.add("metadata --remove", name, clean, "output re-scanned clean" if clean else f"exit {code}")
def check_visible(res: Results, tmp: Path) -> None:
"""Real marked images -> the PRODUCT'S DECISION is honoured, and a removed mark clears.
The success criterion is not "the mark is always gone" -- it is "the product did what it
decided, and the decision is right". Two designed behaviours make a blind re-detect
misleading:
* The pill is GATED (`_keep_pill`): a low-confidence pill with no corroboration is
deliberately NOT removed, so `visible` correctly writes nothing and exits 2. That is
the gate working, not a miss -- checked by asking `remove_auto_marks` whether it
chose to act.
* Samsung is the faintest mark (peak alpha ~0.38) at a razor-thin 0.40 gate, so a
borderline positive can be reduced yet re-detect just above threshold. Reported as
the measured before->after confidence, not a bare pass/fail.
"""
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.watermark_registry import detect_marks, get_mark, remove_auto_marks
print("\nvisible --mark auto -- real marked image per mark, product decision then re-detect")
for mark in ("doubao", "jimeng", "gemini", "samsung", "jimeng_pill"):
src = find_visible_positive(mark)
if src is None:
res.add("visible", mark, True, "no live positive in bucket (skipped, not a failure)")
continue
img = imread(str(src))
conf_before = next((d.confidence for d in detect_marks(img) if d.detected and d.key == mark), 0.0)
# What does the product DECIDE to do? (labels lists the marks it removed.)
_out_img, labels = remove_auto_marks(img, sensitivity="auto", provenance=frozenset(), backend="cv2")
acted = get_mark(mark).label in labels
outp = tmp / f"visible_{mark}{src.suffix}"
code, out = run(["visible", str(src), "--mark", "auto", "-o", str(outp)])
if not acted:
# The gate declined this mark. The CLI must then write nothing and exit 2 --
# that is the correct outcome, so verify the CLI agrees with the decision.
ok = code == 2 and not outp.exists()
res.add(
"visible", mark, ok, f"gate declined (conf {conf_before:.2f}); CLI exit {code}, no output -- correct"
)
continue
if not outp.exists():
res.add("visible", mark, False, f"product removed it but CLI wrote no output (exit {code})", out)
continue
cleaned = imread(str(outp))
conf_after = next((d.confidence for d in detect_marks(cleaned) if d.detected and d.key == mark), 0.0)
clean = conf_after == 0.0
if clean:
res.add("visible", mark, True, f"removed, re-detect clean ({conf_before:.2f} -> below gate)")
else:
# Reduced but still over the gate: a real residual. Honest partial, flagged.
res.add(
"visible",
mark,
False,
f"reduced {conf_before:.2f} -> {conf_after:.2f} but still over gate (faint-mark residual)",
)
def check_erase(res: Results, tmp: Path) -> None:
"""A real image, each fill backend actually runs and writes an output."""
from remove_ai_watermarks.region_eraser import lama_available, migan_available
print("\nerase --region -- each fill backend on a real image")
src = next((Path(p) for p in sorted(glob.glob(str(CORPUS / "*" / "*"))) if Path(p).stat().st_size > 50_000), None)
if src is None:
res.add("erase", "-", False, "no corpus image")
return
backends = ["cv2"] + (["migan"] if migan_available() else []) + (["lama"] if lama_available() else [])
for backend in backends:
outp = tmp / f"erase_{backend}{src.suffix}"
code, out = run(["erase", str(src), "--region", "10,10,120,60", "--backend", backend, "-o", str(outp)])
ok = outp.exists() and code == 0 and outp.stat().st_size > 0
res.add(f"erase --backend {backend}", src.name[:12], ok, "output written" if ok else f"exit {code}", out)
def check_diffusion(res: Results, tmp: Path, gemini_src: str, openai_src: str) -> None:
"""The GPU path: invisible + all on MPS at reduced resolution -> a CHANGED image."""
import numpy as np
from remove_ai_watermarks.image_io import imread
print("\ninvisible / all -- real SynthID image on MPS, reduced resolution")
for label, src in (("invisible/gemini", gemini_src), ("invisible/openai", openai_src)):
if not src:
res.add(label, "-", True, "no real positive found (skipped)")
continue
sp = Path(src)
outp = tmp / f"inv_{sp.stem}.png"
code, out = run(
["invisible", src, "-o", str(outp), "--device", "mps", "--max-resolution", "512", "--seed", "0"],
timeout=1200,
)
if not outp.exists():
res.add(label, sp.name[:12], False, f"no output (exit {code})", out)
continue
before, after = imread(src), imread(str(outp))
# Diffusion regenerates every pixel; the output must actually differ from the input.
# A DIFFERENT SHAPE is itself proof it changed (the pipeline resizes), so treat it
# as changed rather than comparing arrays that cannot be compared. The first
# version wrote `array_equal(after, before if shapes match else after)`, which
# compares `after` WITH ITSELF on the mismatch branch and is therefore always
# "unchanged" -- it would have reported a genuinely resized output as a no-op.
changed = (
after is not None
and before is not None
and (before.shape != after.shape or not np.array_equal(before, after))
)
res.add(label, sp.name[:12], changed, "diffusion wrote a changed image" if changed else "output == input", out)
if gemini_src:
sp = Path(gemini_src)
outp = tmp / f"all_{sp.stem}.png"
code, out = run(
["all", gemini_src, "-o", str(outp), "--device", "mps", "--max-resolution", "512", "--seed", "0"],
timeout=1200,
)
ok = outp.exists() and outp.stat().st_size > 0
res.add("all (full pipeline)", sp.name[:12], ok, "output written" if ok else f"exit {code}", out)
def check_batch(res: Results, tmp: Path) -> None:
"""A real directory -> every supported input produces an output."""
print("\nbatch -- a real directory")
indir = tmp / "batch_in"
indir.mkdir(exist_ok=True)
picks = sorted(glob.glob(str(DATASETS / "doubao" / "*")))[:5]
for p in picks:
shutil.copy2(p, indir / Path(p).name)
n_in = len(list(indir.glob("*")))
if n_in == 0:
res.add("batch", "-", False, "no inputs to seed")
return
outdir = tmp / "batch_out"
code, out = run(["batch", str(indir), "-o", str(outdir), "--mode", "visible"], timeout=600)
n_out = len(list(outdir.glob("*"))) if outdir.exists() else 0
ok = code == 0 and n_out >= n_in
res.add("batch --mode visible", f"{n_in} imgs", ok, f"{n_out}/{n_in} outputs (exit {code})", out)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--diffusion", action="store_true", help="also run invisible/all on MPS (slow)")
ap.add_argument("--gemini", default="")
ap.add_argument("--openai", default="")
a = ap.parse_args()
res = Results()
with tempfile.TemporaryDirectory(prefix="raiw_e2e_") as td:
tmp = Path(td)
check_identify(res)
check_metadata(res, tmp)
check_visible(res, tmp)
check_erase(res, tmp)
check_batch(res, tmp)
if a.diffusion:
check_diffusion(res, tmp, a.gemini, a.openai)
else:
print("\n(diffusion skipped -- pass --diffusion to run invisible/all on MPS)")
raise SystemExit(res.report())
if __name__ == "__main__":
main()
scripts/resource_ceilings.py
+156
View File
@@ -0,0 +1,156 @@
"""Tier B4: peak RSS and wall time per fill backend across input sizes.
WHY THIS EXISTS
The docs describe the backends in CAPABILITY prose -- "MI-GAN is the memory-tight pick",
"big-LaMa does not fit a minimal droplet", "~0.6-0.9 GB regardless of upload size". Those
numbers were measured once, informally, and are now load-bearing for a real deployment
decision (the free tier runs `migan`, the paid tier runs `lama`). This measures them.
WHAT IT MEASURES
For each (backend, input size): peak RSS of a FRESH process doing exactly one erase, and
the wall time. A fresh subprocess per measurement is the point -- peak RSS inside a
long-lived process is contaminated by whatever ran before it, and the question here is
what a per-request worker actually needs.
The mask is a fixed, small corner region at every size, because the claim under test is
that the learned backends crop around the mask and so their memory is bounded by the MARK
size, not the image size. If that holds, the curve is flat in input size; if it does not,
it climbs and the droplet sizing is wrong.
READING IT
RSS is the peak resident set of the whole worker, which includes the interpreter, numpy,
cv2 and (for the learned backends) onnxruntime plus the model. That is the honest number
for sizing a container -- not the model tensor alone.
DATA SAFETY
Generates its own synthetic inputs. Reads nothing from the corpus, writes nothing tracked.
uv run python scripts/resource_ceilings.py # cv2 + whatever is installed
uv run python scripts/resource_ceilings.py --max-mp 25 # push to 25 MP
"""
from __future__ import annotations
import argparse
import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
REPO = Path(__file__).resolve().parents[1]
_UV = shutil.which("uv") or "uv"
# (label, width, height) -- 1 MP up to 25 MP, the range a phone photo upload spans.
SIZES = (("1MP", 1000, 1000), ("4MP", 2000, 2000), ("12MP", 4000, 3000), ("25MP", 5000, 5000))
# The child process: build an image, erase one small corner region, report peak RSS.
# Kept as a string so each measurement is a genuinely fresh interpreter.
_CHILD = """
import json, resource, sys, time
import numpy as np
from remove_ai_watermarks.region_eraser import erase
w, h, backend = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3]
rng = np.random.default_rng(0)
# Textured, not flat: a flat image compresses and can let a backend shortcut work.
img = rng.integers(0, 255, (h, w, 3), dtype=np.uint8)
box = (w - 320, h - 90, 300, 70) # a mark-sized corner region at every input size
t0 = time.monotonic()
try:
out = erase(img, boxes=[box], backend=backend)
# Assert the fill actually RAN. A wrong call signature or a silently-declining
# backend would otherwise report the process's numpy footprint as if it were the
# backend's cost -- the first version of this script did exactly that on all 12
# cells and the numbers looked plausible.
# Compare ONLY the mask box. A full-frame `(out != img).any()` allocates a boolean
# temp the size of the image BEFORE getrusage is read -- 75 MB at 25 MP, ~17% of the
# cv2 figure, and it GROWS with the input, so the harness would partly manufacture
# the very "cv2 scales with input size" conclusion it is measuring.
bx, by, bw_, bh_ = box
changed = int((out[by:by+bh_, bx:bx+bw_] != img[by:by+bh_, bx:bx+bw_]).any())
err = "" if changed else "NO-OP: backend did not modify the masked region"
except Exception as e:
err = f"{type(e).__name__}: {e}"[:200]
elapsed = time.monotonic() - t0
peak_kb = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print("RESULT" + json.dumps({"peak_kb": peak_kb, "sec": round(elapsed, 2), "error": err}))
"""
def peak_rss_mb(peak_kb: int) -> float:
"""ru_maxrss is BYTES on macOS and KILOBYTES on Linux -- normalize to MB.
Getting this wrong silently reports 1024x off, which would look like a dramatic
finding rather than a unit bug.
"""
return peak_kb / (1024 * 1024) if sys.platform == "darwin" else peak_kb / 1024
def measure(backend: str, w: int, h: int, timeout: int = 900) -> dict[str, object] | None:
with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as fh:
fh.write(_CHILD)
child = fh.name
try:
p = subprocess.run( # noqa: S603 -- fixed argv, our own child script
[_UV, "run", "python", child, str(w), str(h), backend],
cwd=REPO,
capture_output=True,
text=True,
timeout=timeout,
check=False,
)
except subprocess.TimeoutExpired:
return {"error": f"TIMEOUT >{timeout}s", "sec": float(timeout), "peak_mb": float("nan")}
finally:
Path(child).unlink(missing_ok=True)
line = next((x for x in p.stdout.splitlines() if x.startswith("RESULT")), None)
if line is None:
return {"error": (p.stderr or p.stdout)[-200:], "sec": 0.0, "peak_mb": float("nan")}
data = json.loads(line[len("RESULT") :])
return {"error": data["error"], "sec": data["sec"], "peak_mb": round(peak_rss_mb(data["peak_kb"]), 1)}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--max-mp", type=int, default=25, help="skip sizes above this megapixel count")
a = ap.parse_args()
from remove_ai_watermarks.region_eraser import lama_available, migan_available
backends = ["cv2"] + (["migan"] if migan_available() else []) + (["lama"] if lama_available() else [])
sizes = [s for s in SIZES if (s[1] * s[2]) / 1e6 <= a.max_mp + 0.5]
print(f"backends: {backends}\nsizes: {[s[0] for s in sizes]}")
print("\nOne FRESH process per cell; the mask is a fixed small corner at every size.")
print("If the learned backends really crop around the mask, RSS stays flat in input size.\n")
print(f"{'backend':8s} {'size':6s} {'peak RSS':>11s} {'wall':>8s} note")
rows: list[tuple[str, str, dict[str, object]]] = []
for backend in backends:
for label, w, h in sizes:
r = measure(backend, w, h)
if r is None:
continue
rows.append((backend, label, r))
note = str(r["error"])[:48] or "ok"
print(f"{backend:8s} {label:6s} {r['peak_mb']:8} MB {r['sec']:7}s {note}", flush=True)
print(f"\n{'=' * 72}\nRESOURCE CEILINGS\n{'=' * 72}")
for backend in backends:
cells = [(lbl, r) for b, lbl, r in rows if b == backend and not r["error"]]
if not cells:
continue
peaks = [float(r["peak_mb"]) for _, r in cells] # type: ignore[arg-type]
lo, hi = min(peaks), max(peaks)
growth = "flat in input size" if hi <= lo * 1.6 else f"GROWS {hi / max(lo, 0.1):.1f}x with input size"
print(f" {backend:6s} peak {lo:.0f}-{hi:.0f} MB across {cells[0][0]}..{cells[-1][0]} -- {growth}")
print("\nRSS is the whole worker (interpreter + numpy + cv2 + any model), i.e. the")
print("number to size a container with, not the model tensor alone.")
if __name__ == "__main__":
main()
scripts/robustness_suite.py
+341
View File
@@ -0,0 +1,341 @@
"""Tier E: does the real CLI fail GRACEFULLY on adversarial and degenerate inputs?
WHAT "PASS" MEANS HERE
Not "it succeeded" -- most of these inputs SHOULD be rejected. A pass is a graceful
outcome: a clear message, a sane exit code, and **no unhandled traceback**, within the
timeout. The failure modes this is hunting are the ones a user actually hits and that no
unit test covers, because unit tests feed well-formed fixtures:
* an unhandled traceback -- the library crashed instead of reporting
* a hang -- worse than a crash for a batch caller
* a silent success on garbage -- it "processed" a corrupt file and wrote something
A non-zero exit with a readable error is a PASS. A Python traceback is a FAIL even when
the exit code looks tidy.
WHY THESE INPUTS
Every case is drawn from something real: ~0.2% of corpus uploads are truncated, ~2% carry
a mismatched extension, Unicode filenames were issue #17, and a wrapping service will run
concurrent jobs against one path. Decompression bombs and absurd geometry are the cheap
denial-of-service shapes any tool taking user uploads must survive.
DATA SAFETY
Builds its own inputs (synthetic, or truncated copies of committed fixtures) inside a
temp dir. Reads corpus images read-only for the one large-input case. Writes nothing
tracked.
uv run python scripts/robustness_suite.py
"""
from __future__ import annotations
import argparse
import os
import shutil
import stat
import subprocess
import sys
import tempfile
import zlib
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
REPO = Path(__file__).resolve().parents[1]
SAMPLES = REPO / "data" / "samples"
_UV = shutil.which("uv") or "uv"
# A traceback in the output means the failure escaped the error handling, whatever the
# exit code says.
_CRASH_MARKERS = ("Traceback (most recent call last)", "Fatal Python error", "Segmentation fault")
class Results:
def __init__(self) -> None:
self.rows: list[tuple[str, str, bool, str, str]] = []
def add(self, case: str, cmd: str, ok: bool, detail: str, output: str = "") -> None:
self.rows.append((case, cmd, ok, detail, "" if ok else output[-600:]))
print(f" [{'PASS' if ok else 'FAIL'}] {case:34s} {cmd:10s} {detail}", flush=True)
def report(self) -> int:
bad = [r for r in self.rows if not r[2]]
print(f"\n{'=' * 78}\nROBUSTNESS (Tier E) {len(self.rows) - len(bad)}/{len(self.rows)} graceful")
print(f"{'=' * 78}")
for case, cmd, ok, detail, out in self.rows:
if not ok:
print(f" FAIL {case} ({cmd}) {detail}")
for line in out.strip().splitlines()[-10:]:
print(f" {line}")
if not bad:
print(" every adversarial input was handled without a crash or a hang")
return 1 if bad else 0
def run(args: list[str], timeout: int = 120) -> tuple[int, str, bool]:
"""Returns (exit_code, output, timed_out)."""
try:
p = subprocess.run( # noqa: S603 -- fixed argv, driving our own CLI on purpose
[_UV, "run", "remove-ai-watermarks", *args],
cwd=REPO,
capture_output=True,
text=True,
timeout=timeout,
check=False,
)
except subprocess.TimeoutExpired:
return (-1, "", True)
return (p.returncode, p.stdout + p.stderr, False)
def graceful(res: Results, case: str, cmd: str, args: list[str], timeout: int = 120) -> None:
"""Run one adversarial case and score it on gracefulness, not on success."""
code, out, timed_out = run(args, timeout=timeout)
if timed_out:
res.add(case, cmd, False, f"HUNG (> {timeout}s)", out)
return
crashed = next((m for m in _CRASH_MARKERS if m in out), None)
if crashed:
res.add(case, cmd, False, f"unhandled {crashed} (exit {code})", out)
return
res.add(case, cmd, True, f"handled cleanly, exit {code}")
def make_inputs(tmp: Path) -> dict[str, Path]:
"""Build the adversarial corpus. Each entry is something a real upload can be."""
import numpy as np
from remove_ai_watermarks.image_io import imwrite
made: dict[str, Path] = {}
# A well-formed baseline, so a failure elsewhere is attributable to the input.
good = tmp / "good.png"
imwrite(good, np.full((600, 800, 3), 128, np.uint8))
made["good"] = good
# Truncated: a real PNG cut mid-stream (~0.2% of real uploads).
src = SAMPLES / "chatgpt-1.png"
if src.exists():
raw = src.read_bytes()
trunc = tmp / "truncated.png"
trunc.write_bytes(raw[: len(raw) // 3])
made["truncated"] = trunc
# Corrupt: correct magic bytes, garbage body.
corrupt = tmp / "corrupt.png"
corrupt.write_bytes(b"\x89PNG\r\n\x1a\n" + os.urandom(4096))
made["corrupt"] = corrupt
# Zero-byte file with a valid extension.
empty = tmp / "empty.png"
empty.write_bytes(b"")
made["empty"] = empty
# Not an image at all, but named like one.
text = tmp / "actually_text.jpg"
text.write_bytes(b"this is not an image, it is a text file pretending\n" * 50)
made["not_an_image"] = text
# Degenerate geometry: a 1x1, and a 1-pixel-tall sliver (the shape that once faulted
# cv2's GaussianBlur natively on Windows).
imwrite(tmp / "tiny.png", np.full((1, 1, 3), 200, np.uint8))
made["tiny_1x1"] = tmp / "tiny.png"
imwrite(tmp / "sliver.png", np.full((1, 4000, 3), 200, np.uint8))
made["sliver_1x4000"] = tmp / "sliver.png"
# Decompression bomb: a tiny file that decodes to a huge canvas. Hand-built so the
# on-disk size stays trivial while the declared dimensions are enormous.
bomb = tmp / "bomb.png"
bomb.write_bytes(_png_bomb(16000, 16000))
made["decompression_bomb"] = bomb
# Unicode + RTL filenames (issue #17 was Unicode-safe IO).
uni = tmp / "изображение-测试-🎨.png"
shutil.copy2(good, uni)
made["unicode_filename"] = uni
rtl = tmp / "صورة-اختبار.png"
shutil.copy2(good, rtl)
made["rtl_filename"] = rtl
# Mismatched extension: PNG content named .jpg (~2% of real uploads).
mismatch = tmp / "png_named_jpg.jpg"
shutil.copy2(good, mismatch)
made["mismatched_extension"] = mismatch
return made
def _png_bomb(w: int, h: int) -> bytes:
"""A valid PNG header declaring a huge canvas over highly-compressible data."""
def chunk(tag: bytes, data: bytes) -> bytes:
return len(data).to_bytes(4, "big") + tag + data + zlib.crc32(tag + data).to_bytes(4, "big")
ihdr = w.to_bytes(4, "big") + h.to_bytes(4, "big") + bytes([8, 2, 0, 0, 0]) # 8-bit RGB
raw = b"".join(b"\x00" + b"\x00" * (w * 3) for _ in range(h))
return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw, 9)) + chunk(b"IEND", b"")
def check_bad_inputs(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
print("\nmalformed and degenerate inputs -- every command must refuse, not crash")
out = tmp / "out.png"
for case in ("truncated", "corrupt", "empty", "not_an_image", "tiny_1x1", "sliver_1x4000"):
src = inputs.get(case)
if src is None:
continue
graceful(res, case, "identify", ["identify", str(src)])
graceful(res, case, "visible", ["visible", str(src), "-o", str(out)])
graceful(res, case, "metadata", ["metadata", str(src), "--check"])
def check_bomb(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
print("\ndecompression bomb -- must not exhaust memory or hang")
bomb = inputs.get("decompression_bomb")
if bomb is None:
return
size_kb = bomb.stat().st_size / 1024
print(f" (bomb is {size_kb:.0f} KB on disk, declares 16000x16000)")
graceful(res, "decompression_bomb", "identify", ["identify", str(bomb)], timeout=180)
graceful(res, "decompression_bomb", "visible", ["visible", str(bomb), "-o", str(tmp / "b.png")], timeout=180)
def check_filenames(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
"""Unicode/RTL paths must round-trip on BOTH read and write (issue #17)."""
print("\nunicode / RTL / mismatched-extension paths")
for case in ("unicode_filename", "rtl_filename", "mismatched_extension"):
src = inputs.get(case)
if src is None:
continue
# Write to a Unicode OUTPUT path too -- the read side alone does not prove the IO.
outp = tmp / f"вывод-{case}-📤.png"
code, out, timed = run(["identify", str(src)])
if timed or any(m in out for m in _CRASH_MARKERS):
res.add(case, "identify", False, "crashed or hung", out)
else:
res.add(case, "identify", True, f"read fine, exit {code}")
code, out, timed = run(["erase", str(src), "--region", "5,5,50,30", "-o", str(outp)])
wrote = outp.exists() and outp.stat().st_size > 0
crash = any(m in out for m in _CRASH_MARKERS)
res.add(case, "erase", wrote and not crash and not timed, f"unicode output written={wrote} exit={code}", out)
def check_output_paths(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
print("\nhostile output paths")
good = inputs["good"]
# A nested output dir that does not exist yet -- the CLI should create it.
nested = tmp / "a" / "b" / "c" / "out.png"
code, out, timed = run(["erase", str(good), "--region", "5,5,40,20", "-o", str(nested)])
crash = any(m in out for m in _CRASH_MARKERS)
res.add(
"nonexistent_nested_outdir",
"erase",
nested.exists() and not crash and not timed,
f"created={nested.exists()} exit={code}",
out,
)
# A READ-ONLY output directory -- must report, not traceback.
ro = tmp / "readonly"
ro.mkdir(exist_ok=True)
ro.chmod(stat.S_IRUSR | stat.S_IXUSR)
try:
graceful(
res, "readonly_output_dir", "erase", ["erase", str(good), "--region", "5,5,40,20", "-o", str(ro / "x.png")]
)
finally:
ro.chmod(stat.S_IRWXU) # restore so the temp dir can be cleaned up
# A directory where a file is expected.
graceful(res, "directory_as_input", "identify", ["identify", str(tmp)])
# A path that simply is not there.
graceful(res, "missing_input", "identify", ["identify", str(tmp / "nope.png")])
def check_concurrency(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
"""A wrapping service will run jobs in parallel against one input path."""
print("\nconcurrent runs against a single input")
good = inputs["good"]
def one(i: int) -> tuple[int, str, bool]:
return run(["erase", str(good), "--region", "5,5,40,20", "-o", str(tmp / f"conc_{i}.png")])
with ThreadPoolExecutor(max_workers=4) as ex:
outs = list(ex.map(one, range(4)))
crashed = [o for o in outs if any(m in o[1] for m in _CRASH_MARKERS) or o[2]]
all_written = all((tmp / f"conc_{i}.png").exists() for i in range(4))
res.add(
"4x concurrent on one input",
"erase",
not crashed and all_written,
f"{sum(1 for i in range(4) if (tmp / f'conc_{i}.png').exists())}/4 outputs, {len(crashed)} crashed",
"".join(o[1] for o in crashed),
)
def check_batch_silent_loss(res: Results, tmp: Path, inputs: dict[str, Path]) -> None:
"""The nastiest shape: NO output files AND a success exit code.
`graceful()` cannot see this class -- it scores exit code and traceback markers, and a
run that writes nothing while exiting 0 has neither. Corpus-reproduced 2026-07-20:
`batch --mode visible` into a read-only directory wrote 0 of 2 files and exited 0, so a
wrapping service would treat an empty output directory as a completed run. Any check
for a silent no-op must assert on the ARTIFACTS, not on the status.
"""
print("\nbatch into a read-only output dir -- must NOT exit 0 with nothing written")
indir = tmp / "loss_in"
indir.mkdir(exist_ok=True)
for i in range(2):
shutil.copy2(inputs["good"], indir / f"img{i}.png")
ro = tmp / "loss_out"
ro.mkdir(exist_ok=True)
ro.chmod(stat.S_IRUSR | stat.S_IXUSR)
try:
code, out, timed = run(["batch", str(indir), "-o", str(ro)], timeout=180)
written = len(list(ro.glob("*")))
finally:
ro.chmod(stat.S_IRWXU)
ok = not timed and not (code == 0 and written == 0)
res.add(
"batch_readonly_outdir", "batch", ok, f"exit {code}, {written}/2 written (exit 0 + 0 files = data loss)", out
)
def check_batch_edges(res: Results, tmp: Path) -> None:
print("\nbatch edge cases")
empty_dir = tmp / "empty_dir"
empty_dir.mkdir(exist_ok=True)
graceful(res, "batch_on_empty_dir", "batch", ["batch", str(empty_dir), "-o", str(tmp / "bo")])
# A directory of junk: nothing decodable, must not crash the whole run.
junk = tmp / "junk_dir"
junk.mkdir(exist_ok=True)
(junk / "a.png").write_bytes(b"\x89PNG\r\n\x1a\n" + os.urandom(500))
(junk / "b.jpg").write_bytes(b"not an image at all")
graceful(res, "batch_on_undecodable_dir", "batch", ["batch", str(junk), "-o", str(tmp / "jo")], timeout=180)
def main() -> None:
ap = argparse.ArgumentParser()
ap.parse_args()
res = Results()
with tempfile.TemporaryDirectory(prefix="raiw_robust_") as td:
tmp = Path(td)
print("building adversarial inputs...")
inputs = make_inputs(tmp)
print(f"built {len(inputs)} inputs\n")
check_bad_inputs(res, tmp, inputs)
check_bomb(res, tmp, inputs)
check_filenames(res, tmp, inputs)
check_output_paths(res, tmp, inputs)
check_concurrency(res, tmp, inputs)
check_batch_silent_loss(res, tmp, inputs)
check_batch_edges(res, tmp)
raise SystemExit(res.report())
if __name__ == "__main__":
main()
src/remove_ai_watermarks/api.py
+7 -1
View File
@@ -105,7 +105,13 @@ def _write_visible_result(
shutil.copyfile(source_path, out_path)
else:
image_io.write_bgr_with_alpha(out_path, result, loaded.alpha)
# imwrite is contractually NON-RAISING, so this bool is the only signal the file
# was not created. Unchecked, the metadata strip below ran on a nonexistent path
# and surfaced as a confusing "cannot read image <INPUT>" naming the OUTPUT path
# (Tier E, 2026-07-20). Raise here so a library caller and the CLI both get an
# accurate message about the write.
if not image_io.write_bgr_with_alpha(out_path, result, loaded.alpha):
raise OSError(f"failed to write output (is the destination writable?): {out_path}")
if strip_metadata:
from remove_ai_watermarks import metadata
src/remove_ai_watermarks/cli.py
+42 -19
View File
@@ -410,6 +410,22 @@ def _parse_sensitivity(value: str) -> watermark_registry.Sensitivity:
EXIT_NO_VISIBLE_MARK = 2
def _write_output_or_exit(output: Path, bgr: NDArray[Any], alpha: NDArray[Any] | None) -> None:
"""Write the final image, or fail with a readable error instead of a traceback.
`image_io.imwrite` is contractually NON-RAISING: it returns False when the codec
rejects the image or the path cannot be written. Every caller here follows its write
with `output.stat()` to report the size, so a silently-failed write (read-only
directory, full disk) died with a bare `FileNotFoundError` traceback pointing at the
stat, not at the write. Found by the Tier E adversarial sweep 2026-07-20.
Regression: `tests/test_cli_robustness.py::TestFailedWriteIsReported`.
"""
output.parent.mkdir(parents=True, exist_ok=True)
if not image_io.write_bgr_with_alpha(output, bgr, alpha):
console.print(f" Error: failed to write output (is the destination writable?): {output}")
raise SystemExit(1)
def _no_visible_mark_exit(source: Path) -> NoReturn:
"""Explain why no visible watermark was removed, then exit non-zero.
@@ -566,8 +582,11 @@ def _run_visible_auto(
except RuntimeError as e: # selected migan/lama backend whose extra is absent
console.print(f" Error: {e}")
raise SystemExit(1) from e
except (ValueError, OSError) as e: # unreadable / truncated / non-image input
console.print(f" Error: cannot read image {source.name}: {e}")
except (ValueError, OSError) as e:
# Covers BOTH an unreadable input and an unwritable output, so the message must
# not assert which: it used to say "cannot read image <input>" while quoting the
# OUTPUT path, blaming the wrong file (Tier E, 2026-07-20).
console.print(f" Error: {e}")
raise SystemExit(1) from e
elapsed = time.monotonic() - t0
@@ -630,8 +649,7 @@ def _run_visible_explicit(
raise SystemExit(1) from e
elapsed = time.monotonic() - t0
output.parent.mkdir(parents=True, exist_ok=True)
image_io.write_bgr_with_alpha(output, result, alpha)
_write_output_or_exit(output, result, alpha)
if strip_metadata:
try:
from remove_ai_watermarks.metadata import remove_ai_metadata
@@ -646,7 +664,7 @@ def _run_visible_explicit(
@main.command("visible")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
)
@@ -729,7 +747,7 @@ def _parse_region(spec: str) -> tuple[int, int, int, int]:
@main.command("erase")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--region", "regions", multiple=True, required=True, help="x,y,w,h box to erase (repeatable).")
@click.option(
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
@@ -787,8 +805,7 @@ def cmd_erase(
raise SystemExit(1) from e
elapsed = time.monotonic() - t0
output.parent.mkdir(parents=True, exist_ok=True)
image_io.write_bgr_with_alpha(output, result, alpha)
_write_output_or_exit(output, result, alpha)
if strip_metadata:
try:
@@ -805,7 +822,7 @@ def cmd_erase(
# ── Invisible watermark removal ──
@main.command("invisible")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
)
@@ -941,7 +958,7 @@ def cmd_invisible(
# ── Metadata operations ──
@main.command("metadata")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--check", is_flag=True, help="Check for AI metadata (don't modify).")
@click.option("--remove", is_flag=True, help="Remove AI metadata.")
@click.option(
@@ -1008,7 +1025,7 @@ def cmd_metadata(
# ── Provenance identification ──
@main.command("identify")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
"--no-visible",
is_flag=True,
@@ -1077,7 +1094,7 @@ def cmd_identify(ctx: click.Context, source: Path, no_visible: bool, as_json: bo
# ── Combined "all" mode ──
@main.command("all")
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
"-o", "--output", type=click.Path(path_type=Path), default=None, help="Output path (default: <source>_clean.<ext>)."
)
@@ -1278,12 +1295,11 @@ def cmd_all(
# The invisible step (and downstream cv2.IMREAD_COLOR paths) drops alpha,
# so re-attach the original alpha plane unchanged when writing the final
# output for transparent formats.
output.parent.mkdir(parents=True, exist_ok=True)
final_bgr, _ = image_io.read_bgr_and_alpha(tmp_path)
if final_bgr is None:
console.print(f"Error: Failed to read intermediate file: {tmp_path}")
raise SystemExit(1)
image_io.write_bgr_with_alpha(output, final_bgr, alpha)
_write_output_or_exit(output, final_bgr, alpha)
finally:
# Clean up temp file if it still exists
@@ -1319,8 +1335,10 @@ def _passthrough_copy(img_path: Path, out_path: Path) -> None:
"""Copy the input's pixels through to ``out_path`` unchanged (the invisible-mode skip
paths), so the output dir stays complete without touching the pixels."""
src_bgr, src_alpha = image_io.read_bgr_and_alpha(img_path)
if src_bgr is not None:
image_io.write_bgr_with_alpha(out_path, src_bgr, src_alpha)
if src_bgr is not None and not image_io.write_bgr_with_alpha(out_path, src_bgr, src_alpha):
# The point of this copy is to keep the output dir COMPLETE. A silently-dropped
# copy defeats that and leaves a hole the caller cannot see (Tier E, 2026-07-20).
raise OSError(f"failed to copy input through to output: {out_path}")
@dataclass(frozen=True)
@@ -1455,7 +1473,12 @@ def _process_batch_image(
sensitivity=options.sensitivity,
)
image_io.write_bgr_with_alpha(out_path, result, alpha)
# RAISE, never SystemExit: the batch loop catches per-image exceptions, counts
# them and exits non-zero. Discarding this flag made a read-only output directory
# produce ZERO files and still exit 0 -- silent data loss that also contradicted
# the documented batch contract (Tier E, 2026-07-20).
if not image_io.write_bgr_with_alpha(out_path, result, alpha):
raise OSError(f"failed to write output (is the destination writable?): {out_path}")
saved_alpha = alpha
if mode in ("invisible", "all"):
@@ -1479,8 +1502,8 @@ def _process_batch_image(
# so re-attach the cached alpha when the input had transparency.
if mode == "all" and saved_alpha is not None:
final_bgr, _ = image_io.read_bgr_and_alpha(out_path)
if final_bgr is not None:
image_io.write_bgr_with_alpha(out_path, final_bgr, saved_alpha)
if final_bgr is not None and not image_io.write_bgr_with_alpha(out_path, final_bgr, saved_alpha):
raise OSError(f"failed to re-attach alpha to output: {out_path}")
return synthid_skipped
src/remove_ai_watermarks/image_io.py
+10 -5
View File
@@ -221,8 +221,8 @@ def read_bgr_and_alpha(path: str | Path) -> tuple[NDArray[Any] | None, NDArray[A
return image, None
def write_bgr_with_alpha(path: str | Path, bgr: NDArray[Any], alpha: NDArray[Any] | None) -> None:
"""Write BGR (with optional alpha) to ``path``.
def write_bgr_with_alpha(path: str | Path, bgr: NDArray[Any], alpha: NDArray[Any] | None) -> bool:
"""Write BGR (with optional alpha) to ``path``. Returns ``imwrite``'s success flag.
When ``alpha`` is provided and the output extension supports it, the original
alpha plane is rejoined unchanged. The watermark region is NOT made transparent:
@@ -230,10 +230,15 @@ def write_bgr_with_alpha(path: str | Path, bgr: NDArray[Any], alpha: NDArray[Any
transparent hole that renders as a white box on any non-transparent viewer
(issue #30). Preserving the input alpha keeps genuinely transparent backgrounds
intact without inventing new holes.
Returning the flag is load-bearing: :func:`imwrite` is contractually non-raising, so
this is the ONLY signal a caller gets that the file was not created. Discarding it let
a failed write (read-only directory, full disk) run on to ``output.stat()`` and die
with a bare ``FileNotFoundError`` traceback instead of a readable error.
Regression: ``tests/test_cli_robustness.py::TestFailedWriteIsReported``.
"""
import numpy as np
if alpha is None or Path(path).suffix.lower() not in ALPHA_FORMATS:
imwrite(path, bgr)
return
imwrite(path, np.dstack([bgr, alpha]))
return imwrite(path, bgr)
return imwrite(path, np.dstack([bgr, alpha]))
tests/test_cli_robustness.py
+114
View File
@@ -0,0 +1,114 @@
"""A failed write and a directory input must REPORT, not raise a traceback.
Both found 2026-07-20 by the Tier E adversarial sweep (`scripts/robustness_suite.py`),
which drives the real CLI over degenerate inputs. Neither was reachable from the 849-test
suite, because unit tests feed well-formed fixtures into writable directories.
1. READ-ONLY OUTPUT DIRECTORY. `image_io.imwrite` is contractually non-raising -- it
returns False when the codec rejects the image or the path cannot be written. But
`write_bgr_with_alpha` discarded that bool and returned None, so no caller could tell a
failed write from a successful one. `cmd_erase` then ran `output.stat()` on a file that
was never created and died with `FileNotFoundError`. The signal existed the whole way
down and was thrown away by the wrapper.
2. A DIRECTORY PASSED WHERE A FILE IS EXPECTED. `click.Path(exists=True)` accepts
directories unless told otherwise, so `identify <dir>` reached the metadata scanner and
raised `IsADirectoryError` from `open()`. (`batch` was already correct: it declares
`file_okay=False`.)
"""
from __future__ import annotations
import numpy as np
import pytest
from click.testing import CliRunner
from remove_ai_watermarks import image_io
from remove_ai_watermarks.cli import main
@pytest.fixture
def bgr() -> np.ndarray:
return np.full((60, 80, 3), 128, np.uint8)
class TestFailedWriteIsReported:
def test_write_bgr_with_alpha_reports_failure(self, tmp_path, bgr):
"""The wrapper must PROPAGATE imwrite's failure signal, not swallow it.
Uses a nonexistent directory rather than `chmod`: CI runs windows-latest, where
`os.chmod` cannot make a DIRECTORY unwritable, so a chmod-based assertion would be
a platform-dependent flake. `write_bgr_with_alpha` does no mkdir of its own, so a
missing parent is a genuine write failure on every platform.
"""
assert image_io.write_bgr_with_alpha(tmp_path / "no-such-dir" / "x.png", bgr, None) is False
def test_write_bgr_with_alpha_reports_success(self, tmp_path, bgr):
"""The other direction, so the assertion above cannot pass by always being False."""
assert image_io.write_bgr_with_alpha(tmp_path / "ok.png", bgr, None) is True
def test_erase_reports_a_failed_write_without_a_traceback(self, tmp_path, bgr, monkeypatch):
"""The CLI must exit non-zero with a readable message, not raise FileNotFoundError.
The write is forced to fail by patching, not by `chmod`: the CLI mkdirs the parent,
so a missing directory would not reproduce it, and chmod on a directory is a no-op
on Windows. Patching states the condition under test directly -- "the write failed".
"""
src = tmp_path / "in.png"
image_io.imwrite(src, bgr)
monkeypatch.setattr(image_io, "write_bgr_with_alpha", lambda *a, **k: False)
result = CliRunner().invoke(main, ["erase", str(src), "--region", "5,5,20,10", "-o", str(tmp_path / "x.png")])
assert result.exit_code != 0
assert not isinstance(result.exception, FileNotFoundError), "write failure escaped as a traceback"
assert "write" in result.output.lower() or "failed" in result.output.lower()
def test_batch_counts_a_failed_write_instead_of_exiting_zero(self, tmp_path, bgr, monkeypatch):
"""The worst shape of this bug: no output files AND a success exit code.
Corpus-reproduced 2026-07-20 -- `batch --mode visible` into a read-only directory
wrote ZERO files for 2 inputs and exited 0, so a wrapping service would treat an
empty output directory as a completed run. The batch loop counts per-image
exceptions, so the write must RAISE there, never `SystemExit` (which would abort
the whole run instead of failing one image).
"""
indir = tmp_path / "in"
indir.mkdir()
for i in range(2):
image_io.imwrite(indir / f"img{i}.png", bgr)
monkeypatch.setattr(image_io, "write_bgr_with_alpha", lambda *a, **k: False)
result = CliRunner().invoke(main, ["batch", str(indir), "-o", str(tmp_path / "out"), "--mode", "visible"])
assert result.exit_code != 0, "a batch that wrote nothing must not exit 0"
class TestApiReportsFailedWrite:
"""The same bug lived one layer down, in the library API, with a misleading message.
`api._write_visible_result` also discarded the write flag, then ran the metadata strip
on a file that was never created. The resulting `FileNotFoundError` surfaced through
the CLI as `cannot read image <INPUT>: ... <OUTPUT path>` -- it blamed the input while
quoting the output. A library caller got the same confusing error with no CLI at all.
"""
def test_remove_visible_raises_a_clear_error_on_unwritable_output(self, tmp_path, bgr, monkeypatch):
"""Patched rather than chmod'd: `_write_visible_result` mkdirs the parent, so a
missing directory would not reproduce it, and chmod on a directory is a no-op on
the Windows CI runner."""
from remove_ai_watermarks.api import remove_visible
monkeypatch.setattr(image_io, "write_bgr_with_alpha", lambda *a, **k: False)
with pytest.raises(OSError, match="failed to write output"):
remove_visible(bgr, tmp_path / "out.png", strip_metadata=False)
class TestDirectoryInputIsRejected:
@pytest.mark.parametrize("cmd", ["identify", "visible", "erase", "metadata", "invisible", "all"])
def test_directory_as_source_is_a_clean_usage_error(self, tmp_path, cmd):
"""A directory must be refused by argument parsing, never reach the scanners."""
args = [cmd, str(tmp_path)]
if cmd == "erase":
args += ["--region", "1,1,5,5"]
if cmd == "metadata":
args += ["--check"]
result = CliRunner().invoke(main, args)
assert result.exit_code != 0
assert not isinstance(result.exception, IsADirectoryError), "directory reached the file reader"