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

Add full-surface verification harnesses and corpus sweep plan

Browse Source 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
2026-07-20 08:14:50 -07:00
co-authored by Claude Opus 4.8
parent cfefd9d819
commit c150180acf
6 changed files with 2021 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/fill_quality.py
+323
View File
@@ -0,0 +1,323 @@
"""Tier B1: measure visible-mark FILL quality against a constructed ground truth.
THE PROBLEM THIS SOLVES
A real watermarked image has no clean counterpart, so fill quality has only ever been
eyeballed ("cv2 smears on texture, LaMa is best") and never carried a number. The docs
state a preference order for `--backend auto` that rests on no measurement.
THE CONSTRUCTION
Take a VERIFIED-CLEAN corpus image, stamp a known mark onto it using the mark's own
captured alpha map (the forward model the reverse-alpha work established:
`stamped = (1-a)*bg + a*white`), then remove it and compare against the TRUE original.
The original is the answer key that reality never provides.
TWO MEASUREMENT DECISIONS THAT MAKE OR BREAK THIS
* Score INSIDE the footprint only. The fill touches a tiny corner, so whole-frame PSNR
sits near 60 dB whatever the backend does and would rank them all "excellent".
* Report the DAMAGE baseline (stamped vs original) next to the recovery (filled vs
original). Without it a PSNR is uninterpretable: 30 dB is a triumph if the mark cost
12 dB and a failure if it cost 29 dB. The honest metric is how much of the gap the
fill closes.
WHAT IT DOES NOT MEASURE
Detection. The mark is localized with `force=True`, so a miss cannot contaminate the
fill numbers -- this isolates the FILL. Detection accuracy is Tier C's job.
DATA SAFETY
Corpus images are user uploads: read-only, local analysis, output under a gitignored
data/spaces/ path. No image content is written into the report.
uv run python scripts/fill_quality.py --n 60
"""
from __future__ import annotations
import argparse
import collections
import glob
import json
import random
import sys
from pathlib import Path
from typing import Any
import cv2
import numpy as np
sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent))
from invisible_quality_audit import _ssim # reuse, do not reimplement a third SSIM
REPO = Path(__file__).resolve().parents[1]
CORPUS = REPO / "data" / "spaces" / "originals"
OUT = REPO / "data" / "spaces" / "_fill_quality.jsonl"
# Text marks: a bundled alpha PNG plus the engine's own corner geometry.
STAMPABLE = ("doubao", "jimeng", "samsung")
# Marks with no bundled alpha asset. Both expose a default footprint via
# `footprint_mask(force=True)`, so they are stamped by fitting their own alpha source
# into that slot: Gemini's alpha is derived from its background captures, the pill's
# from its synthetic font-rendered silhouette.
SLOT_STAMPABLE = ("gemini", "jimeng_pill")
def psnr(a: np.ndarray, b: np.ndarray) -> float:
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
return float("inf") if mse == 0 else 10 * float(np.log10(255.0**2 / mse))
def texture_of(box: np.ndarray) -> float:
"""Median Sobel magnitude -- the same texture proxy the pill's flatness gate uses."""
g = cv2.cvtColor(box, cv2.COLOR_BGR2GRAY) if box.ndim == 3 else box
gx = cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)
gy = cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)
return float(np.median(cv2.magnitude(gx, gy)))
def engine_for(mark_key: str) -> Any:
"""The TextMarkEngine instance for a mark key (the registry does not expose one)."""
import importlib
mod = importlib.import_module(f"remove_ai_watermarks.{mark_key}_engine")
cls = next(
obj
for name, obj in vars(mod).items()
# Exclude the imported base class: it takes a `config` arg, the subclasses do not.
if isinstance(obj, type) and name.endswith("Engine") and name != "TextMarkEngine"
)
return cls()
def stamp(image: np.ndarray, mark_key: str) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
"""Composite a mark's alpha glyph into its canonical corner. Returns (stamped, bbox)."""
from remove_ai_watermarks._text_mark_engine import load_alpha_template
engine = engine_for(mark_key)
cfg = engine.config
alpha = load_alpha_template(cfg.asset_name)
if alpha is None:
return None
loc = engine.locate(image)
base = engine.scale_base(image)
gw = max(cfg.min_gw, int(cfg.alpha_width_frac * base))
gh = max(4, int(cfg.alpha_height_frac * base))
if gw < 8 or gh < 4 or gw > loc.w or gh > loc.h:
return None
a = cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_AREA).astype(np.float32)
x = loc.x + (loc.w - gw) // 2
y = loc.y + (loc.h - gh) // 2
out = image.copy()
roi = out[y : y + gh, x : x + gw].astype(np.float32)
a3 = a[..., None]
out[y : y + gh, x : x + gw] = np.clip(roi * (1 - a3) + 255.0 * a3, 0, 255).astype(np.uint8)
return out, (x, y, gw, gh)
def _slot_alpha(mark_key: str) -> np.ndarray | None:
"""The alpha source for a mark that ships no bundled alpha PNG."""
if mark_key == "gemini":
from remove_ai_watermarks.gemini_engine import _shared_engine
# _shared_engine is the package's own lru_cache singleton: constructing
# GeminiEngine() per image reloads its captures, recomputes both alpha maps and
# rebuilds the whole template ladder -- the exact cost that singleton exists to
# avoid. Private attribute access is deliberate (no accessor exists); a
# measurement script may reach in, product code must not.
return np.asarray(_shared_engine()._alpha_large, dtype=np.float32)
if mark_key == "jimeng_pill":
from remove_ai_watermarks._text_mark_engine import load_alpha_template
return load_alpha_template("jimeng_pill.png")
return None
def stamp_slot(image: np.ndarray, mark_key: str) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
"""Stamp a mark into its OWN default footprint slot (the `--no-detect` geometry).
Used for gemini and the pill, which have no bundled alpha asset with corner
fractions. The mark is fitted into the middle of its slot so the fill has to
recover the same kind of region it would in production.
"""
from remove_ai_watermarks.watermark_registry import get_mark
alpha = _slot_alpha(mark_key)
if alpha is None:
return None
loc = get_mark(mark_key).localize(image, force=True)
if loc.mask is None:
return None
ys, xs = np.nonzero(loc.mask)
if ys.size == 0:
return None
y0, y1, x0, x1 = int(ys.min()), int(ys.max()), int(xs.min()), int(xs.max())
bw, bh = x1 - x0 + 1, y1 - y0 + 1
gw, gh = max(8, int(bw * 0.7)), max(8, int(bh * 0.7))
if gw < 8 or gh < 8 or gw > image.shape[1] or gh > image.shape[0]:
return None
a = cv2.resize(alpha, (gw, gh), interpolation=cv2.INTER_AREA).astype(np.float32)
a = np.clip(a, 0.0, 1.0)
x, y = x0 + (bw - gw) // 2, y0 + (bh - gh) // 2
out = image.copy()
roi = out[y : y + gh, x : x + gw].astype(np.float32)
a3 = a[..., None]
out[y : y + gh, x : x + gw] = np.clip(roi * (1 - a3) + 255.0 * a3, 0, 255).astype(np.uint8)
return out, (x, y, gw, gh)
def stamp_any(image: np.ndarray, mark_key: str) -> tuple[np.ndarray, tuple[int, int, int, int]] | None:
return stamp(image, mark_key) if mark_key in STAMPABLE else stamp_slot(image, mark_key)
def clean_sources(n: int, seed: int = 11) -> list[Path]:
"""Corpus images with NO metadata signal and NO mark detection -- verified clean."""
from remove_ai_watermarks.identify import identify
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.watermark_registry import detect_marks
pool = glob.glob(str(CORPUS / "*" / "*"))
random.Random(seed).shuffle(pool) # noqa: S311 -- deterministic sampling, not crypto
out: list[Path] = []
for p in pool:
if len(out) >= n:
break
path = Path(p)
try:
img = imread(str(path))
if img is None or min(img.shape[:2]) < 400:
continue
# check_visible=False: identify would otherwise decode the file a second
# time and run the very same visible detectors as detect_marks below.
if identify(path, check_visible=False).signals:
continue
if any(d.detected for d in detect_marks(img)):
continue
except Exception: # noqa: S112 -- a bad corpus file just is not a candidate
continue
out.append(path)
return out
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--n", type=int, default=60, help="clean source images")
ap.add_argument("--out", type=Path, default=OUT)
a = ap.parse_args()
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.region_eraser import lama_available, migan_available
from remove_ai_watermarks.watermark_registry import fill, get_mark
backends = ["cv2"] + (["migan"] if migan_available() else []) + (["lama"] if lama_available() else [])
print(f"backends under test: {backends}")
print(f"selecting {a.n} verified-clean sources...", flush=True)
sources = clean_sources(a.n)
print(f"got {len(sources)}\n", flush=True)
rows: list[dict[str, Any]] = []
with open(a.out, "w", encoding="utf-8") as fh:
for i, src in enumerate(sources, 1):
base = imread(str(src))
if base is None:
continue
for key in (*STAMPABLE, *SLOT_STAMPABLE):
st = stamp_any(base, key)
if st is None:
continue
stamped, (x, y, w, h) = st
# Score a slightly padded box: the fill dilates, so a tight glyph box
# would miss damage the backend does just outside the glyph.
pad = 8
y0, y1 = max(0, y - pad), min(base.shape[0], y + h + pad)
x0, x1 = max(0, x - pad), min(base.shape[1], x + w + pad)
truth = base[y0:y1, x0:x1]
dmg_psnr = psnr(stamped[y0:y1, x0:x1], truth)
dmg_ssim = _ssim(
cv2.cvtColor(stamped[y0:y1, x0:x1], cv2.COLOR_BGR2GRAY), cv2.cvtColor(truth, cv2.COLOR_BGR2GRAY)
)
tex = texture_of(truth)
mark = get_mark(key)
loc = mark.localize(stamped, force=True)
if loc.mask is None:
continue
for backend in backends:
try:
filled = fill(stamped, loc.mask, backend=backend)
except Exception as e:
rows.append({"src": src.name, "mark": key, "backend": backend, "error": str(e)[:150]})
continue
# Invariant: the fill must not touch anything outside its mask.
outside = loc.mask == 0
untouched = bool(np.array_equal(filled[outside], stamped[outside]))
rec = {
"src": src.name,
"mark": key,
"backend": backend,
"texture": round(tex, 2),
"damage_psnr": round(dmg_psnr, 2),
"damage_ssim": round(dmg_ssim, 4),
"filled_psnr": round(psnr(filled[y0:y1, x0:x1], truth), 2),
"filled_ssim": round(
_ssim(
cv2.cvtColor(filled[y0:y1, x0:x1], cv2.COLOR_BGR2GRAY),
cv2.cvtColor(truth, cv2.COLOR_BGR2GRAY),
),
4,
),
"outside_mask_untouched": untouched,
}
rows.append(rec)
fh.write(json.dumps(rec) + "\n")
if i % 10 == 0:
fh.flush()
print(f" {i}/{len(sources)}", flush=True)
report(rows, a.out)
def report(rows: list[dict[str, Any]], out: Path) -> None:
good = [r for r in rows if "error" not in r]
if not good:
print("no measurements")
return
tex = sorted(r["texture"] for r in good)
t1, t2 = tex[len(tex) // 3], tex[2 * len(tex) // 3]
def bucket(t: float) -> str:
return "flat" if t <= t1 else ("mid" if t <= t2 else "textured")
print(f"\n{'=' * 78}\nFILL QUALITY n={len(good)} measurements (texture terciles at {t1:.1f} / {t2:.1f})")
print(f"{'=' * 78}")
print("recovered = filled PSNR minus damaged PSNR; how much of the mark's damage the fill undoes\n")
hdr = f"{'mark':9s} {'backend':8s} {'bg':9s} {'n':>4s} {'damaged':>9s} {'filled':>9s}"
print(f"{hdr} {'recovered':>10s} {'ssim':>7s}")
agg: dict[tuple[str, str, str], list[dict[str, Any]]] = collections.defaultdict(list)
for r in good:
agg[(r["mark"], r["backend"], bucket(r["texture"]))].append(r)
for (mark, backend, bg), rs in sorted(agg.items()):
# MEDIAN, not mean: a fill that reproduces a flat background EXACTLY scores
# PSNR=inf, and a single inf makes the mean inf -- which reported "+inf" for
# every flat bucket on the first run and hid the real numbers.
finite = [r for r in rs if np.isfinite(r["filled_psnr"])]
perfect = len(rs) - len(finite)
d = float(np.median([r["damage_psnr"] for r in rs]))
f = float(np.median([r["filled_psnr"] for r in finite])) if finite else float("inf")
sv = float(np.median([r["filled_ssim"] for r in rs]))
# Recovery is the MEDIAN OF THE PER-IMAGE DIFFERENCES, not the difference of the
# two medians. Those are not the same statistic on skewed data and they can even
# disagree in SIGN: the first version of this report subtracted medians and
# produced a backend ranking that a paired comparison did not support.
rec = float(np.median([r["filled_psnr"] - r["damage_psnr"] for r in finite])) if finite else float("inf")
tail = f" ({perfect} exact)" if perfect else ""
print(f"{mark:9s} {backend:8s} {bg:9s} {len(rs):4d} {d:9.2f} {f:9.2f} {rec:+10.2f} {sv:7.4f}{tail}")
bad = [r for r in good if not r["outside_mask_untouched"]]
print(f"\noutside-mask invariant violations: {len(bad)}" + (" <-- BUG" if bad else " (none)"))
print(f"records: {out}")
if __name__ == "__main__":
main()
scripts/pill_gate_audit.py
+224
View File
@@ -0,0 +1,224 @@
"""Measure the Jimeng pill on the PRODUCT path, not the raw detector.
WHY THIS EXISTS
`visible_removal_audit.py` calls `get_mark(key).remove` directly, which bypasses
`_keep_pill`. For the pill that is the wrong path and it reads as a disaster: on a
300-image sample the raw path was "detector still fires after removal" 68-75% of the
time. The product does not do that -- `--mark auto` gates the pill hard (never on
Doubao; unrestricted only when the bottom-right wordmark fired; otherwise only on a
flat footprint). Measured through the gate on the same sample, 7 of 80 raw detections
survived to removal and 6 of those were corroborated real.
That sample was too small to state a precision (95% CI 49-97%). This script runs the
gated path over EVERY pill positive in the corpus so the interval is usable.
THE CORROBORATION PROXY AND ITS BIAS
"Real pill" here means: TC260 metadata names Jimeng, OR the bottom-right "★ 即梦AI"
wordmark is detected. Both are independent of the pill detector, which is the point.
But the proxy MISSES a real pill on a metadata-stripped screenshot with no visible
wordmark -- so measured precision is a LOWER BOUND, not a point estimate. Do not quote
it as if it were exact.
Corpus images are user uploads: read-only, local analysis, gitignored output.
uv run python scripts/pill_gate_audit.py --jobs 7
"""
from __future__ import annotations
import argparse
import collections
import json
import math
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from concurrent.futures import TimeoutError as FutureTimeout
from concurrent.futures.process import BrokenProcessPool
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
REPO = Path(__file__).resolve().parents[1]
POSITIVES = REPO / "data" / "spaces" / "_visible_positives.jsonl"
OUT = REPO / "data" / "spaces" / "_pill_gate_audit.jsonl"
def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]:
if n == 0:
return (0.0, 0.0)
p = k / n
d = 1 + z * z / n
c = (p + z * z / (2 * n)) / d
h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d
return (100 * (c - h), 100 * (c + h))
def _one(path_str: str) -> dict[str, object]:
from remove_ai_watermarks.api import visible_provenance
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.pill_engine import PillEngine
from remove_ai_watermarks.watermark_registry import detect_marks, get_mark, remove_auto_marks
rec: dict[str, object] = {"path": Path(path_str).name, "status": "ok"}
try:
img = imread(path_str)
if img is None:
return {**rec, "status": "unreadable"}
prov = visible_provenance(Path(path_str))
dets = {d.key: d for d in detect_marks(img) if d.detected}
wordmark = "jimeng" in dets
tc260 = "jimeng" in prov
eng = PillEngine()
_, labels = remove_auto_marks(img, sensitivity="auto", provenance=prov, backend="cv2")
# Match the EXACT pill label. A substring test on "AI生成" also matches Doubao's
# label ("Doubao 豆包AI生成 text"), which counted Doubao removals as pill removals
# and inflated the pill's measured precision.
pill_label = get_mark("jimeng_pill").label
removed = any(str(x) == pill_label for x in labels)
# The gate reads the ARBITRATED keys, not strict detection: jimeng can be accepted
# on its relaxed (provenance-confirmed) arm, which `detect_marks` never shows. Take
# the arm from what was actually removed, or every such case lands in "other".
jimeng_decided = any(str(x) == get_mark("jimeng").label for x in labels)
return {
**rec,
"removed_by_product": bool(removed),
"corroborated": bool(wordmark or tc260),
"wordmark": bool(wordmark),
"jimeng_decided": bool(jimeng_decided),
"tc260": bool(tc260),
"doubao_fired": "doubao" in dets,
"footprint_flat": bool(eng.footprint_is_flat(img)),
}
except Exception as e:
return {**rec, "status": f"error:{type(e).__name__}"}
def _batch(paths: list[str], jobs: int, timeout: int) -> list[dict[str, object]]:
try:
with ProcessPoolExecutor(max_workers=jobs) as ex:
futs = [ex.submit(_one, p) for p in paths]
return [f.result() for f in as_completed(futs, timeout=timeout)]
except (FutureTimeout, BrokenProcessPool, OSError, RuntimeError):
out = []
for p in paths:
try:
out.append(_one(p))
except BaseException: # a native crash costs one file, not the sweep
out.append({"path": Path(p).name, "status": "crashed"})
return out
def report(rows: list[dict[str, object]]) -> None:
ok = [r for r in rows if r.get("status") == "ok"]
if not ok:
print("no usable rows")
return
removed = [r for r in ok if r["removed_by_product"]]
corro = [r for r in ok if r["corroborated"]]
tp = [r for r in removed if r["corroborated"]]
missed = [r for r in corro if not r["removed_by_product"]]
print(f"\n{'=' * 76}\nJIMENG PILL ON THE PRODUCT PATH n={len(ok)} raw detections\n{'=' * 76}")
print(f"raw detections corroborated as real : {len(corro):5d} ({100 * len(corro) / len(ok):.1f}%)")
print(f"the gate lets through to removal : {len(removed):5d} ({100 * len(removed) / len(ok):.1f}%)")
if removed:
lo, hi = wilson(len(tp), len(removed))
pct = 100 * len(tp) / len(removed)
print(f" of those, corroborated real : {len(tp):5d} -> precision {pct:.1f}% (95% CI {lo:.1f}-{hi:.1f})")
if corro:
lo, hi = wilson(len(tp), len(corro))
pct = 100 * len(tp) / len(corro)
head = f"corroborated pills removed : {len(tp):5d}/{len(corro)}"
print(f"{head} -> recall {pct:.1f}% (95% CI {lo:.1f}-{hi:.1f})")
print(f"corroborated pills the gate suppressed: {len(missed):5d}")
print("\nprecision is a LOWER BOUND: the corroboration proxy cannot see a real pill on a")
print("metadata-stripped image whose wordmark is absent or missed.\n")
print("which gate arm let it through:")
arms: collections.Counter[str] = collections.Counter()
for r in removed:
if r.get("jimeng_decided"):
arms["jimeng wordmark accepted -> unrestricted"] += 1
elif r["tc260"]:
arms["tc260 metadata + flat footprint"] += 1
else:
arms["other (unexpected -- the gate has no third arm)"] += 1
for k, v in arms.most_common():
print(f" {v:5d} {k}")
print("\nwhy the rest were suppressed:")
sup: collections.Counter[str] = collections.Counter()
for r in ok:
if r["removed_by_product"]:
continue
if r["doubao_fired"]:
sup["doubao fired (pill never rides on it)"] += 1
elif not r["corroborated"]:
sup["no corroboration (no wordmark, no TC260)"] += 1
elif not r["footprint_flat"]:
sup["corroborated but footprint too textured"] += 1
else:
sup["other"] += 1
for k, v in sup.most_common():
print(f" {v:5d} {k}")
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--jobs", type=int, default=max(1, (os.cpu_count() or 4) - 2))
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--batch", type=int, default=200)
ap.add_argument("--timeout", type=int, default=900)
ap.add_argument("--out", type=Path, default=OUT)
ap.add_argument("--restart", action="store_true")
a = ap.parse_args()
paths = []
with open(POSITIVES, encoding="utf-8") as fh:
for line in fh:
try:
r = json.loads(line)
except Exception: # noqa: S112
continue
if "jimeng_pill" in (r.get("keys") or []):
paths.append(r["path"])
paths.sort()
if a.limit:
paths = paths[: a.limit]
done: set[str] = set()
if a.restart and a.out.exists():
a.out.unlink() # --restart must TRUNCATE; appending kept the old rows in the report
if a.out.exists() and not a.restart:
with open(a.out, encoding="utf-8") as fh:
for line in fh:
try:
done.add(json.loads(line)["path"])
except Exception: # noqa: S112
continue
todo = [p for p in paths if Path(p).name not in done]
print(f"pill positives {len(paths)} done {len(done)} to do {len(todo)} jobs {a.jobs}", flush=True)
rows: list[dict[str, object]] = []
with open(a.out, "a", encoding="utf-8") as fh:
for i in range(0, len(todo), a.batch):
for rec in _batch(todo[i : i + a.batch], a.jobs, a.timeout):
fh.write(json.dumps(rec) + "\n")
rows.append(rec)
fh.flush()
print(f" {len(rows)}/{len(todo)}", flush=True)
with open(a.out, encoding="utf-8") as fh:
allrows = []
for line in fh:
try:
allrows.append(json.loads(line))
except Exception: # noqa: S112
continue
report(allrows)
print(f"records -> {a.out}")
if __name__ == "__main__":
main()
scripts/sidecar_regression.py
+216
View File
@@ -0,0 +1,216 @@
"""Tier A1: diff today's `identify` against the verdicts recorded in the corpus sidecars.
`data/spaces/identify/<day>/<uid>.json` holds the verdict a past run produced for
`data/spaces/originals/<day>/<uid>_src.<ext>`. Re-running identify and diffing turns the
corpus into a ~39k-image behavioral regression suite that needs no new labelling.
WHAT A DIFF MEANS -- READ THIS BEFORE PANICKING
The sidecars were written by OLDER versions, so an intended improvement shows up as a
diff exactly like a regression does. The output is therefore CLASSIFIED, not pass/fail:
lost_ai verdict was AI, now is not <- the alarm; almost always a bug
lost_signal a signal family stopped firing <- the alarm
new_ai verdict was not AI, now is <- usually an improvement
new_signal a signal family started firing <- usually an improvement
platform attribution changed
confidence confidence level changed
unchanged identical on every compared axis
Only `lost_*` is a regression by default. Everything else needs a look before the
baseline is moved.
WHY FAMILIES, NOT RAW STRINGS
Watermark descriptions are human-readable prose and their WORDING has changed between
versions ("C2PA Content Credentials (OpenAI)" vs "... (OpenAI, Truepic)"). Diffing raw
strings would report every rewording as a lost+new signal pair and bury the real
regressions. Signals are normalized to families (c2pa, synthid, visible_sparkle, ...)
so the comparison tracks BEHAVIOR, not phrasing.
DATA SAFETY
Corpus images are user uploads: read-only, local analysis. Output goes to a gitignored
path under data/spaces/ and records uids, never image content.
uv run python scripts/sidecar_regression.py --sample 500 # representative trial
uv run python scripts/sidecar_regression.py # full corpus, resumable
Use `--sample` for a trial, never `--limit`: the sidecar list is sorted by day, so
`--limit N` reads only the EARLIEST day and its result does not generalize (the first
trial run of this script took 200 files that were all 2026-05-29).
"""
from __future__ import annotations
import argparse
import collections
import glob
import json
import os
import random
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
REPO = Path(__file__).resolve().parents[1]
IDENTIFY_DIR = REPO / "data" / "spaces" / "identify"
ORIGINALS = REPO / "data" / "spaces" / "originals"
OUT = REPO / "data" / "spaces" / "_sidecar_regression.jsonl"
# Map a watermark description to a stable behavior family. Order matters: the first
# matching pattern wins, so put the specific tokens above the generic ones.
_FAMILIES: tuple[tuple[str, tuple[str, ...]], ...] = (
("synthid", ("synthid",)),
("visible_sparkle", ("sparkle",)),
("visible_doubao", ("豆包", "doubao")),
("visible_jimeng", ("即梦", "jimeng")),
("visible_samsung", ("galaxy ai", "generati dall")),
("aigc_tc260", ("aigc", "tc260")),
("iptc", ("iptc", "digitalsourcetype", "made with ai", "made-with-ai")),
("c2pa", ("c2pa", "content credentials")),
("trustmark", ("trustmark",)),
("open_invisible", ("open invisible", "dwt", "stable diffusion xl")),
("xai_signature", ("xai", "grok signature")),
("exif_generator", ("exif", "software tag", "png text")),
)
def family_of(description: str) -> str:
low = description.lower()
for name, tokens in _FAMILIES:
if any(t in low for t in tokens):
return name
return "other"
def families(descriptions: list[str]) -> set[str]:
return {family_of(d) for d in descriptions or []}
def compare(sidecar: dict, report: object) -> dict:
"""Classify the difference between a recorded verdict and a fresh one."""
old_ai, new_ai = sidecar.get("is_ai_generated"), getattr(report, "is_ai_generated", None)
old_fam = families(sidecar.get("watermarks") or [])
new_fam = families(list(getattr(report, "watermarks", []) or []))
old_plat, new_plat = sidecar.get("platform"), getattr(report, "platform", None)
old_conf, new_conf = sidecar.get("confidence"), getattr(report, "confidence", None)
classes: list[str] = []
# Treat only a real True->not-True transition as lost. None means "unknown", and the
# library never asserts False, so None->None is not a change.
if bool(old_ai) and not bool(new_ai):
classes.append("lost_ai")
if not bool(old_ai) and bool(new_ai):
classes.append("new_ai")
if old_fam - new_fam:
classes.append("lost_signal")
if new_fam - old_fam:
classes.append("new_signal")
if old_plat != new_plat:
classes.append("platform")
if old_conf != new_conf:
classes.append("confidence")
return {
"classes": classes or ["unchanged"],
"old_ai": old_ai,
"new_ai": new_ai,
"lost_families": sorted(old_fam - new_fam),
"new_families": sorted(new_fam - old_fam),
"old_platform": old_plat,
"new_platform": new_plat,
"old_confidence": old_conf,
"new_confidence": new_conf,
}
def _one(sidecar_path: str) -> dict:
uid = os.path.basename(sidecar_path)[:-5]
day = os.path.basename(os.path.dirname(sidecar_path))
try:
with open(sidecar_path, encoding="utf-8") as fh:
sidecar = json.load(fh)
except Exception as e: # a corrupt sidecar must not kill the sweep
return {"uid": uid, "day": day, "status": "sidecar_unreadable", "error": str(e)[:200]}
src = sidecar.get("src") or ""
image = ORIGINALS / day / src
if not src or not image.exists():
found = glob.glob(str(ORIGINALS / day / f"{uid}*"))
if not found:
return {"uid": uid, "day": day, "status": "image_missing"}
image = Path(found[0])
try:
from remove_ai_watermarks.identify import identify
report = identify(image)
except Exception as e: # record and continue; a crash IS a finding
return {"uid": uid, "day": day, "status": "identify_raised", "error": f"{type(e).__name__}: {e}"[:300]}
return {"uid": uid, "day": day, "status": "ok", **compare(sidecar, report)}
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--limit", type=int, default=0, help="only N sidecars (trial run)")
ap.add_argument(
"--sample",
type=int,
default=0,
help="randomly sample N across ALL days -- use this for a trial, not --limit: "
"the sidecar list is sorted, so --limit takes one day and is not representative",
)
ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 2))
ap.add_argument("--out", type=Path, default=OUT)
ap.add_argument("--restart", action="store_true", help="ignore existing output and start over")
a = ap.parse_args()
sidecars = sorted(glob.glob(str(IDENTIFY_DIR / "*" / "*.json")))
if a.sample:
random.Random(7).shuffle(sidecars) # noqa: S311 -- deterministic sampling, not crypto
sidecars = sorted(sidecars[: a.sample])
elif a.limit:
sidecars = sidecars[: a.limit]
if not sidecars:
raise SystemExit(f"no sidecars under {IDENTIFY_DIR}")
# Resume: a full sweep is ~1 h, so never redo work an interrupted run already did.
done: set[str] = set()
if a.restart and a.out.exists():
a.out.unlink() # --restart must TRUNCATE; the file is reopened in append mode below
if a.out.exists() and not a.restart:
with open(a.out, encoding="utf-8") as fh:
for line in fh:
try:
done.add(json.loads(line)["uid"])
except Exception: # noqa: S112 -- tolerate a torn last line from an interrupted run
continue
todo = [s for s in sidecars if os.path.basename(s)[:-5] not in done]
print(f"sidecars {len(sidecars)} already done {len(done)} to do {len(todo)} workers {a.workers}")
counts: collections.Counter[str] = collections.Counter()
a.out.parent.mkdir(parents=True, exist_ok=True)
with open(a.out, "a", encoding="utf-8") as fh, ProcessPoolExecutor(max_workers=a.workers) as ex:
futures = {ex.submit(_one, s): s for s in todo}
for i, fut in enumerate(as_completed(futures), 1):
rec = fut.result()
fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
if rec["status"] != "ok":
counts[rec["status"]] += 1
else:
for c in rec["classes"]:
counts[c] += 1
if i % 500 == 0:
fh.flush()
print(f" {i}/{len(todo)} " + " ".join(f"{k}={v}" for k, v in counts.most_common(6)), flush=True)
print(f"\n{'=' * 70}\nSIDECAR REGRESSION processed={len(todo)}\n{'=' * 70}")
for k, v in counts.most_common():
flag = " <-- REGRESSION" if k.startswith("lost_") or k in ("identify_raised",) else ""
print(f" {k:22s} {v:7d}{flag}")
print(f"\nfull records: {a.out}")
if __name__ == "__main__":
main()
scripts/smoke_matrix.py
+621
View File
@@ -0,0 +1,621 @@
"""Release smoke matrix: exercise every CLI parameter CHOICE on REAL data.
Run before a release. This is not a unit-test substitute -- it drives the real CLI as a
subprocess, so it covers argument parsing, exit-code semantics and the file-writing
contracts that unit tests with fakes cannot.
WHAT IT COVERS AND WHY THAT SHAPE
* Every choice-valued flag at least once (`--backend`, `--sensitivity`, `--mark`,
`--mode`, `--inpaint-method`, ...). Full permutation is combinatorially large and
mostly meaningless; dead branches and typos hide in the CHOICES, not in their
products.
* Every input FORMAT and shape edge case through the pixel paths -- PNG/JPEG/WebP/
HEIC/AVIF, alpha, unicode names, misnamed extensions, truncated files, tiny and
landscape frames. Real-world breakage lives here far more than in flag combos.
* CONTRACTS, not just exit codes: `visible` must write NO output and exit 2 when no
mark is found (re-serving the input reads as success -- the recurring "it didn't
work" report); a no-op must be byte-identical; `metadata --remove` must actually
strip; a JPEG strip must not touch the pixels. Note the pixel-lossless contract is
the DEFAULT path's -- `--remove-all` deliberately re-encodes (see metadata.py).
* The diffusion bodies under `--diffusion`, at `--max-resolution 512` so they fit MPS
(~1 min/image on 32 GB unified memory). Not just exit codes: `invisible` must
restore the input resolution and must NOT re-stamp SDXL's own open watermark.
WHAT IT DOES NOT COVER, DELIBERATELY AND LOUDLY
* Without `--diffusion`, the model-running bodies are reported as SKIPPED with a
reason, never as passes -- a green run that quietly skipped half the surface is
worse than a red one.
* Removal STRENGTH is never certified here. Whether a watermark is actually gone
needs the per-vendor oracles (docs/known-limitations.md); these rows prove the
paths run and keep their contracts, nothing more.
* The re-embed row is gated on a POSITIVE CONTROL. imwatermark is positive-only and
fails to round-trip on some pristine carriers, so "no watermark found" proves
nothing there; the row degrades to a skip rather than a false pass.
uv run python scripts/smoke_matrix.py # corpus + fixtures
uv run python scripts/smoke_matrix.py --quick # fixtures only, no corpus
uv run python scripts/smoke_matrix.py --diffusion # + the SDXL model paths
"""
from __future__ import annotations
import argparse
import json
import random
import shutil
import subprocess
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
SAMPLES = REPO / "data" / "samples"
CORPUS = REPO / "data" / "spaces" / "originals"
EXIT_NO_VISIBLE_MARK = 2
def _capture(args: list[str]) -> str:
"""Run the CLI and return stdout (for the rows that inspect output, not exit code)."""
exe = shutil.which("uv") or "uv"
p = subprocess.run( # noqa: S603
[exe, "run", "remove-ai-watermarks", *args], capture_output=True, text=True, cwd=REPO, check=False
)
return p.stdout
@dataclass
class Result:
name: str
status: str # pass | FAIL | skip
detail: str = ""
cmd: str = ""
@dataclass
class Runner:
tmp: Path
results: list[Result] = field(default_factory=list)
def run(self, name: str, args: list[str], *, expect_exit: int | None = 0, timeout: int = 180) -> Result:
exe = shutil.which("uv") or "uv"
cmd = [exe, "run", "remove-ai-watermarks", *args]
try:
p = subprocess.run( # noqa: S603
cmd, capture_output=True, text=True, timeout=timeout, cwd=REPO, check=False
)
except subprocess.TimeoutExpired:
r = Result(name, "FAIL", f"timeout after {timeout}s", " ".join(args))
self.results.append(r)
return r
ok = expect_exit is None or p.returncode == expect_exit
detail = "" if ok else f"exit {p.returncode} (want {expect_exit}): {(p.stderr or p.stdout).strip()[-200:]}"
r = Result(name, "pass" if ok else "FAIL", detail, " ".join(args))
self.results.append(r)
return r
def check(self, name: str, ok: bool, detail: str = "") -> None:
self.results.append(Result(name, "pass" if ok else "FAIL", "" if ok else detail))
def skip(self, name: str, why: str) -> None:
self.results.append(Result(name, "skip", why))
def corpus_pick(n: int, suffixes: tuple[str, ...]) -> list[Path]:
"""Real uploads, chosen deterministically so a failure is reproducible."""
if not CORPUS.exists():
return []
pool = [p for p in CORPUS.glob("*/*") if p.suffix.lower() in suffixes]
random.Random(7).shuffle(pool) # noqa: S311 -- deterministic sampling, not cryptography
return pool[:n]
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--quick", action="store_true", help="fixtures only; skip the corpus rows")
ap.add_argument(
"--diffusion", action="store_true", help="also run the model-running paths (SDXL weights, ~1 min/image)"
)
a = ap.parse_args()
tmp = Path(tempfile.mkdtemp(prefix="raiw-smoke-"))
r = Runner(tmp)
doubao = SAMPLES / "doubao-1.png"
chatgpt = SAMPLES / "chatgpt-1.png"
grok = SAMPLES / "grok-1.jpg"
# ---- identify: every flag, and JSON must actually parse -------------------
r.run("identify plain", ["identify", str(doubao)])
r.run("identify --no-visible", ["identify", str(doubao), "--no-visible"])
try:
json.loads(_capture(["identify", str(doubao), "--json"]))
r.check("identify --json parses", True)
except Exception as e:
r.check("identify --json parses", False, str(e))
# ---- metadata: check/remove, and the strip must be REAL ------------------
r.run("metadata --check", ["metadata", str(doubao), "--check"])
out = tmp / "meta.png"
r.run("metadata --remove", ["metadata", str(doubao), "--remove", "-o", str(out)])
if out.exists():
try:
rep = json.loads(_capture(["identify", str(out), "--json"]))
# Scope the assert to METADATA. doubao-1 also carries a visible pixel mark,
# which `metadata --remove` must NOT touch -- demanding an empty watermark
# list would fail on correct behavior (and hide a real metadata leak behind
# a permanently-red row).
leftover = [s for s in rep.get("signals", []) if not s.get("name", "").startswith("visible_")]
r.check(
"metadata --remove strips every metadata signal",
not leftover and not rep.get("ai_from_metadata"),
f"still reports metadata signals {[s.get('name') for s in leftover]}",
)
r.check(
"metadata --remove leaves the visible mark alone",
any(s.get("name", "").startswith("visible_") for s in rep.get("signals", [])),
"the pixel mark vanished -- a metadata strip must not touch pixels",
)
except Exception as e:
r.check("metadata --remove strips every metadata signal", False, str(e))
# The pixel-lossless contract belongs to the DEFAULT path. `--remove-all`
# (keep_standard=False) deliberately falls through to the full PIL re-encode
# (metadata.py: the lossless marker walk preserves standard segments, so it cannot
# serve a strip-everything caller), which is lossy for JPEG -- measured ~49 dB.
# Asserting losslessness there tested a contract the code never made.
import numpy as np
from remove_ai_watermarks.image_io import imread
out2 = tmp / "meta2.jpg"
r.run("metadata --remove (jpeg, default)", ["metadata", str(grok), "--remove", "-o", str(out2)])
if out2.exists():
a_, b_ = imread(str(grok)), imread(str(out2))
same = a_ is not None and b_ is not None and a_.shape == b_.shape and bool(np.array_equal(a_, b_))
r.check("jpeg metadata strip is pixel-lossless", same, "pixels changed -- the strip re-encoded")
out3 = tmp / "meta3.jpg"
r.run("metadata --remove --remove-all (jpeg)", ["metadata", str(grok), "--remove", "--remove-all", "-o", str(out3)])
if out3.exists():
rep3 = json.loads(_capture(["identify", str(out3), "--json"]))
r.check(
"--remove-all strips the AI metadata too",
not [s for s in rep3.get("signals", []) if not s.get("name", "").startswith("visible_")],
f"still reports {rep3.get('signals')}",
)
# ---- visible: every --mark choice, every --backend, every --sensitivity ---
for mark in ("auto", "gemini", "doubao", "jimeng", "samsung", "jimeng_pill"):
# doubao-1 carries only the doubao mark, so every other --mark must exit 2
want = 0 if mark in ("auto", "doubao") else EXIT_NO_VISIBLE_MARK
r.run(
f"visible --mark {mark}",
["visible", str(doubao), "--mark", mark, "-o", str(tmp / f"m_{mark}.png")],
expect_exit=want,
)
# Every backend CHOICE, not just the two that need no extra. migan/lama were
# measured at library level (scripts/fill_quality.py) but had never been driven
# through the CLI, which is a different code path (resolve_backend + the warning).
from remove_ai_watermarks.region_eraser import lama_available, migan_available
for backend in ("auto", "cv2", "migan", "lama"):
have = {"migan": migan_available(), "lama": lama_available()}.get(backend, True)
if not have:
r.skip(f"visible --backend {backend}", f"the `{backend}` extra is not installed")
continue
r.run(
f"visible --backend {backend}",
["visible", str(doubao), "--backend", backend, "-o", str(tmp / f"b_{backend}.png")],
timeout=600,
)
for sens in ("auto", "strict"):
r.run(
f"visible --sensitivity {sens}",
["visible", str(doubao), "--sensitivity", sens, "-o", str(tmp / f"s_{sens}.png")],
)
r.run("visible --keep-metadata", ["visible", str(doubao), "--keep-metadata", "-o", str(tmp / "keep.png")])
r.run("visible --no-detect (forced)", ["visible", str(doubao), "--no-detect", "-o", str(tmp / "force.png")])
for removed in ("assume-ai", "aggressive"):
r.run(
f"visible rejects --sensitivity {removed}",
["visible", str(doubao), "--sensitivity", removed],
expect_exit=2,
)
# CONTRACT: no mark -> no output file, exit 2 (never re-serve the input as success)
noout = tmp / "must_not_exist.png"
r.run("visible no-mark exits 2", ["visible", str(chatgpt), "-o", str(noout)], expect_exit=EXIT_NO_VISIBLE_MARK)
r.check("visible no-mark writes NO output", not noout.exists(), "wrote an output for an undetected mark")
# ---- erase: backends, methods, repeated regions, dilate ------------------
for method in ("telea", "ns"):
r.run(
f"erase --inpaint-method {method}",
[
"erase",
str(doubao),
"--region",
"10,10,60,30",
"--inpaint-method",
method,
"-o",
str(tmp / f"e_{method}.png"),
],
)
r.run(
"erase repeated --region",
["erase", str(doubao), "--region", "10,10,40,20", "--region", "80,80,40,20", "-o", str(tmp / "e_multi.png")],
)
r.run(
"erase --dilate",
["erase", str(doubao), "--region", "10,10,40,20", "--dilate", "5", "-o", str(tmp / "e_dil.png")],
)
for backend in ("cv2", "migan", "lama"):
have = {"migan": migan_available(), "lama": lama_available()}.get(backend, True)
if not have:
r.skip(f"erase --backend {backend}", f"the `{backend}` extra is not installed")
continue
r.run(
f"erase --backend {backend}",
[
"erase",
str(doubao),
"--region",
"10,10,40,20",
"--backend",
backend,
"-o",
str(tmp / f"e_{backend}.png"),
],
timeout=600,
)
r.run(
"erase --dilate 0 (no dilation)",
["erase", str(doubao), "--region", "10,10,40,20", "--dilate", "0", "-o", str(tmp / "e_d0.png")],
)
r.run("erase rejects a malformed --region", ["erase", str(doubao), "--region", "not,a,box"], expect_exit=2)
# ---- global + explicit-default flags that had never been exercised -------
r.run("--verbose", ["--verbose", "identify", str(doubao)])
r.run("--version", ["--version"])
r.run("metadata --keep-standard (explicit)", ["metadata", str(doubao), "--check", "--keep-standard"])
r.run(
"visible --strip-metadata (explicit)",
["visible", str(doubao), "--strip-metadata", "-o", str(tmp / "sm.png")],
)
r.run("visible --detect (explicit default)", ["visible", str(doubao), "--detect", "-o", str(tmp / "det.png")])
# ---- batch: the non-diffusion modes -------------------------------------
bd = tmp / "batch_in"
bd.mkdir()
for f in (doubao, chatgpt):
shutil.copy(f, bd / f.name)
for mode in ("visible", "metadata"):
r.run(
f"batch --mode {mode}",
["batch", str(bd), "--mode", mode, "-o", str(tmp / f"batch_{mode}")],
expect_exit=None,
)
# ---- diffusion: argument handling always; the model body under --diffusion ----
r.run("invisible --help parses full knob set", ["invisible", "--help"])
r.run("all --help parses full knob set", ["all", "--help"])
_media_rows(r, tmp)
if a.diffusion:
_diffusion_rows(r, tmp, doubao)
else:
for name in ("invisible", "all", "batch --mode invisible"):
r.skip(f"{name} (model-running body)", "pass --diffusion to exercise it (needs the SDXL weights)")
# ---- real-data formats and shapes ---------------------------------------
if not a.quick:
picks: list[tuple[str, Path]] = []
for suf, label in ((".heic", "heic"), (".avif", "avif"), (".webp", "webp"), (".jpeg", "jpeg")):
picks += [(label, p) for p in corpus_pick(2, (suf,))]
picks += [("png", p) for p in corpus_pick(3, (".png",))]
if not picks:
r.skip("real-format rows", "corpus not present (data/spaces/originals)")
for label, p in picks:
r.run(f"identify real {label}", ["identify", str(p), "--json"])
r.run(
f"visible auto real {label}",
["visible", str(p), "-o", str(tmp / f"r_{label}_{p.stem[:8]}.png")],
expect_exit=None,
) # 0 or 2 are both correct; a CRASH is not
# unicode + misnamed extension + truncated: the documented real-world traps
if picks:
src = picks[0][1]
uni = tmp / "тест изображение 测试.png"
shutil.copy(src, uni)
r.run("unicode filename", ["identify", str(uni), "--json"])
mis = tmp / "actually_png.jpg" # content PNG, extension JPEG
shutil.copy(SAMPLES / "chatgpt-1.png", mis)
r.run("misnamed extension", ["identify", str(mis), "--json"])
trunc = tmp / "truncated.png"
trunc.write_bytes((SAMPLES / "chatgpt-1.png").read_bytes()[:4096])
r.run("truncated file does not crash", ["identify", str(trunc), "--json"], expect_exit=None)
# ---- report --------------------------------------------------------------
bad = [x for x in r.results if x.status == "FAIL"]
skipped = [x for x in r.results if x.status == "skip"]
ok = [x for x in r.results if x.status == "pass"]
print(f"\n{'=' * 74}\nSMOKE MATRIX pass={len(ok)} FAIL={len(bad)} skipped={len(skipped)}\n{'=' * 74}")
for x in skipped:
print(f" SKIP {x.name:46s} {x.detail}")
for x in bad:
print(f" FAIL {x.name:46s} {x.detail}")
if x.cmd:
print(f" cmd: {x.cmd}")
if not bad:
print(" no failures")
print(f"\ntmp artifacts: {tmp}")
raise SystemExit(1 if bad else 0)
def _knob_rows(r: Runner, tmp: Path, img: Path) -> None:
"""Every diffusion knob the matrix never touched.
Deliberately cheap (`--steps 4`, `--max-resolution 384`): these rows answer "is the
knob accepted and does the run complete", NOT "is the output good". Quality per knob
needs a per-knob oracle and most of them have none (`--humanize` has no oracle at
all), so claiming more here would be dishonest.
"""
from remove_ai_watermarks import upscaler
# --steps 20 is the floor that WORKS, not an arbitrary choice: effective timesteps
# are int(steps * strength), so at the default strength 0.15 anything below
# --steps 7 rounds to ZERO and the pipeline dies inside torch. The first version of
# these rows used --steps 4 and every single one failed with the same reshape error.
fast = ["--max-resolution", "384", "--min-resolution", "0", "--steps", "20", "--force", "--seed", "0"]
def run(name: str, extra: list[str], *, tag: str, expect: int | None = 0) -> None:
r.run(
name,
["invisible", str(img), "-o", str(tmp / f"k_{tag}.png"), *fast, *extra],
expect_exit=expect,
timeout=2400,
)
run("--pipeline sdxl", ["--pipeline", "sdxl"], tag="sdxl")
run("--pipeline controlnet", ["--pipeline", "controlnet"], tag="cnet")
run("--strength", ["--strength", "0.2"], tag="strength")
run("--guidance-scale", ["--guidance-scale", "5.0"], tag="gs")
run("--controlnet-scale", ["--controlnet-scale", "0.5"], tag="cns")
run("--humanize", ["--humanize", "0.3"], tag="hum")
run("--unsharp", ["--unsharp", "0.5"], tag="uns")
run("--no-adaptive-polish", ["--no-adaptive-polish"], tag="nap")
run("--tile", ["--tile", "--tile-size", "256", "--tile-overlap", "64"], tag="tile")
run("--device mps", ["--device", "mps"], tag="mps")
run("--upscaler lanczos", ["--upscaler", "lanczos"], tag="lanczos")
run("--auto (deprecated no-op)", ["--auto"], tag="auto")
if upscaler.is_available():
run("--upscaler esrgan", ["--upscaler", "esrgan"], tag="esrgan")
else:
r.skip("--upscaler esrgan", "the `esrgan` extra is not installed")
# CPU is correctness-relevant (it is the documented MPS-OOM fallback) but slow, so
# it gets the smallest possible run rather than being skipped.
r.run(
"--device cpu",
[
"invisible",
str(img),
"-o",
str(tmp / "k_cpu.png"),
"--device",
"cpu",
"--max-resolution",
"256",
"--min-resolution",
"0",
"--steps",
"20",
"--force",
"--seed",
"0",
],
timeout=3600,
)
# qwen is CUDA-class by design (bf16 MMDiT, no MPS fallback). On this host the
# honest outcome is a CLEAN failure, not a crash -- assert it does not hang or
# dump a traceback at the user.
try:
import torch
cuda = bool(torch.cuda.is_available())
except Exception:
cuda = False
if cuda:
run("--pipeline qwen", ["--pipeline", "qwen"], tag="qwen")
else:
r.skip("--pipeline qwen", "CUDA-class pipeline; no CUDA device on this host")
# --model and --hf-token are deliberately not exercised: one would download a second
# multi-GB checkpoint, the other needs a real credential. Skipped loudly, not passed.
r.skip("--model", "would download a second multi-GB checkpoint")
r.skip("--hf-token", "needs a real credential; cannot be exercised meaningfully here")
# CONTRACT, not just execution: the same seed must reproduce the same pixels.
a_out, b_out = tmp / "seed_a.png", tmp / "seed_b.png"
for out in (a_out, b_out):
r.run(
f"seed determinism run ({out.name})",
["invisible", str(img), "-o", str(out), *fast],
timeout=2400,
)
if a_out.exists() and b_out.exists():
import numpy as np
from remove_ai_watermarks.image_io import imread
x, y = imread(str(a_out)), imread(str(b_out))
r.check(
"same --seed reproduces identical pixels",
x is not None and y is not None and x.shape == y.shape and bool(np.array_equal(x, y)),
"two runs with the same seed differed",
)
def _media_rows(r: Runner, tmp: Path) -> None:
"""Audio/video metadata strip via ffmpeg -- a supported path with no corpus coverage.
The corpus is images only, so the media is synthesized here with ffmpeg rather than
left untested.
"""
if not shutil.which("ffmpeg"):
r.skip("audio/video metadata strip", "ffmpeg not on PATH")
return
for name, gen in (
("mp4", ["-f", "lavfi", "-i", "testsrc=duration=1:size=128x128:rate=8", "-pix_fmt", "yuv420p"]),
("mp3", ["-f", "lavfi", "-i", "sine=frequency=440:duration=1"]),
):
src = tmp / f"media.{name}"
ff = shutil.which("ffmpeg") or "ffmpeg"
made = subprocess.run( # noqa: S603
[ff, "-y", *gen, "-metadata", "comment=Made with AI", str(src)],
capture_output=True,
check=False,
)
if made.returncode != 0 or not src.exists():
r.skip(f"{name} metadata strip", "ffmpeg could not synthesize the fixture")
continue
out = tmp / f"media_clean.{name}"
r.run(f"{name} metadata strip runs", ["metadata", str(src), "--remove", "-o", str(out)], expect_exit=None)
if out.exists():
r.check(f"{name} strip produced a non-empty file", out.stat().st_size > 0, "empty output")
def _sdxl_watermark_bits(img: object) -> float:
"""Bits of the open SDXL DWT-DCT watermark recovered from `img` (128 = perfect)."""
import numpy as np
from imwatermark import WatermarkDecoder
truth = np.frombuffer(b"StableDiffusionV1"[:16], dtype=np.uint8)
rec = WatermarkDecoder("bytes", 128).decode(img, "dwtDct")
return float(128 - np.unpackbits(truth ^ np.frombuffer(bytes(rec), dtype=np.uint8)).sum())
def _diffusion_rows(r: Runner, tmp: Path, doubao: Path) -> None:
"""Exercise the model-running bodies at a reduced resolution (MPS-friendly).
Bounded with `--max-resolution 512` and a fixed seed: the point is that the paths
RUN and keep their contracts, not to certify removal strength (that needs the
per-vendor oracles, see docs/known-limitations.md).
"""
import numpy as np
from remove_ai_watermarks.image_io import imread
small = ["--max-resolution", "512", "--seed", "0"]
# `invisible` must restore the original resolution and NOT re-stamp SDXL's own
# open watermark (add_watermarker=False; a remover that re-marks its output is the
# regression this row exists for).
#
# The carrier matters: imwatermark is positive-only and fails to round-trip on some
# pristine images, so an "absent" verdict on a fragile carrier proves NOTHING. mj-1
# is used because it round-trips at 128/128; the control is re-checked on the actual
# OUTPUT and the row degrades to a skip rather than a false pass if it goes fragile.
mj = SAMPLES / "mj-1.png"
inv = tmp / "inv_mj.png"
res = r.run("invisible runs (mps, 512px)", ["invisible", str(mj), "-o", str(inv), "--force", *small], timeout=1800)
if res.status == "pass" and inv.exists():
src_img, out_img = imread(str(mj)), imread(str(inv))
r.check(
"invisible restores the input resolution",
src_img is not None and out_img is not None and src_img.shape == out_img.shape,
f"{None if src_img is None else src_img.shape} -> {None if out_img is None else out_img.shape}",
)
try:
from imwatermark import WatermarkEncoder
enc = WatermarkEncoder()
enc.set_watermark("bytes", b"StableDiffusionV1"[:16])
control = _sdxl_watermark_bits(enc.encode(np.array(out_img).copy(), "dwtDct"))
if control < 118:
r.skip("invisible does not re-embed an SDXL watermark", f"carrier fragile (control {control:.0f}/128)")
else:
bits = _sdxl_watermark_bits(out_img)
r.check(
"invisible does not re-embed an SDXL watermark",
bits < 118,
f"re-embedded: {bits:.0f}/128 recovered (control {control:.0f}/128)",
)
except ImportError:
r.skip("invisible does not re-embed an SDXL watermark", "imwatermark absent (extra `detect`)")
# `all`: every stage must land -- the visible mark AND the metadata both gone.
allout = tmp / "all_out.png"
res = r.run("all runs (mps, 512px)", ["all", str(doubao), "-o", str(allout), *small], timeout=1800)
if res.status == "pass" and allout.exists():
rep = json.loads(_capture(["identify", str(allout), "--json"]))
r.check(
"all clears visible + metadata in one pass",
not rep.get("signals"),
f"still reports {[s.get('name') for s in rep.get('signals', [])]}",
)
_knob_rows(r, tmp, mj)
# `batch --mode invisible`: every input must produce an output (a silent short
# write is the failure this row guards).
bd = tmp / "batch_inv"
bd.mkdir(exist_ok=True)
for f in (SAMPLES / "chatgpt-2.png", mj):
shutil.copy(f, bd / f.name)
bout = tmp / "batch_inv_out"
res = r.run(
"batch --mode invisible runs (mps, 512px)",
["batch", str(bd), "--mode", "invisible", "-o", str(bout), *small],
expect_exit=None,
timeout=3600,
)
if res.status == "pass":
produced = len(list(bout.glob("*"))) if bout.exists() else 0
r.check("batch invisible writes one output per input", produced == 2, f"{produced} outputs for 2 inputs")
# `batch --mode all` -- the only --mode value the matrix never ran.
aout = tmp / "batch_all_out"
r.run(
"batch --mode all runs (mps, 512px)",
["batch", str(bd), "--mode", "all", "-o", str(aout), *small],
expect_exit=None,
timeout=3600,
)
# The AI-enhanced composite path: regenerate ONLY a region and feather it back,
# leaving everything outside the box pixel-exact. Library-level -- the CLI has no
# flag for it, so it would otherwise never be exercised on real data.
try:
import numpy as np
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.noai.watermark_remover import WatermarkRemover
src = imread(str(mj))
h, w = src.shape[:2]
box = (w // 4, h // 4, w // 4, h // 4)
rem = WatermarkRemover(pipeline="controlnet")
rout = tmp / "region_composite.png"
rem.remove_watermark(mj, rout, strength=0.15, num_inference_steps=20, seed=0, region=box)
got = imread(str(rout))
if got is None or got.shape != src.shape:
r.check("region composite keeps the frame outside the box", False, "shape changed or unreadable")
else:
mask = np.ones(src.shape[:2], dtype=bool)
x, y, bw, bh = box
# Outside the box PLUS the feather margin must be untouched.
pad = 96
mask[max(0, y - pad) : y + bh + pad, max(0, x - pad) : x + bw + pad] = False
r.check(
"region composite keeps the frame outside the box",
bool(np.array_equal(src[mask], got[mask])),
"pixels changed outside the regenerated region",
)
except Exception as e:
r.skip("region composite (remove_watermark(region=...))", f"{type(e).__name__}: {e}"[:120])
if __name__ == "__main__":
main()
scripts/visible_positives.py
+141
View File
@@ -0,0 +1,141 @@
"""Parallel detection pass: list every corpus image carrying a known visible mark.
Why this exists separately from `visible_removal_audit.py`: that audit is single-process,
so a full-corpus sweep costs ~10 h and running it once per backend costs ~30 h. But its
expensive half is DETECTION, and detection does not depend on the fill backend. Splitting
it out means detecting once in parallel and then feeding the positives to the audit via
its `--paths-file` seam, over a few thousand images instead of forty thousand.
CRASH TOLERANCE IS NOT OPTIONAL AT THIS SCALE
cv2/libpng decode native-crash on some real uploads. A plain `ProcessPoolExecutor.map`
over 39k files then DEADLOCKS: the worker dies without a Python traceback and the parent
waits forever on a result that never arrives (observed 2026-07-19 -- 26 min of work lost
because results were only written at the end). So this script:
* writes every result to JSONL as it arrives -- a kill never costs more than a batch;
* is resumable -- an interrupted run skips what it already recorded;
* runs a FRESH pool per batch with a timeout, so one poisoned file costs one batch,
and that batch is retried serially to find and record the offender.
Corpus images are user uploads: read-only, local analysis, gitignored output.
uv run python scripts/visible_positives.py --jobs 6
"""
from __future__ import annotations
import argparse
import collections
import glob
import json
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from concurrent.futures import TimeoutError as FutureTimeout
from concurrent.futures.process import BrokenProcessPool
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
# The package's own format set. An inlined copy here silently skipped .heif, which
# CLAUDE.md documents as supported.
from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS as _EXTS
REPO = Path(__file__).resolve().parents[1]
CORPUS = REPO / "data" / "spaces" / "originals"
OUT = REPO / "data" / "spaces" / "_visible_positives.jsonl"
PATHS = REPO / "data" / "spaces" / "_visible_positives.txt"
def _one(path: str) -> dict[str, object]:
from remove_ai_watermarks.image_io import imread
from remove_ai_watermarks.watermark_registry import detect_marks
try:
img = imread(path)
if img is None:
return {"path": path, "keys": [], "status": "unreadable"}
return {"path": path, "keys": [d.key for d in detect_marks(img) if d.detected], "status": "ok"}
except Exception as e:
return {"path": path, "keys": [], "status": f"error:{type(e).__name__}"}
def _run_batch(batch: list[str], jobs: int, timeout: int) -> list[dict[str, object]]:
"""One batch in a fresh pool. On a native worker crash or timeout, retry serially so
the poisoned file is identified and recorded instead of stalling the whole sweep."""
try:
with ProcessPoolExecutor(max_workers=jobs) as ex:
futs = {ex.submit(_one, p): p for p in batch}
return [f.result() for f in as_completed(futs, timeout=timeout)]
except (FutureTimeout, BrokenProcessPool, OSError, RuntimeError):
out: list[dict[str, object]] = []
for p in batch:
try:
out.append(_one(p))
except BaseException: # a native crash here kills only this file, not the sweep
out.append({"path": p, "keys": [], "status": "crashed"})
return out
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--jobs", type=int, default=max(1, (os.cpu_count() or 4) - 2))
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--batch", type=int, default=200)
ap.add_argument("--timeout", type=int, default=600, help="seconds per batch before falling back to serial")
ap.add_argument("--out", type=Path, default=OUT)
ap.add_argument("--restart", action="store_true")
a = ap.parse_args()
files = sorted(p for p in glob.glob(str(CORPUS / "*" / "*")) if Path(p).suffix.lower() in _EXTS)
if a.limit:
files = files[: a.limit]
done: set[str] = set()
if a.restart and a.out.exists():
a.out.unlink() # --restart must TRUNCATE; the file is reopened in append mode below
if a.out.exists() and not a.restart:
with open(a.out, encoding="utf-8") as fh:
for line in fh:
try:
done.add(json.loads(line)["path"])
except Exception: # noqa: S112 -- tolerate a torn last line
continue
todo = [p for p in files if p not in done]
print(f"images {len(files)} already done {len(done)} to do {len(todo)} jobs {a.jobs}", flush=True)
counts: collections.Counter[str] = collections.Counter()
status: collections.Counter[str] = collections.Counter()
seen = 0
with open(a.out, "a", encoding="utf-8") as fh:
for start in range(0, len(todo), a.batch):
for rec in _run_batch(todo[start : start + a.batch], a.jobs, a.timeout):
fh.write(json.dumps(rec) + "\n")
status[str(rec["status"])] += 1
for k in rec["keys"]: # type: ignore[union-attr]
counts[str(k)] += 1
seen += 1
fh.flush()
print(f" {seen}/{len(todo)} {dict(counts)} {dict(status)}", flush=True)
# Rebuild the paths file from the FULL record set, not just this run's slice.
hits = []
with open(a.out, encoding="utf-8") as fh:
for line in fh:
try:
r = json.loads(line)
except Exception: # noqa: S112
continue
if r.get("keys"):
hits.append(r["path"])
# Derive the paths file from --out so a trial run with a scratch --out cannot
# overwrite the shared list a full sweep produced.
paths_out = a.out.with_suffix(".txt")
paths_out.write_text("\n".join(sorted(set(hits))) + "\n", encoding="utf-8")
print(f"\npositives: {len(set(hits))} images")
for k, v in counts.most_common():
print(f" {v:6d} {k}")
print(f"statuses: {dict(status)}\npaths -> {paths_out}\nrecords -> {a.out}")
if __name__ == "__main__":
main()