mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-06 22:18:36 +02:00
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 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
d9204e2e1f
commit
7f22852593
@@ -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()
|
||||
Reference in New Issue
Block a user