mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-10 08:00:32 +02:00
Add a portable metadata record so collection and verdict can run apart
`collect_metadata_record` returns a JSON-safe record carrying an image's provenance metadata regions -- never its pixels -- and the existing `evidence_from_metadata_record` + `identify_from_evidence` build the verdict from it without opening the file. The contract is equality with `identify(path, metadata only)`, verified over the tracked fixtures and over a local corpus of 3,478 images (every file carrying a rare signal, plus a random slice): zero differences. Three placements defeated earlier drafts and each is now a rule with a test: the `scan_head` buffer is the head CONCATENATED with late metadata, so a structural walk must read the raw head instead; Samsung splits its evidence between a post-EOI trailer and the coded scan; and PIL's info keys must be emitted in the file path's candidate order, since the first token match wins. Also fix a real detection gap found while establishing that equality: a label the decoder can read but a raw byte scan cannot -- a compressed PNG `zTXt` packet, or a WebP XMP chunk past the scan window -- was invisible to `identify`. Eight corpus files carrying a China TC260 AIGC label or an IPTC "Made with AI" tag were reported as no signal at all. `scripts/detection_timing.py` and its report script measure the metadata path per method; they write outside the repository and are read-only over a dataset. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
f481e6f944
commit
0c5961a0ed
@@ -0,0 +1,239 @@
|
||||
"""Per-method wall time for metadata extraction and the verdict built on it.
|
||||
|
||||
WHY THIS EXISTS
|
||||
The detection path is documented by CAPABILITY -- which signals it reads, in what
|
||||
order, with what confidence -- but not by COST. Batch throughput and the question
|
||||
"which probe dominates, and does it depend on container or on file size" have never
|
||||
been measured.
|
||||
|
||||
WHAT IT MEASURES
|
||||
Only the file-backed half: metadata extraction, then the verdict evaluated on the
|
||||
extracted evidence. The pixel-backed detectors are deliberately out of scope.
|
||||
|
||||
COLD pass: one measurement on a file the process has never touched --
|
||||
``extract_provenance_evidence``. Only the first read of a file is genuinely cold,
|
||||
so it buys exactly one number, and that is the one a single-shot run pays.
|
||||
|
||||
WARM pass, with the filesystem cache now hot:
|
||||
1. ``extract_provenance_evidence`` as a whole.
|
||||
2. Its nine components, timed IN THE ORDER the dataclass constructs them and with
|
||||
the per-file caches cleared once beforehand -- so ``extract_c2pa_info`` carries
|
||||
the Rust manifest reader and the later ``get_ai_metadata`` sees the same warm
|
||||
cache it sees in production. Timing them in any other order moves that cost to
|
||||
a different row and flatters whichever ran second.
|
||||
3. ``identify_from_evidence``: pure verdict logic, the source is never reopened.
|
||||
4. ``identify(check_visible=False, check_invisible=False)`` -- extraction plus
|
||||
verdict as one call, the cross-check that 1 + 3 is the whole metadata path.
|
||||
|
||||
Every ``@lru_cache`` in the metadata and C2PA modules is cleared before each timed
|
||||
unit. Without that the second measurement of a file answers from the memo and
|
||||
reports a cost of zero -- the caches are keyed on (path, mtime, size) and this
|
||||
script reads each file several times.
|
||||
|
||||
READING IT
|
||||
Component times do NOT sum to the ``extract_provenance_evidence`` total for free:
|
||||
they come from a separate cache-cleared run, so the sum is a cross-check. A gap
|
||||
means a component is missing from the list. Both numbers are written.
|
||||
|
||||
DATA SAFETY
|
||||
Read-only over a local dataset. Writes only to the given output prefix, which
|
||||
belongs outside the repository.
|
||||
|
||||
uv run python scripts/detection_timing.py data/spaces/originals .local-eval/timing/run
|
||||
uv run python scripts/detection_timing.py <dataset> <prefix> --limit 200
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Iterator
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
|
||||
from remove_ai_watermarks import identify as identify_mod
|
||||
from remove_ai_watermarks import metadata as metadata_mod
|
||||
from remove_ai_watermarks._internal import c2pa as c2pa_mod
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED = frozenset({".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif", ".avif"})
|
||||
|
||||
# Timed in the order ``ProvenanceEvidence`` is constructed in ``identify.py``. The
|
||||
# report script imports this to label its columns, so the order is the contract.
|
||||
COMPONENTS: tuple[tuple[str, Callable[[Path], Any]], ...] = (
|
||||
("c2pa_info", c2pa_mod.extract_c2pa_info),
|
||||
("ai_metadata", metadata_mod.get_ai_metadata),
|
||||
("scan_head", lambda p: metadata_mod.scan_head(p, identify_mod._SCAN_BYTES)),
|
||||
("iptc_ai_system", metadata_mod.iptc_ai_system),
|
||||
("aigc_label", metadata_mod.aigc_label),
|
||||
("exif_generator", metadata_mod.exif_generator),
|
||||
("xai_signature", metadata_mod.xai_signature),
|
||||
("huggingface_job", metadata_mod.huggingface_job),
|
||||
("samsung_genai", metadata_mod.samsung_genai),
|
||||
)
|
||||
|
||||
|
||||
def _cache_clearers() -> tuple[Callable[[], None], ...]:
|
||||
"""Every per-file memo in the metadata path, found by attribute, not by name list.
|
||||
|
||||
A hand-written list silently goes stale the next time a probe gains a cache, and a
|
||||
stale entry shows up as a suspiciously fast row rather than as an error.
|
||||
"""
|
||||
found: list[Callable[[], None]] = []
|
||||
for module in (metadata_mod, c2pa_mod):
|
||||
for name in dir(module):
|
||||
clear = getattr(getattr(module, name, None), "cache_clear", None)
|
||||
if callable(clear):
|
||||
found.append(clear)
|
||||
return tuple(found)
|
||||
|
||||
|
||||
_CLEARERS = _cache_clearers()
|
||||
|
||||
|
||||
def _clear() -> None:
|
||||
for clear in _CLEARERS:
|
||||
clear()
|
||||
|
||||
|
||||
def _ms(fn: Callable[[], Any]) -> tuple[float, Any]:
|
||||
"""Wall time in milliseconds plus the call's result."""
|
||||
start = time.perf_counter_ns()
|
||||
value = fn()
|
||||
return (time.perf_counter_ns() - start) / 1e6, value
|
||||
|
||||
|
||||
def _pixel_geometry(path: Path) -> tuple[str | None, int | None, int | None]:
|
||||
"""Container format and pixel dimensions from the header alone."""
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(path) as img:
|
||||
return img.format, img.width, img.height
|
||||
except Exception: # unreadable or an unsupported container
|
||||
return None, None, None
|
||||
|
||||
|
||||
def _warm_breakdown(path: Path, row: dict[str, Any]) -> None:
|
||||
"""Fill ``row`` with the warm-cache per-method breakdown."""
|
||||
_clear()
|
||||
row["extract_evidence_ms"], evidence = _ms(lambda: identify_mod.extract_provenance_evidence(path))
|
||||
|
||||
_clear()
|
||||
component_sum = 0.0
|
||||
for name, fn in COMPONENTS:
|
||||
elapsed, _ = _ms(lambda fn=fn: fn(path))
|
||||
row[f"meta_{name}_ms"] = elapsed
|
||||
component_sum += elapsed
|
||||
row["meta_components_sum_ms"] = component_sum
|
||||
|
||||
row["verdict_from_evidence_ms"], report = _ms(lambda: identify_mod.identify_from_evidence(evidence))
|
||||
|
||||
_clear()
|
||||
row["identify_metadata_only_ms"], full = _ms(
|
||||
lambda: identify_mod.identify(path, check_visible=False, check_invisible=False)
|
||||
)
|
||||
|
||||
row["scan_bytes"] = len(evidence.scan)
|
||||
row["has_c2pa"] = bool(evidence.c2pa_info)
|
||||
row["has_ai_metadata"] = bool(evidence.ai_metadata)
|
||||
row["is_ai_generated"] = full.is_ai_generated
|
||||
row["confidence"] = full.confidence
|
||||
row["platform"] = full.platform
|
||||
row["signals"] = [signal.name for signal in full.signals]
|
||||
# The two verdict paths must agree; a mismatch means the breakdown timed a
|
||||
# different code path than the end-to-end call and the rows are not comparable.
|
||||
row["verdict_agrees"] = (report.confidence, report.platform) == (full.confidence, full.platform)
|
||||
|
||||
|
||||
def _measure(path: Path) -> dict[str, Any]:
|
||||
row: dict[str, Any] = {"path": str(path), "ext": path.suffix.lower()}
|
||||
try:
|
||||
row["bytes"] = path.stat().st_size
|
||||
except OSError as exc:
|
||||
return {**row, "error": f"stat: {exc}"}
|
||||
|
||||
# COLD first: this is the only moment the file is untouched by this process.
|
||||
_clear()
|
||||
try:
|
||||
row["cold_extract_evidence_ms"], _ = _ms(lambda: identify_mod.extract_provenance_evidence(path))
|
||||
except Exception as exc:
|
||||
return {**row, "error": f"cold extract: {type(exc).__name__}: {exc}"}
|
||||
|
||||
row["format"], row["width"], row["height"] = _pixel_geometry(path)
|
||||
width, height = row["width"], row["height"]
|
||||
row["megapixels"] = round(width * height / 1e6, 3) if width and height else None
|
||||
|
||||
try:
|
||||
_warm_breakdown(path, row)
|
||||
except Exception as exc:
|
||||
row["error"] = f"warm pass: {type(exc).__name__}: {exc}"
|
||||
return row
|
||||
|
||||
|
||||
def _iter_images(root: Path) -> Iterator[Path]:
|
||||
for path in sorted(root.rglob("*")):
|
||||
if path.is_file() and path.suffix.lower() in SUPPORTED:
|
||||
yield path
|
||||
|
||||
|
||||
def _done_paths(out_path: Path) -> set[str]:
|
||||
"""Paths already recorded, so a long run resumes instead of restarting."""
|
||||
if not out_path.exists():
|
||||
return set()
|
||||
done: set[str] = set()
|
||||
with out_path.open(encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
try:
|
||||
done.add(json.loads(line)["path"])
|
||||
except (ValueError, KeyError):
|
||||
continue
|
||||
return done
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("dataset", type=Path, help="directory of images, scanned recursively")
|
||||
parser.add_argument("out_prefix", type=Path, help="output prefix; writes <prefix>.jsonl")
|
||||
parser.add_argument("--limit", type=int, default=0, help="stop after N files (0 = all)")
|
||||
parser.add_argument("--progress-every", type=int, default=200)
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
|
||||
out_path = args.out_prefix.with_suffix(".jsonl")
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
done = _done_paths(out_path)
|
||||
if done:
|
||||
log.info("resuming: %d files already recorded", len(done))
|
||||
|
||||
processed = 0
|
||||
started = time.monotonic()
|
||||
with out_path.open("a", encoding="utf-8") as handle:
|
||||
for path in _iter_images(args.dataset):
|
||||
if str(path) in done:
|
||||
continue
|
||||
row = _measure(path)
|
||||
handle.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
|
||||
handle.flush()
|
||||
processed += 1
|
||||
if processed % args.progress_every == 0:
|
||||
log.info("%d files, %.2f files/s", processed, processed / (time.monotonic() - started))
|
||||
if args.limit and processed >= args.limit:
|
||||
break
|
||||
|
||||
elapsed = time.monotonic() - started
|
||||
rate = processed / max(elapsed, 1e-9)
|
||||
log.info("done: %d files in %.1f s (%.2f files/s) -> %s", processed, elapsed, rate, out_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user