CalvinBackup/remove-ai-watermarks
1
0
Fork 0
mirror of https://github.com/wiltodelta/remove-ai-watermarks.git synced 2026-08-07 06:28: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
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()