From 7f2285259395905e21084228fd3da6b9a7b23e5e Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Tue, 14 Jul 2026 17:16:36 +0300 Subject: [PATCH] fix(metadata): route strip output format by content, not extension remove_ai_metadata chose its save format (and the lossless-JPEG fast path) from the OUTPUT file extension. On the ~2% of real uploads whose extension lies about their content (a PNG served as .jpg is the common case, ~0.9% of the corpus), the default flow -- which inherits the source's own extension -- re-encoded a lossless PNG/WebP into a real JPEG, silently degrading the pixels and breaking the "work with originals" invariant. Sniff the actual container from magic bytes (_sniff_image_format, reusing the 12-byte head already read for the ISOBMFF check) and route on content: a misnamed lossless source (source-extension format != content) is preserved in its true format, while a correctly-named source still honors a deliberate output-extension conversion (source.png -> output.jpg). The JPEG-lossless gate is likewise content-gated. Found by a new metadata-removal parity audit over the local corpus (scripts/metadata_removal_audit.py): 18170/18173 carriers strip cleanly, and this fix takes the 208 pixel-integrity failures (all misnamed PNGs) to 0. Co-Authored-By: Claude Opus 4.8 --- scripts/metadata_removal_audit.py | 256 +++++++++++++++++++++++++++ src/remove_ai_watermarks/metadata.py | 55 ++++-- tests/test_metadata.py | 31 ++++ 3 files changed, 329 insertions(+), 13 deletions(-) create mode 100644 scripts/metadata_removal_audit.py diff --git a/scripts/metadata_removal_audit.py b/scripts/metadata_removal_audit.py new file mode 100644 index 0000000..de2a0f4 --- /dev/null +++ b/scripts/metadata_removal_audit.py @@ -0,0 +1,256 @@ +"""Audit AI-metadata REMOVAL over a local image corpus (detection<->removal parity). + +`corpus_gap_scan.py` proves the DETECTOR sees a marker; this proves the STRIPPER +reaches it. For every file that carries an AI-metadata signal, run +``remove_ai_metadata`` and re-scan the output with the SAME oracle +(``get_ai_metadata``). Any signal that survives is a real parity bug: a re-served +file still reads as AI. Also assert the strip is lossless -- decoded pixels +(and alpha) bit-identical before/after -- since the removal must only touch +metadata, never the coded image. + +A no-op control set (clean images with no AI metadata) verifies the stripper +neither ADDS a signal nor corrupts pixels on files it should leave alone. + +Operates on gitignored data only (data/spaces/...); writes nothing tracked. + + uv run python scripts/metadata_removal_audit.py \ + --corpus data/spaces/originals --identify data/spaces/identify \ + --out data/spaces/_metadata_removal_audit.csv --jobs 8 +""" + +from __future__ import annotations + +import csv +import json +import logging +import random +import tempfile +from collections import Counter +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +import click +import numpy as np + +from remove_ai_watermarks.noai.constants import SUPPORTED_FORMATS + +log = logging.getLogger(__name__) + +# identify-JSON watermark substrings that imply a METADATA-borne signal (as +# opposed to a purely visual sparkle/text mark). Used only to pick the candidate +# population fast; get_ai_metadata is the per-file ground truth. +_META_HINTS = ( + "C2PA", + "Content Credentials", + "IPTC", + "Made with AI", + "AIGC", + "TC260", + "EXIF", + "Signature", + "SynthID", + "hf-job", + "HuggingFace", + "Samsung", + "soft-binding", + "metadata", +) + + +def _pixels(path: Path) -> tuple[np.ndarray | None, np.ndarray | None]: + """Decoded BGR + alpha, for the lossless-strip integrity check.""" + from remove_ai_watermarks import image_io + + return image_io.read_bgr_and_alpha(path) + + +def _same_pixels(a: Path, b: Path) -> bool | None: + """True/False if both decode; None if either is undecodable (can't compare).""" + try: + bgr_a, al_a = _pixels(a) + bgr_b, al_b = _pixels(b) + except Exception: + return None + if bgr_a is None or bgr_b is None: + return None + if bgr_a.shape != bgr_b.shape or not np.array_equal(bgr_a, bgr_b): + return False + if (al_a is None) != (al_b is None): + return False + return al_a is None or al_b is None or np.array_equal(al_a, al_b) + + +def _audit_one(path_str: str) -> dict[str, object]: + """Worker: detect -> strip -> re-detect + pixel-integrity for one file.""" + from remove_ai_watermarks.metadata import get_ai_metadata, remove_ai_metadata + + path = Path(path_str) + row: dict[str, object] = { + "path": path.name, + "ext": path.suffix.lower(), + "carrier": False, + "before": "", + "after": "", + "parity_ok": "", + "pixels_identical": "", + "status": "ok", + } + try: + before = get_ai_metadata(path) + except Exception as exc: + row["status"] = f"scan_error:{type(exc).__name__}" + return row + row["before"] = "|".join(sorted(before)) + row["carrier"] = bool(before) + + with tempfile.TemporaryDirectory() as td: + out = Path(td) / f"clean{path.suffix.lower()}" + try: + remove_ai_metadata(path, out) + except Exception as exc: + row["status"] = f"strip_error:{type(exc).__name__}" + return row + if not out.exists(): + row["status"] = "no_output" + return row + try: + after = get_ai_metadata(out) + except Exception as exc: + row["status"] = f"rescan_error:{type(exc).__name__}" + return row + row["after"] = "|".join(sorted(after)) + row["parity_ok"] = not after # every AI signal must be gone + same = _same_pixels(path, out) + row["pixels_identical"] = "" if same is None else same + return row + + +def _candidate_paths(corpus: Path, identify: Path | None, clean_sample: int) -> tuple[list[Path], list[Path]]: + """Return (carriers, clean_controls). Uses identify JSONs when present to pick + metadata carriers fast; falls back to scanning every file.""" + if identify is None or not identify.exists(): + files = sorted(p for p in corpus.rglob("*") if p.is_file() and p.suffix.lower() in SUPPORTED_FORMATS) + return files, [] + + carriers: list[Path] = [] + clean: list[Path] = [] + for jf in identify.rglob("*.json"): + try: + d = json.loads(jf.read_text()) + except Exception: # noqa: S112 -- skip an unreadable identify JSON, not security-relevant + continue + src = d.get("src") + if not src: + continue + img = corpus / jf.parent.name / src + if not img.exists() or img.suffix.lower() not in SUPPORTED_FORMATS: + continue + wm = " | ".join(d.get("watermarks") or []) + if any(h in wm for h in _META_HINTS): + carriers.append(img) + elif not d.get("is_ai_generated"): + clean.append(img) + rng = random.Random(0) # noqa: S311 -- deterministic sampling seed, not cryptographic + rng.shuffle(clean) + return carriers, clean[:clean_sample] + + +@click.command() +@click.option( + "--corpus", type=click.Path(exists=True, file_okay=False, path_type=Path), default=Path("data/spaces/originals") +) +@click.option( + "--identify", + type=click.Path(path_type=Path), + default=Path("data/spaces/identify"), + help="identify-JSON dir to pick carriers (skip = scan all).", +) +@click.option("--out", type=click.Path(path_type=Path), default=Path("data/spaces/_metadata_removal_audit.csv")) +@click.option( + "--clean-sample", type=int, default=1500, help="No-op control: N clean images to prove the strip is a no-op." +) +@click.option("--limit", type=int, default=0, help="Cap carriers scanned (0 = all).") +@click.option("--jobs", type=int, default=8) +def main(corpus: Path, identify: Path | None, out: Path, clean_sample: int, limit: int, jobs: int) -> None: + logging.basicConfig(level=logging.ERROR, format="%(message)s") + carriers, clean = _candidate_paths(corpus, identify, clean_sample) + if limit: + carriers = carriers[:limit] + tasks = [(p, "carrier") for p in carriers] + [(p, "clean") for p in clean] + click.echo(f"Carriers: {len(carriers)} | clean controls: {len(clean)} | jobs {jobs}") + + rows: list[dict[str, object]] = [] + with ProcessPoolExecutor(max_workers=jobs) as ex: + futs = {ex.submit(_audit_one, str(p)): k for p, k in tasks} + for done, fut in enumerate(as_completed(futs), 1): + row = fut.result() + row["kind"] = futs[fut] + rows.append(row) + if done % 500 == 0: + click.echo(f" {done}/{len(tasks)}") + + out.parent.mkdir(parents=True, exist_ok=True) + fields = ["kind", "path", "ext", "carrier", "before", "after", "parity_ok", "pixels_identical", "status"] + with out.open("w", newline="") as f: + w = csv.DictWriter(f, fieldnames=fields) + w.writeheader() + for r in rows: + w.writerow({k: r.get(k, "") for k in fields}) + + # ---- Summary ---- + car = [r for r in rows if r["kind"] == "carrier"] + ctl = [r for r in rows if r["kind"] == "clean"] + real_carriers = [r for r in car if r["carrier"] and r["status"] == "ok"] + parity_fail = [r for r in real_carriers if r["parity_ok"] is False] + pixel_fail = [r for r in real_carriers if r["pixels_identical"] is False] + errors = [r for r in rows if r["status"] != "ok"] + + click.echo("\n===== METADATA REMOVAL PARITY =====") + click.echo(f"Carrier candidates: {len(car)}") + click.echo(f"Confirmed carriers: {len(real_carriers)} (get_ai_metadata non-empty)") + click.echo(f"Parity FAILS (signal survives strip): {len(parity_fail)}") + click.echo(f"Pixel-integrity FAILS (strip altered pixels): {len(pixel_fail)}") + click.echo(f"Errors (scan/strip/decode): {len(errors)}") + + # Per-signal parity breakdown. + sig_total: Counter[str] = Counter() + sig_fail: Counter[str] = Counter() + for r in real_carriers: + for s in str(r["before"]).split("|"): + if s: + sig_total[s] += 1 + if r["parity_ok"] is False: + for s in str(r["after"]).split("|"): + if s: + sig_fail[s] += 1 + click.echo("\nPer-signal (carriers / surviving-after-strip):") + for s, n in sig_total.most_common(): + click.echo(f" {s:24} {n:6} survived: {sig_fail.get(s, 0)}") + + click.echo("\n===== NO-OP CONTROL (clean images) =====") + ctl_ok = [r for r in ctl if r["status"] == "ok"] + added = [r for r in ctl_ok if r["after"]] # strip must not ADD a signal + corrupted = [r for r in ctl_ok if r["pixels_identical"] is False] + click.echo(f"Clean controls scanned: {len(ctl_ok)}") + click.echo(f"Strip ADDED a signal: {len(added)}") + click.echo(f"Strip corrupted pixels: {len(corrupted)}") + + if parity_fail: + click.echo("\n--- Parity failures (first 30) ---") + for r in parity_fail[:30]: + click.echo(f" {r['ext']:6} survived=[{r['after']}] {r['path']}") + if pixel_fail: + click.echo("\n--- Pixel-integrity failures (first 30) ---") + for r in pixel_fail[:30]: + click.echo(f" {r['ext']:6} {r['path']}") + if errors: + ec: Counter[str] = Counter(str(r["status"]) for r in errors) + click.echo("\n--- Errors by kind ---") + for k, n in ec.most_common(): + click.echo(f" {n:5} {k}") + + click.echo(f"\nReport: {out}") + + +if __name__ == "__main__": + main() diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index a368c3b..8bb8b33 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -1054,6 +1054,28 @@ def _strip_jpeg_metadata_lossless(source_path: Path, output_path: Path) -> bool: return True +# Fallback extension -> PIL save format, used only when the content sniff is +# inconclusive (never for JPEG re-encode of lossless content). +_EXT_TO_PIL_FORMAT = {".jpg": "JPEG", ".jpeg": "JPEG", ".webp": "WEBP", ".png": "PNG"} + + +def _sniff_image_format(head: bytes) -> str | None: + """Actual raster format from a file's leading magic bytes (>= 12 bytes), as a PIL + format name ("JPEG"/"PNG"/"WEBP"), or None when unrecognized. The file EXTENSION is + unreliable: ~2% of real uploads carry a mismatched one (a PNG served as ``.jpg`` is + common). Choosing the save format by extension re-encodes a lossless PNG/WebP into a + real JPEG, silently degrading the pixels -- so the strip routes on content instead. + ISOBMFF/GIF are handled before this point or fall through to PNG; only the + lossy-vs-lossless distinction that matters here is resolved.""" + if head[:2] == b"\xff\xd8": + return "JPEG" + if head[:8] == b"\x89PNG\r\n\x1a\n": + return "PNG" + if head[:4] == b"RIFF" and head[8:12] == b"WEBP": + return "WEBP" + return None + + def remove_ai_metadata( source_path: Path, output_path: Path | None = None, @@ -1121,6 +1143,12 @@ def remove_ai_metadata( if source_path.suffix.lower() in _FFMPEG_STRIP_EXTS: return _strip_with_ffmpeg(source_path, output_path) + # Route on the ACTUAL content format, not the extension (which lies on ~2% of real + # uploads -- a PNG served as .jpg, etc.). Trusting the extension would push a + # lossless PNG/WebP through the lossy JPEG re-encode below just because its name + # ends .jpg, breaking the "work with originals" invariant. + true_fmt = _sniff_image_format(head) # reuse the 12 bytes already read above + # JPEG: strip AI metadata at the byte level so the DCT scan (the pixels) is NOT # re-encoded. The PIL open+save path below is lossy for JPEG (a q95 re-encode that # would undo the quality-preserving writes of the removal pipelines); this keeps a @@ -1128,11 +1156,7 @@ def remove_ai_metadata( # non-parseable JPEG. Only when keep_standard: the lossless walk drops AI segments # but preserves standard ones, so a keep_standard=False caller (strip EVERYTHING) # must use the full re-encode path below instead. - if ( - keep_standard - and output_path.suffix.lower() in (".jpg", ".jpeg") - and _strip_jpeg_metadata_lossless(source_path, output_path) - ): + if keep_standard and true_fmt == "JPEG" and _strip_jpeg_metadata_lossless(source_path, output_path): return output_path # Fail-safe for a truncated / corrupt image: PIL raises OSError when it decodes a @@ -1155,11 +1179,19 @@ def remove_ai_metadata( # Read image and filter metadata with Image.open(source_path) as img: img = img.copy() - fmt = output_path.suffix.lower() + # Pick the save format. Honor the caller's output extension (so a deliberate + # source.png -> output.jpg conversion still works) UNLESS the SOURCE is misnamed + # -- a lossless PNG/WebP whose extension lies (served as .jpg). There the output + # extension only inherited the source's wrong name, so re-encoding to JPEG would + # silently degrade an original; preserve the true content format instead. + source_ext_fmt = _EXT_TO_PIL_FORMAT.get(source_path.suffix.lower()) + if true_fmt is not None and true_fmt != source_ext_fmt: + fmt = true_fmt # misnamed source: never let a lying extension force a re-encode + else: + fmt = _EXT_TO_PIL_FORMAT.get(output_path.suffix.lower()) or true_fmt or "PNG" - save_kwargs: dict[str, Any] = {} - if fmt in (".jpg", ".jpeg"): - save_kwargs["format"] = "JPEG" + save_kwargs: dict[str, Any] = {"format": fmt} + if fmt == "JPEG": # JPEG output is unavoidably lossy, so minimize the loss: high quality # and no chroma subsampling (4:4:4). Without these PIL defaults to # quality 75 + 4:2:0, which visibly degrades a re-saved image. @@ -1167,15 +1199,12 @@ def remove_ai_metadata( save_kwargs["subsampling"] = 0 if img.mode in ("RGBA", "P"): img = img.convert("RGB") - elif fmt == ".webp": + elif fmt == "WEBP": # Preserve the WebP container losslessly instead of silently rewriting # it as PNG (which changes the format and bloats the file). - save_kwargs["format"] = "WEBP" save_kwargs["lossless"] = True if img.mode == "P": # WebP cannot encode palette mode img = img.convert("RGBA" if "transparency" in img.info else "RGB") - else: - save_kwargs["format"] = "PNG" # Collect non-AI metadata kept_meta: dict[str, str] = {} diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 7f55724..926f540 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -179,6 +179,37 @@ class TestHasAiMetadata: assert np.array_equal(before, after), f"{name}: pixels changed (DCT was re-encoded)" assert not has_ai_metadata(out), f"{name}: AI metadata survived the strip" + def test_strip_preserves_lossless_content_with_mismatched_extension(self, tmp_path: Path): + """F1 regression: the save format is chosen by CONTENT, not the file extension. + A PNG served with a .jpg name (common on real uploads -- ~1% of the corpus is a + PNG/WebP under a .jpg extension) must be stripped losslessly as PNG, NOT + re-encoded into a real JPEG. The extension-driven path silently degraded it, + breaking the 'work with originals' invariant.""" + import numpy as np + from PIL import Image + from PIL.PngImagePlugin import PngInfo + + from remove_ai_watermarks import image_io + from remove_ai_watermarks.metadata import remove_ai_metadata + + # Random noise: JPEG (even q95 4:4:4) provably shifts it; PNG stays exact. + arr = np.random.default_rng(0).integers(0, 256, (64, 64, 3), dtype=np.uint8) + info = PngInfo() + info.add_text("parameters", "Stable Diffusion prompt") # an AI-provenance key + src = tmp_path / "actually_png.jpg" # PNG bytes under a .jpg name + Image.fromarray(arr).save(src, format="PNG", pnginfo=info) + assert has_ai_metadata(src) + + out = tmp_path / "cleaned.jpg" + remove_ai_metadata(src, out) + + before = image_io.imread(str(src)) + after = image_io.imread(str(out)) + assert before is not None + assert after is not None + assert np.array_equal(before, after), "misnamed PNG was lossily re-encoded to JPEG" + assert not has_ai_metadata(out), "AI metadata survived the strip" + @staticmethod def _xmp_iptc_jpeg(tmp_path: Path, name: str, marker: bytes) -> Path: """A real (decodable) JPEG carrying the IPTC AI marker in a well-formed APP1