mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +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
@@ -23,6 +23,8 @@ Run `bash maintain.sh` from the repository root. The authoritative type gate is
|
||||
|
||||
Boundary modules for cv2, Torch, and Diffusers may carry narrow per-file relaxations for unknown third-party types. Keep pure-logic files strict, preserve the local piexif stub, and fix real errors before widening a pragma.
|
||||
|
||||
From a worktree, `uv run` imports the package from the MAIN checkout -- that is where the editable install points. A script measuring a worktree's edit must insert that worktree's `src` at `sys.path[0]` and assert `module.__file__` resolves inside it, or it silently compares unmodified code against itself.
|
||||
|
||||
## Model-adjacent tests
|
||||
|
||||
Do not classify an entire module as untestable because its main path downloads a model. Keep pure behavior covered without downloads, including:
|
||||
@@ -78,8 +80,24 @@ rules follow, and both were broken in practice before they were written down:
|
||||
`TextMarkDetection.match_box` and the registry threads the detection into the mask
|
||||
builder; a mask path that re-runs its own sweep is how the two drift apart.
|
||||
|
||||
The C2PA manifest-store JSON is NOT stable across reads: the reader regenerates manifest
|
||||
URNs and instance ids, so two reads of one unchanged file matched in 54 of 120 measured
|
||||
cases. Compare the derived `c2pa_info`, never the raw store.
|
||||
|
||||
A third seam reaches the same verdict: `collect_metadata_record` ->
|
||||
`evidence_from_metadata_record` -> `identify_from_evidence`, the path a caller uses when
|
||||
collection and verdict run on different machines. Its contract is equality with
|
||||
`identify(path, check_visible=False, check_invisible=False)` on the same image, and it
|
||||
can break from EITHER side -- a region the collector stops walking, or a placement the
|
||||
file path learns to read and the record does not. `tests/test_metadata_record.py` pins
|
||||
it over the tracked fixtures; the corpus comparison is what actually finds the gaps
|
||||
(three real ones so far, recorded in `docs/module-internals.md`). Change either side and
|
||||
re-run both.
|
||||
|
||||
Before changing anything in the detection path, record the detectors' exact verdicts
|
||||
over a local sample first and diff them after. A refactor here is only correct if that
|
||||
record is byte-identical, and a green test suite does not establish that on its own.
|
||||
record is byte-identical, and a green test suite does not establish that on its own. A
|
||||
change that is meant to FIX detection is the exception that proves the rule: the diff
|
||||
must then be exactly the files you intended to change, named in advance.
|
||||
|
||||
Environment setup, dependency recovery, CI behavior, and fixture policy: [`../../docs/development.md`](../../docs/development.md).
|
||||
|
||||
@@ -386,6 +386,31 @@ metadata extraction from verdict logic:
|
||||
- `identify` preserves the path-based API and adds the optional registered
|
||||
visible-mark and open invisible-watermark decoders after extraction.
|
||||
|
||||
### Portable metadata record
|
||||
|
||||
[`metadata_record.py`](../src/remove_ai_watermarks/metadata_record.py) produces the
|
||||
record `evidence_from_metadata_record` consumes, so collection and verdict can run
|
||||
on different machines. Its contract is equality with the file path, and the three
|
||||
defects found while establishing that equality are the reason each rule exists:
|
||||
|
||||
- It walks the file's RAW head, never the `scan_head` buffer. That buffer is the
|
||||
head concatenated with late metadata payloads, so a structural walk runs off the
|
||||
end of the real head and parses appended bytes as chunks — which produced 11 MB
|
||||
records and a phantom AIGC signal.
|
||||
- Samsung Galaxy AI splits its evidence: the `PhotoEditor_Re_Edit_Data` marker sits
|
||||
in the post-EOI trailer while the `genAIType` value it is gated on can sit inside
|
||||
the entropy-coded scan (measured on one file: marker at 580 619, value at
|
||||
382 953). A marked file therefore keeps the whole tail window, not just the
|
||||
trailer.
|
||||
- PIL's info keys are emitted in the file path's own candidate order
|
||||
(`Software`, `Source`, `Title`, `Description`, then EXIF). `generator_from_metadata`
|
||||
returns the FIRST candidate carrying a known token, so dict order alone changed
|
||||
the reported platform on nine NovelAI files.
|
||||
|
||||
Pixel forensics are deliberately absent: nothing in the provenance path reads them.
|
||||
Verified over 3,609 research-scan records — dropping every pixel section changed no
|
||||
verdict.
|
||||
|
||||
The DWT-DCT detector and the visible-mark stage share a single decode of the
|
||||
source, held by
|
||||
a per-call `_SharedDecode`. It exposes two accessors because the two arms need
|
||||
|
||||
+39
-2
@@ -182,8 +182,45 @@ evidence = extract_provenance_evidence(Path("input.png"))
|
||||
report = identify_from_evidence(evidence)
|
||||
```
|
||||
|
||||
If metadata was collected by another component, normalize its nested record
|
||||
without reopening the original file:
|
||||
### Collect once, judge elsewhere
|
||||
|
||||
`collect_metadata_record` splits the two halves apart: it is the only step that
|
||||
touches the file, and it returns a JSON-serializable record the verdict can be
|
||||
built from on another machine, in another process, or later.
|
||||
|
||||
```python
|
||||
import json
|
||||
|
||||
from remove_ai_watermarks.identify import evidence_from_metadata_record, identify_from_evidence
|
||||
from remove_ai_watermarks.metadata_record import collect_metadata_record
|
||||
|
||||
record = collect_metadata_record(Path("input.png")) # reads the file
|
||||
blob = json.dumps(record) # ship it anywhere
|
||||
|
||||
evidence = evidence_from_metadata_record(json.loads(blob), path=Path("input.png"))
|
||||
report = identify_from_evidence(evidence) # reads nothing
|
||||
```
|
||||
|
||||
The verdict is the same one `identify(path, check_visible=False,
|
||||
check_invisible=False)` returns for that file. That equality is the record's whole
|
||||
contract, and it is verified two ways: over the tracked provenance fixtures in
|
||||
`tests/test_metadata_record.py`, and over a local corpus, where 3,478 images
|
||||
(every file carrying a rare signal, plus a random slice) produced identical
|
||||
reports through both paths.
|
||||
|
||||
A record carries metadata regions, never pixels: marker segments before the JPEG
|
||||
scan, every PNG chunk except `IDAT`, RIFF chunks except the coded image, the
|
||||
ISOBMFF provenance boxes, the container's trailer, the parsed EXIF tags the
|
||||
verdict reads by name, PIL's info mapping, and the C2PA manifest store. Typical
|
||||
size is 13 kB (p90 93 kB); the tail belongs to images carrying a large embedded
|
||||
manifest, where the store itself dominates.
|
||||
|
||||
The `path` argument is metadata: it labels the report and is never opened by
|
||||
either function, so a record collected elsewhere can be judged against a path that
|
||||
does not exist locally.
|
||||
|
||||
If metadata was collected by another component instead, normalize its nested
|
||||
record the same way:
|
||||
|
||||
```python
|
||||
from remove_ai_watermarks.identify import (
|
||||
|
||||
@@ -68,7 +68,7 @@ payloads. Removal remuxes either container through ffmpeg with stream copy.
|
||||
Metadata stripping for supported audio containers is a separate implemented
|
||||
path.
|
||||
|
||||
**Box detection window — now handled (v0.6.8):** detection no longer relies on a fixed first-MB read. `metadata.scan_head(path, size)` reads the first `size` bytes and, for ISOBMFF, appends the payloads of late provenance boxes found by `isobmff.scan_c2pa_region` (a file-seeking top-level box walker that skips past `mdat` by size without reading it), so a C2PA/AIGC/IPTC manifest placed AFTER a large `mdat` in a streaming/non-faststart MP4 is now caught. Every C2PA/marker byte scan (`has_ai_metadata`, `aigc_label`, `iptc_ai_system`, `synthid_source`, `exif_generator` XMP, `get_ai_metadata` soft-binding, and `identify`) goes through `scan_head`; for PNG it likewise appends the payloads of `tEXt` / `iTXt` / `zTXt` / `eXIf` / `iCCP` chunks that start beyond the window (`_png_late_metadata`, seeking past `IDAT`), which is how a TC260 AIGC label appended after the pixel stream is caught; for every other input, and for any file that fits inside `size`, it is exactly `f.read(size)`.
|
||||
**Box detection window — now handled (v0.6.8):** detection no longer relies on a fixed first-MB read. `metadata.scan_head(path, size)` reads the first `size` bytes and, for ISOBMFF, appends the payloads of late provenance boxes found by `isobmff.scan_c2pa_region` (a file-seeking top-level box walker that skips past `mdat` by size without reading it), so a C2PA/AIGC/IPTC manifest placed AFTER a large `mdat` in a streaming/non-faststart MP4 is now caught. Every C2PA/marker byte scan (`has_ai_metadata`, `aigc_label`, `iptc_ai_system`, `synthid_source`, `exif_generator` XMP, `get_ai_metadata` soft-binding, and `identify`) goes through `scan_head`; for PNG it likewise appends the payloads of `tEXt` / `iTXt` / `zTXt` / `eXIf` / `iCCP` chunks that start beyond the window (`_png_late_metadata`, seeking past `IDAT`), which is how a TC260 AIGC label appended after the pixel stream is caught; for a file at least `size` bytes long it also appends the metadata text the decoder reaches but a raw read cannot (`_decoder_visible_text`) — a compressed PNG `zTXt` packet, or a WebP XMP chunk past the window, both of which hid real AI labels in a corpus; for any file that fits inside `size`, it is exactly `f.read(size)`.
|
||||
|
||||
Native TC260 MP4/MOV tags do not live in those top-level provenance boxes.
|
||||
`tc260_aigc_payloads` separately seeks through `moov.udta.meta.keys/ilst`, so
|
||||
|
||||
@@ -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())
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Aggregate ``detection_timing.py`` records into per-method, per-segment tables.
|
||||
|
||||
WHAT IT PRODUCES
|
||||
Reading it back: the per-file JSONL is one row per image with a millisecond field
|
||||
per method. This collapses it into percentiles per method, then repeats that per
|
||||
segment -- container, megapixels, file size, C2PA presence, verdict confidence --
|
||||
because a single median hides a path whose cost is carried entirely by one
|
||||
container or by the files that actually have a manifest.
|
||||
|
||||
Writes ``<prefix>_summary.csv`` (long form: segment_kind, segment, method, n, p50,
|
||||
p90, p99, mean) and prints a markdown report.
|
||||
|
||||
Percentiles are computed by nearest-rank on the sorted sample, not interpolated:
|
||||
every reported number is a time some real file actually took.
|
||||
|
||||
DATA SAFETY
|
||||
Reads and writes only the given prefix, which belongs outside the repository.
|
||||
|
||||
uv run python scripts/detection_timing_report.py .local-eval/timing/full
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
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().parent))
|
||||
|
||||
from detection_timing import COMPONENTS as _TIMED
|
||||
|
||||
# Taken from the script that WROTE the rows, in its order, so a probe added, removed
|
||||
# or reordered there cannot silently leave a column missing or mislabelled here.
|
||||
COMPONENTS = tuple(name for name, _ in _TIMED)
|
||||
METHODS = (
|
||||
("cold_extract_evidence_ms", "extract_provenance_evidence (cold)"),
|
||||
("extract_evidence_ms", "extract_provenance_evidence (warm)"),
|
||||
*((f"meta_{name}_ms", f" {name}") for name in COMPONENTS),
|
||||
("meta_components_sum_ms", " (sum of components)"),
|
||||
("verdict_from_evidence_ms", "identify_from_evidence (verdict)"),
|
||||
("identify_metadata_only_ms", "identify(metadata only), end to end"),
|
||||
)
|
||||
# Median-only columns in the per-segment tables; the full percentile set goes to the CSV.
|
||||
HEADLINE_LABELS = (
|
||||
("extract_evidence_ms", "extract p50"),
|
||||
("verdict_from_evidence_ms", "verdict p50"),
|
||||
("identify_metadata_only_ms", "end-to-end p50"),
|
||||
)
|
||||
HEADLINE = tuple(key for key, _ in HEADLINE_LABELS)
|
||||
|
||||
|
||||
def _bucket(value: float | None, edges: tuple[float, ...], labels: tuple[str, ...]) -> str:
|
||||
if value is None:
|
||||
return "unknown"
|
||||
for edge, label in zip(edges, labels, strict=False):
|
||||
if value < edge:
|
||||
return label
|
||||
return labels[-1]
|
||||
|
||||
|
||||
SEGMENTS: tuple[tuple[str, Callable[[dict[str, Any]], str]], ...] = (
|
||||
("container", lambda r: str(r.get("format") or "unknown")),
|
||||
(
|
||||
"megapixels",
|
||||
lambda r: _bucket(r.get("megapixels"), (1, 4, 12), ("<1 MP", "1-4 MP", "4-12 MP", ">12 MP")),
|
||||
),
|
||||
(
|
||||
"file size",
|
||||
lambda r: _bucket(
|
||||
(r["bytes"] / 1e6) if r.get("bytes") is not None else None,
|
||||
(1, 5, 20),
|
||||
("<1 MB", "1-5 MB", "5-20 MB", ">20 MB"),
|
||||
),
|
||||
),
|
||||
("c2pa", lambda r: "with C2PA" if r.get("has_c2pa") else "no C2PA"),
|
||||
("verdict", lambda r: f"confidence={r.get('confidence') or 'n/a'}"),
|
||||
)
|
||||
|
||||
|
||||
def _numeric(records: list[dict[str, Any]], key: str) -> list[float]:
|
||||
"""The numeric values of ``key``, skipping rows where it is absent or non-numeric.
|
||||
|
||||
One definition of "counts as a measurement", so the printed table and the CSV can
|
||||
never disagree about which rows a method was measured on.
|
||||
"""
|
||||
return [float(r[key]) for r in records if isinstance(r.get(key), (int, float))]
|
||||
|
||||
|
||||
def _percentile(sorted_values: list[float], q: float) -> float:
|
||||
"""Nearest-rank percentile: the reported value is one a real file produced."""
|
||||
index = max(0, math.ceil(q * len(sorted_values)) - 1)
|
||||
return sorted_values[index]
|
||||
|
||||
|
||||
def _stats(values: list[float]) -> dict[str, float | int]:
|
||||
ordered = sorted(values)
|
||||
return {
|
||||
"n": len(ordered),
|
||||
"p50": _percentile(ordered, 0.50),
|
||||
"p90": _percentile(ordered, 0.90),
|
||||
"p99": _percentile(ordered, 0.99),
|
||||
"mean": statistics.fmean(ordered),
|
||||
}
|
||||
|
||||
|
||||
def _rows(path: Path) -> Iterator[dict[str, Any]]:
|
||||
with path.open(encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
try:
|
||||
yield json.loads(line)
|
||||
except ValueError: # a truncated tail while the run is still writing
|
||||
continue
|
||||
|
||||
|
||||
def _table(rows: list[dict[str, Any]], methods: tuple[tuple[str, str], ...]) -> list[str]:
|
||||
out = ["| method | n | p50 | p90 | p99 | mean |", "|---|---:|---:|---:|---:|---:|"]
|
||||
for key, label in methods:
|
||||
values = _numeric(rows, key)
|
||||
if not values:
|
||||
continue
|
||||
s = _stats(values)
|
||||
out.append(f"| {label} | {s['n']} | {s['p50']:.3f} | {s['p90']:.3f} | {s['p99']:.3f} | {s['mean']:.3f} |")
|
||||
return out
|
||||
|
||||
|
||||
def _correlation(rows: list[dict[str, Any]], x_key: str, y_key: str) -> float | None:
|
||||
both = [r for r in rows if isinstance(r.get(x_key), (int, float)) and isinstance(r.get(y_key), (int, float))]
|
||||
if len(both) < 3:
|
||||
return None
|
||||
xs = _numeric(both, x_key)
|
||||
ys = _numeric(both, y_key)
|
||||
try:
|
||||
return statistics.correlation(xs, ys)
|
||||
except statistics.StatisticsError: # a constant column has no correlation
|
||||
return None
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("prefix", type=Path, help="prefix given to detection_timing.py")
|
||||
args = parser.parse_args()
|
||||
|
||||
jsonl = args.prefix.with_suffix(".jsonl")
|
||||
all_rows = list(_rows(jsonl))
|
||||
failed = [r for r in all_rows if "error" in r]
|
||||
rows = [r for r in all_rows if "error" not in r]
|
||||
if not rows:
|
||||
print(f"no usable records in {jsonl}")
|
||||
return 1
|
||||
|
||||
disagreed = [r for r in rows if r.get("verdict_agrees") is False]
|
||||
|
||||
lines = [
|
||||
f"# Detection timing over {len(rows)} images",
|
||||
"",
|
||||
f"Source: `{jsonl}`. Failed rows: {len(failed)}. Verdict-path mismatches: {len(disagreed)}.",
|
||||
"",
|
||||
"All times in milliseconds. Percentiles are nearest-rank.",
|
||||
"",
|
||||
"## Whole corpus",
|
||||
"",
|
||||
*_table(rows, METHODS),
|
||||
"",
|
||||
]
|
||||
|
||||
csv_rows: list[dict[str, Any]] = []
|
||||
for key, label in METHODS:
|
||||
values = _numeric(rows, key)
|
||||
if values:
|
||||
csv_rows.append({"segment_kind": "all", "segment": "all", "method": label.strip(), **_stats(values)})
|
||||
|
||||
for kind, classify in SEGMENTS:
|
||||
groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in rows:
|
||||
groups[classify(row)].append(row)
|
||||
lines += [f"## By {kind}", ""]
|
||||
header = ["| segment | n | " + " | ".join(label for _, label in HEADLINE_LABELS) + " |"]
|
||||
header.append("|---|---:|" + "---:|" * len(HEADLINE))
|
||||
lines += header
|
||||
for segment, group in sorted(groups.items(), key=lambda item: -len(item[1])):
|
||||
cells = []
|
||||
for key in HEADLINE:
|
||||
values = _numeric(group, key)
|
||||
cells.append(f"{_percentile(sorted(values), 0.5):.2f}" if values else "-")
|
||||
lines.append(f"| {segment} | {len(group)} | " + " | ".join(cells) + " |")
|
||||
for key, label in METHODS:
|
||||
values = _numeric(group, key)
|
||||
if values:
|
||||
csv_rows.append(
|
||||
{"segment_kind": kind, "segment": segment, "method": label.strip(), **_stats(values)}
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
corr = _correlation(rows, "scan_bytes", "verdict_from_evidence_ms")
|
||||
corr_size = _correlation(rows, "bytes", "extract_evidence_ms")
|
||||
lines += [
|
||||
"## Scaling",
|
||||
"",
|
||||
f"- verdict time vs scan-buffer size: r = {corr:.3f}" if corr is not None else "- verdict correlation: n/a",
|
||||
(
|
||||
f"- extraction time vs file size: r = {corr_size:.3f}"
|
||||
if corr_size is not None
|
||||
else "- extraction correlation: n/a"
|
||||
),
|
||||
"",
|
||||
]
|
||||
|
||||
summary_csv = args.prefix.with_name(args.prefix.name + "_summary.csv")
|
||||
columns = ["segment_kind", "segment", "method", "n", "p50", "p90", "p99", "mean"]
|
||||
with summary_csv.open("w", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
writer.writerows(csv_rows)
|
||||
|
||||
report = "\n".join(lines)
|
||||
args.prefix.with_name(args.prefix.name + "_report.md").write_text(report + "\n", encoding="utf-8")
|
||||
print(report)
|
||||
print(f"\nwrote {summary_csv}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -306,11 +306,16 @@ def scan_head(image_path: Path, size: int = 1024 * 1024) -> bytes:
|
||||
past large boxes like ``mdat``) and PNG ``tEXt`` / ``iTXt`` / ``eXIf`` chunks
|
||||
(seeking past ``IDAT``).
|
||||
|
||||
A file at least ``size`` bytes long additionally gets the metadata text its
|
||||
decoder can reach but a raw read cannot (:func:`_decoder_visible_text`): a
|
||||
compressed PNG ``zTXt`` packet, or a chunk past the window in a container with no
|
||||
late-chunk reader here. A file that fits inside ``size`` is exactly
|
||||
``f.read(size)``, since the raw read already holds every byte.
|
||||
|
||||
This is the shared input for every C2PA / AIGC / IPTC byte scan. The
|
||||
extensions catch a manifest or XMP packet placed AFTER the media data -- a
|
||||
non-faststart MP4 manifest, or a PNG XMP packet appended after the pixels --
|
||||
which a fixed first-MB read would miss. For other inputs, and for files that
|
||||
fit within ``size``, it is exactly ``f.read(size)`` -- behavior-neutral.
|
||||
which a fixed first-MB read would miss.
|
||||
|
||||
The result is memoized per (path, size, mtime): one ``identify``/``get_ai_metadata``
|
||||
call fans out to ~8 byte-scan detectors that each call this on the same file, so
|
||||
@@ -347,9 +352,62 @@ def _scan_head_impl(image_path: Path, size: int) -> bytes:
|
||||
# len(head) == size means the file is at least `size` bytes, so metadata
|
||||
# chunks may lie beyond the window; otherwise the whole PNG is in `head`.
|
||||
head += _png_late_metadata(image_path, size)
|
||||
if len(head) >= size:
|
||||
head += _decoder_visible_text(image_path, head)
|
||||
return head
|
||||
|
||||
|
||||
# Text values the image decoder can reach that a raw byte read cannot. Bounded: a
|
||||
# packet larger than this is not a provenance label.
|
||||
_DECODED_TEXT_LIMIT = 512 * 1024
|
||||
# Decoder values that are binary payloads with their own readers, not metadata text.
|
||||
# An ICC profile is colour data and can run to hundreds of kilobytes; appending it
|
||||
# would bloat the buffer every later detector re-scans, for no signal.
|
||||
_DECODER_BINARY_KEYS = frozenset({"icc_profile"})
|
||||
|
||||
|
||||
def _decoder_visible_text(image_path: Path, head: bytes) -> bytes:
|
||||
"""Metadata text PIL can decode but the raw window does not contain.
|
||||
|
||||
Two placements defeat a fixed byte read, and both were found in a real corpus
|
||||
rather than imagined:
|
||||
|
||||
* COMPRESSED -- a PNG ``zTXt`` chunk is zlib-deflated, so an XMP packet carrying
|
||||
a TC260 AIGC label is unreadable as bytes while PIL inflates it on open. Five
|
||||
corpus files carried a China AIGC label that ``identify`` reported as no signal
|
||||
at all.
|
||||
* BEYOND THE WINDOW in a container with no late-chunk reader -- a WebP XMP chunk
|
||||
at offset 1 093 039 sits 44 kB past the 1 MiB window, and unlike PNG and
|
||||
ISOBMFF, RIFF has no seek-past-the-pixels extension here. Three corpus files
|
||||
hid an IPTC "Made with AI" tag and a C2PA ``trainedAlgorithmicMedia`` that way.
|
||||
|
||||
Only text ALREADY MISSING from ``head`` is appended, so the common case adds
|
||||
nothing and no detector sees a value twice. Skipped entirely when the file fits
|
||||
inside the window, since then the raw read already holds every byte.
|
||||
"""
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(image_path) as img:
|
||||
values = [value for key, value in img.info.items() if key not in _DECODER_BINARY_KEYS]
|
||||
except Exception as exc: # a container PIL cannot open: the raw scan stands alone
|
||||
logger.debug("decoder-visible text unavailable for %s: %s", image_path, exc)
|
||||
return b""
|
||||
|
||||
out = bytearray()
|
||||
for value in values:
|
||||
if isinstance(value, str):
|
||||
encoded = value.encode("utf-8", "replace")
|
||||
elif isinstance(value, bytes):
|
||||
encoded = value
|
||||
else:
|
||||
continue
|
||||
if len(encoded) > _DECODED_TEXT_LIMIT or not encoded or encoded in head:
|
||||
continue
|
||||
out += b"\x00" + encoded
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def has_ai_metadata(image_path: Path) -> bool:
|
||||
"""Check if an image contains AI-generation metadata.
|
||||
|
||||
@@ -1485,3 +1543,15 @@ def xai_signature(image_path: Path) -> bool:
|
||||
if key is None:
|
||||
return _xai_signature_impl(image_path)
|
||||
return _xai_signature_cached(*key)
|
||||
|
||||
|
||||
# ── Shared with the portable metadata record ────────────────────────
|
||||
# `metadata_record` must read exactly the windows and markers the file path reads: a
|
||||
# record built from a different window is a record whose verdict can disagree with
|
||||
# `identify` on the same image. Aliased rather than renamed because the private names
|
||||
# are load-bearing in this module's own tests and in a corpus script.
|
||||
QUICK_SCAN_BYTES = _QUICK_SCAN_BYTES
|
||||
SAMSUNG_EDITOR_MARKER = _SAMSUNG_EDITOR_MARKER
|
||||
read_file_tail = _read_file_tail
|
||||
png_late_metadata = _png_late_metadata
|
||||
exif_text = _exif_text
|
||||
|
||||
@@ -0,0 +1,338 @@
|
||||
"""Collect one image's provenance metadata into a portable, JSON-safe record.
|
||||
|
||||
WHY THIS EXISTS
|
||||
|
||||
``extract_provenance_evidence`` reads a file and hands back evidence in memory, so
|
||||
collection and verdict must happen in the same process, on the machine holding the
|
||||
image. This module splits them: collect here, judge anywhere, from a record that
|
||||
survives JSON.
|
||||
|
||||
record = collect_metadata_record(path) # touches the file
|
||||
evidence = evidence_from_metadata_record(record, path=path)
|
||||
report = identify_from_evidence(evidence) # touches nothing
|
||||
|
||||
WHAT GOES IN, AND WHY NOT SIMPLY THE FILE HEAD
|
||||
|
||||
The verdict reads a scan buffer that ``scan_head`` fills with the first mebibyte of
|
||||
the file. Shipping that verbatim would make a record larger than a phone photo's
|
||||
worth of metadata by two orders of magnitude, because for a PNG almost all of that
|
||||
mebibyte is compressed pixel data in ``IDAT`` -- bytes no provenance token can ever
|
||||
live in. A record carries the metadata REGIONS instead, walked per container: the
|
||||
JPEG marker segments before the coded scan, every PNG chunk but ``IDAT``, the RIFF
|
||||
chunks that are not coded image, the ISOBMFF provenance boxes, and in every case the
|
||||
container's trailer.
|
||||
|
||||
COMPLETENESS IS A MEASURED PROPERTY, NOT A CLAIM
|
||||
|
||||
A region walker is only correct if nothing the verdict reads falls outside the
|
||||
regions it keeps, and no test over fixtures can establish that: the failure mode is
|
||||
a container placement nobody thought of. The contract is therefore ALSO verified
|
||||
against the file path over a real corpus -- same image, both paths, identical
|
||||
``ProvenanceReport``.
|
||||
|
||||
The placements that defeated an earlier draft of this collector, and the reason each
|
||||
rule below exists, are recorded in ``docs/module-internals.md`` under "Portable
|
||||
metadata record".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import struct
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from remove_ai_watermarks._internal.constants import PNG_SIGNATURE
|
||||
from remove_ai_watermarks.metadata import (
|
||||
QUICK_SCAN_BYTES,
|
||||
SAMSUNG_EDITOR_MARKER,
|
||||
exif_text,
|
||||
read_file_tail,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# The structural walk covers the same window the file path reads raw, so the two
|
||||
# cannot disagree about a chunk type inside it. A smaller window would be cheaper but
|
||||
# opens a blind spot: past the window only ``png_late_metadata``'s ALLOWLIST is
|
||||
# collected, while the file path still sees every chunk type up to its own window --
|
||||
# and a C2PA ``caBX`` chunk is in neither that allowlist nor ``IDAT``. Walking here
|
||||
# costs little because the payload of the pixel stream is skipped, not copied.
|
||||
HEAD_WINDOW = 1024 * 1024
|
||||
# The window searched for the container's end marker. Matches the quick-scan window
|
||||
# the file path uses when it goes looking for a Samsung trailer, so a trailer visible
|
||||
# to one path is visible to the other.
|
||||
TAIL_WINDOW = QUICK_SCAN_BYTES
|
||||
# Kept from the tail when no end marker is found, so an unrecognized container still
|
||||
# contributes its last bytes without carrying half a photo.
|
||||
UNKNOWN_TRAILER_WINDOW = 64 * 1024
|
||||
|
||||
# PNG text keys the file path reads for a generator tag, in ITS order. NovelAI stamps
|
||||
# Software/Source/Title rather than EXIF, and the first match wins, so order matters.
|
||||
_GENERATOR_TEXT_KEYS = ("Software", "Source", "Title", "Description")
|
||||
# RIFF chunks holding coded pixels rather than metadata.
|
||||
_RIFF_IMAGE_CHUNKS = frozenset({b"VP8 ", b"VP8L", b"ALPH"})
|
||||
|
||||
|
||||
def _jpeg_regions(data: bytes) -> bytes:
|
||||
"""Every marker segment up to the entropy-coded scan, plus the trailer after EOI.
|
||||
|
||||
The scan itself is skipped by walking to SOS and then jumping to the trailing
|
||||
EOI, so a 20 MB photo contributes only its markers.
|
||||
"""
|
||||
out = bytearray()
|
||||
index, size = 2, len(data)
|
||||
while index + 1 < size:
|
||||
if data[index] != 0xFF:
|
||||
break # malformed boundary: keep what was collected, the tail still follows
|
||||
marker = data[index + 1]
|
||||
if marker in (0xDA, 0xD9): # SOS / EOI: the coded scan follows
|
||||
break
|
||||
if 0xD0 <= marker <= 0xD7 or marker == 0x01: # standalone, no length
|
||||
index += 2
|
||||
continue
|
||||
if index + 4 > size:
|
||||
break
|
||||
segment_length = int.from_bytes(data[index + 2 : index + 4], "big")
|
||||
end = index + 2 + segment_length
|
||||
if segment_length < 2 or end > size:
|
||||
break
|
||||
out += data[index:end]
|
||||
index = end
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _png_regions(data: bytes) -> bytes:
|
||||
"""Every chunk except the ``IDAT`` payloads, plus whatever follows IEND."""
|
||||
out = bytearray()
|
||||
size = len(data)
|
||||
position = len(PNG_SIGNATURE)
|
||||
while position + 8 <= size:
|
||||
(length,) = struct.unpack(">I", data[position : position + 4])
|
||||
chunk_type = data[position + 4 : position + 8]
|
||||
start = position + 8
|
||||
# Clamp the length to the bytes that remain: a malformed 32-bit length must
|
||||
# not push the walk past EOF and abandon a genuine label chunk after it.
|
||||
safe_length = max(0, min(length, size - start))
|
||||
if chunk_type != b"IDAT":
|
||||
out += chunk_type + data[start : start + safe_length]
|
||||
position = start + safe_length + 4 # payload + CRC
|
||||
if chunk_type == b"IEND":
|
||||
out += data[position:] # a trailer past IEND is metadata too
|
||||
break
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _riff_regions(data: bytes) -> bytes:
|
||||
"""Every RIFF chunk except the coded image payloads."""
|
||||
out = bytearray(data[:12]) # 'RIFF' + size + 'WEBP'
|
||||
size = len(data)
|
||||
position = 12
|
||||
while position + 8 <= size:
|
||||
chunk_type = data[position : position + 4]
|
||||
(length,) = struct.unpack("<I", data[position + 4 : position + 8])
|
||||
start = position + 8
|
||||
safe_length = max(0, min(length, size - start))
|
||||
if chunk_type not in _RIFF_IMAGE_CHUNKS:
|
||||
out += chunk_type + data[start : start + safe_length]
|
||||
position = start + safe_length + (safe_length & 1) # chunks are word-aligned
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _isobmff_regions(image_path: Path, head: bytes) -> bytes:
|
||||
"""Header window plus the provenance regions the bounded box walkers find.
|
||||
|
||||
ISOBMFF hides a manifest in a ``uuid``/``jumb`` box that can sit after a
|
||||
multi-megabyte ``mdat``, and a TC260 label in ``moov.udta``. Both walkers seek
|
||||
rather than read the media, so neither pulls the payload in.
|
||||
"""
|
||||
from remove_ai_watermarks._internal.isobmff import scan_c2pa_region, tc260_aigc_payloads
|
||||
|
||||
out = bytearray(head[:HEAD_WINDOW])
|
||||
try:
|
||||
out += scan_c2pa_region(image_path)
|
||||
except Exception as exc:
|
||||
logger.debug("ISOBMFF C2PA region scan failed on %s: %s", image_path, exc)
|
||||
try:
|
||||
for payload in tc260_aigc_payloads(image_path):
|
||||
out += payload
|
||||
except Exception as exc:
|
||||
logger.debug("ISOBMFF TC260 scan failed on %s: %s", image_path, exc)
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _container_regions(image_path: Path, head: bytes) -> tuple[str, bytes]:
|
||||
"""(container label, metadata bytes) for the container ``head`` starts with.
|
||||
|
||||
``head`` must be the file's raw first bytes. Handing this the ``scan_head``
|
||||
buffer instead is a trap that was walked into once: that buffer is the head
|
||||
CONCATENATED with late metadata payloads, so a structural walk runs off the end
|
||||
of the real head and parses the appended bytes as chunks -- which produced 11 MB
|
||||
records and a phantom AIGC signal before the two were separated.
|
||||
"""
|
||||
from remove_ai_watermarks._internal.isobmff import is_isobmff
|
||||
from remove_ai_watermarks.metadata import png_late_metadata
|
||||
|
||||
if head.startswith(b"\xff\xd8"):
|
||||
return "jpeg", _jpeg_regions(head)
|
||||
if head.startswith(PNG_SIGNATURE):
|
||||
# Chunks placed after the pixel stream (an XMP packet at 2.7 MB, say) are
|
||||
# past the window; the same seek-past-IDAT reader the file path uses gets them.
|
||||
return "png", _png_regions(head) + png_late_metadata(image_path, HEAD_WINDOW)
|
||||
if head.startswith(b"RIFF") and head[8:12] == b"WEBP":
|
||||
return "webp", _riff_regions(head)
|
||||
if is_isobmff(head):
|
||||
return "isobmff", _isobmff_regions(image_path, head)
|
||||
return "unknown", head
|
||||
|
||||
|
||||
def _raw_head(image_path: Path) -> bytes:
|
||||
"""The file's first bytes, unmodified -- the input every structural walk needs."""
|
||||
try:
|
||||
with open(image_path, "rb") as handle:
|
||||
return handle.read(HEAD_WINDOW)
|
||||
except OSError as exc:
|
||||
logger.debug("head read failed for %s: %s", image_path, exc)
|
||||
return b""
|
||||
|
||||
|
||||
def _trailer(image_path: Path, container: str) -> bytes:
|
||||
"""The bytes that follow the container's end marker, and nothing else.
|
||||
|
||||
A fixed-size tail read would be almost entirely pixels: the trailer of a 20 MB
|
||||
photo is a few kilobytes at most. So the end marker is located in the tail window
|
||||
and only what follows it is kept. When no marker is found (an unknown container,
|
||||
or one whose end lies before the window) the window is kept as-is, bounded --
|
||||
that is what a byte scan of the same file would have seen anyway.
|
||||
"""
|
||||
tail = read_file_tail(image_path, TAIL_WINDOW)
|
||||
if SAMSUNG_EDITOR_MARKER in tail:
|
||||
# Galaxy AI splits its evidence: the marker sits in the post-EOI trailer, but
|
||||
# the `genAIType` value it is gated on can sit INSIDE the entropy-coded scan
|
||||
# (measured: marker at 580 619, value at 382 953 in the same file). Keeping
|
||||
# only the trailer therefore carries the marker without the value and the
|
||||
# verdict silently drops the Samsung signal, so a marked file keeps the whole
|
||||
# window. Only Samsung-marked files pay for it.
|
||||
return tail
|
||||
marker = {"jpeg": b"\xff\xd9", "png": b"IEND\xae\x42\x60\x82"}.get(container)
|
||||
if marker is None:
|
||||
return tail[-UNKNOWN_TRAILER_WINDOW:]
|
||||
index = tail.rfind(marker)
|
||||
return tail[index + len(marker) :] if index >= 0 else tail[-UNKNOWN_TRAILER_WINDOW:]
|
||||
|
||||
|
||||
def _decoder_info(image_path: Path) -> dict[str, Any]:
|
||||
"""PIL's ``info`` mapping, read once.
|
||||
|
||||
One open for both consumers below. They want different parts of the same mapping
|
||||
(the text keys, and the raw EXIF blob), and opening twice repeats the container
|
||||
header parse and, for a PNG carrying ``zTXt``, the zlib inflate with it.
|
||||
"""
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
with Image.open(image_path) as img:
|
||||
# PIL types this mapping with a non-string key union (a DPI tuple key
|
||||
# exists), so the keys are normalized here rather than assumed.
|
||||
return {str(key): value for key, value in img.info.items()}
|
||||
except Exception as exc: # a container PIL cannot open
|
||||
logger.debug("PIL info unavailable for %s: %s", image_path, exc)
|
||||
return {}
|
||||
|
||||
|
||||
def _exif_pairs(info: dict[str, Any]) -> dict[str, str]:
|
||||
"""The 0th-IFD tags the verdict reads, under their tag NAMES.
|
||||
|
||||
Not a convenience: two probes key on names rather than on the raw bytes already
|
||||
in the regions. ``xai_signature_pair`` wants an (ImageDescription, Artist) pair,
|
||||
and ``_external_exif_generator`` looks for Software / Make / Artist /
|
||||
ImageDescription. Ship the bytes alone and both silently return nothing, which
|
||||
is exactly how a first draft of this collector lost every Grok and NovelAI
|
||||
verdict in a corpus run.
|
||||
"""
|
||||
exif_bytes = info.get("exif")
|
||||
if not exif_bytes:
|
||||
return {}
|
||||
try:
|
||||
import piexif
|
||||
|
||||
tags = piexif.load(exif_bytes).get("0th", {})
|
||||
except Exception as exc: # malformed EXIF
|
||||
logger.debug("EXIF parse failed: %s", exc)
|
||||
return {}
|
||||
|
||||
return {
|
||||
name: text
|
||||
for name, tag in (
|
||||
("Software", piexif.ImageIFD.Software),
|
||||
("Make", piexif.ImageIFD.Make),
|
||||
("Artist", piexif.ImageIFD.Artist),
|
||||
("ImageDescription", piexif.ImageIFD.ImageDescription),
|
||||
)
|
||||
if (text := exif_text(tags, tag))
|
||||
}
|
||||
|
||||
|
||||
def _pil_info(info: dict[str, Any]) -> dict[str, str]:
|
||||
"""PIL's ``info`` mapping as strings, the source of PNG text keys and ``hf-job-id``."""
|
||||
|
||||
def text_of(value: Any) -> str:
|
||||
return value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value)
|
||||
|
||||
# Emitted in the file path's own candidate order. ``generator_from_metadata``
|
||||
# returns the FIRST candidate carrying a known token, and a record that listed
|
||||
# PIL's keys in their natural dict order picked a different string for the same
|
||||
# image -- nine NovelAI files came back as "NovelAI generated image" where the
|
||||
# file path says "NovelAI". Same verdict, different platform text, and the two
|
||||
# paths are supposed to be indistinguishable.
|
||||
out: dict[str, str] = {}
|
||||
for key in _GENERATOR_TEXT_KEYS:
|
||||
value = info.get(key)
|
||||
if value is not None and not isinstance(value, (dict, list, tuple)):
|
||||
out[f"info:{key}"] = text_of(value)
|
||||
for key, value in info.items():
|
||||
if key in _GENERATOR_TEXT_KEYS or isinstance(value, (dict, list, tuple)):
|
||||
continue
|
||||
out[f"info:{key}"] = text_of(value)
|
||||
return out
|
||||
|
||||
|
||||
def collect_metadata_record(image_path: Path) -> dict[str, Any]:
|
||||
"""Collect everything the provenance verdict reads, as a JSON-safe record.
|
||||
|
||||
The record is the transport format for
|
||||
:func:`identify.evidence_from_metadata_record`: it carries the metadata regions
|
||||
(base64), the C2PA manifest store, and PIL's info mapping, and it never carries
|
||||
pixel data.
|
||||
|
||||
Args:
|
||||
image_path: Path to the image.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable dict. ``metadata_base64`` holds the concatenated
|
||||
container regions, ``tail_base64`` the file trailer.
|
||||
"""
|
||||
from remove_ai_watermarks._internal.c2pa import read_manifest_store_json
|
||||
|
||||
container, regions = _container_regions(image_path, _raw_head(image_path))
|
||||
|
||||
info = _decoder_info(image_path)
|
||||
record: dict[str, Any] = {
|
||||
"container": container,
|
||||
"name": image_path.name,
|
||||
"metadata_base64": base64.b64encode(regions).decode("ascii"),
|
||||
# Always collected: Samsung's Galaxy AI marker is a post-EOI trailer, and a
|
||||
# record without it loses that verdict outright.
|
||||
"tail_base64": base64.b64encode(_trailer(image_path, container)).decode("ascii"),
|
||||
# PIL info BEFORE exif: the file path prefers a PNG text tag over an EXIF
|
||||
# one, and the normalizer walks the record in insertion order.
|
||||
"pil": _pil_info(info),
|
||||
"exif": _exif_pairs(info),
|
||||
}
|
||||
store = read_manifest_store_json(image_path)
|
||||
if store is not None:
|
||||
record["c2pa_store"] = store
|
||||
return record
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Tests for the portable metadata record.
|
||||
|
||||
The contract is one property: a verdict built from the record equals the verdict
|
||||
built from the file. Everything here exists to pin that, or to pin a placement the
|
||||
record could plausibly drop.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import struct
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from remove_ai_watermarks.identify import (
|
||||
ProvenanceReport,
|
||||
evidence_from_metadata_record,
|
||||
identify,
|
||||
identify_from_evidence,
|
||||
)
|
||||
from remove_ai_watermarks.metadata_record import HEAD_WINDOW, collect_metadata_record
|
||||
|
||||
FIXTURES = Path(__file__).resolve().parent.parent / "data" / "fixtures" / "provenance"
|
||||
COMPARED = ("is_ai_generated", "platform", "confidence", "ai_source_kind", "ai_from_metadata")
|
||||
|
||||
|
||||
def _verdict_via_record(path: Path) -> ProvenanceReport:
|
||||
"""The contractor's path: collect, serialize, judge -- without the file."""
|
||||
record = json.loads(json.dumps(collect_metadata_record(path)))
|
||||
return identify_from_evidence(evidence_from_metadata_record(record, path=path))
|
||||
|
||||
|
||||
def _assert_same_verdict(path: Path) -> None:
|
||||
via_record = _verdict_via_record(path)
|
||||
via_file = identify(path, check_visible=False, check_invisible=False)
|
||||
|
||||
for field in COMPARED:
|
||||
assert getattr(via_record, field) == getattr(via_file, field), field
|
||||
assert sorted(s.name for s in via_record.signals) == sorted(s.name for s in via_file.signals)
|
||||
assert sorted(via_record.watermarks) == sorted(via_file.watermarks)
|
||||
|
||||
|
||||
def _noise_png(path: Path, size: tuple[int, int] = (700, 700)) -> Path:
|
||||
"""A PNG whose IDAT is incompressible, so the file exceeds the head window."""
|
||||
rng = np.random.default_rng(0)
|
||||
Image.fromarray(rng.integers(0, 255, (size[1], size[0], 3), dtype=np.uint8)).save(path)
|
||||
return path
|
||||
|
||||
|
||||
def _insert_png_chunk(path: Path, chunk_type: bytes, payload: bytes) -> None:
|
||||
"""Splice a chunk in just before IEND, i.e. AFTER the whole pixel stream."""
|
||||
data = path.read_bytes()
|
||||
end = data.rindex(b"IEND") - 4
|
||||
chunk = struct.pack(">I", len(payload)) + chunk_type + payload
|
||||
chunk += struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
|
||||
path.write_bytes(data[:end] + chunk + data[end:])
|
||||
|
||||
|
||||
class TestRecordReproducesTheFileVerdict:
|
||||
"""The whole point of the record. Every tracked provenance fixture is compared,
|
||||
which covers C2PA, the China TC260 label, the xAI EXIF pair and the IPTC tag."""
|
||||
|
||||
@pytest.mark.skipif(not FIXTURES.is_dir(), reason="provenance fixtures not present")
|
||||
@pytest.mark.parametrize("name", sorted(p.name for p in FIXTURES.iterdir()) if FIXTURES.is_dir() else [])
|
||||
def test_fixture(self, name: str):
|
||||
_assert_same_verdict(FIXTURES / name)
|
||||
|
||||
@pytest.mark.skipif(not FIXTURES.is_dir(), reason="provenance fixtures not present")
|
||||
def test_the_fixtures_actually_exercise_several_signals(self):
|
||||
"""Guards the test above: if the fixture set ever narrows to one signal
|
||||
family, equality across it stops meaning much."""
|
||||
found = {
|
||||
signal.name
|
||||
for path in FIXTURES.iterdir()
|
||||
for signal in identify(path, check_visible=False, check_invisible=False).signals
|
||||
}
|
||||
assert len(found) >= 4, found
|
||||
|
||||
|
||||
class TestPlacementsTheRecordCouldDrop:
|
||||
def test_a_trailer_after_eoi_survives(self, tmp_path: Path):
|
||||
"""Samsung Galaxy AI appends its marker past the JPEG EOI, and the value it is
|
||||
gated on can sit further back still. A record that stopped at the last
|
||||
structural marker would report no signal at all."""
|
||||
path = tmp_path / "edited.jpg"
|
||||
Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "JPEG")
|
||||
with path.open("ab") as handle:
|
||||
handle.write(b'PhotoEditor_Re_Edit_Data{"genAIType":17}')
|
||||
|
||||
assert identify(path, check_visible=False, check_invisible=False).confidence == "medium"
|
||||
_assert_same_verdict(path)
|
||||
|
||||
def test_a_metadata_chunk_past_the_head_window_survives(self, tmp_path: Path):
|
||||
"""A PNG encoder may put the label chunk after the pixels. The file is larger
|
||||
than the record's head window, so only the seek-past-IDAT reader finds it."""
|
||||
path = _noise_png(tmp_path / "late.png")
|
||||
assert path.stat().st_size > HEAD_WINDOW
|
||||
# An XMP packet in the namespaced TC260 form, which is how the label travels
|
||||
# when it is not a bare ``AIGC`` keyword chunk.
|
||||
label = json.dumps({"Label": "1", "ContentProducer": "001191110102MACQD9K64010000"})
|
||||
xmp = (
|
||||
b'<x:xmpmeta xmlns:x="adobe:ns:meta/"><rdf:RDF><rdf:Description '
|
||||
b'xmlns:TC260="http://www.tc260.org.cn/ns/AIGC/1.0/"><TC260:AIGC>'
|
||||
+ label.encode()
|
||||
+ b"</TC260:AIGC></rdf:Description></rdf:RDF></x:xmpmeta>"
|
||||
)
|
||||
_insert_png_chunk(path, b"iTXt", b"XML:com.adobe.xmp\x00\x00\x00\x00\x00" + xmp)
|
||||
|
||||
assert "aigc" in {s.name for s in identify(path, check_visible=False, check_invisible=False).signals}
|
||||
_assert_same_verdict(path)
|
||||
|
||||
def test_an_exif_pair_survives(self, tmp_path: Path):
|
||||
"""xAI is recognized from an (ImageDescription, Artist) PAIR, and both are read
|
||||
by tag NAME. A record carrying only raw bytes loses it."""
|
||||
import piexif
|
||||
|
||||
path = tmp_path / "grok.jpg"
|
||||
Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "JPEG")
|
||||
exif = {
|
||||
"0th": {
|
||||
piexif.ImageIFD.ImageDescription: b"Signature: " + b"A" * 80,
|
||||
piexif.ImageIFD.Artist: b"3f2504e0-4f89-11d3-9a0c-0305e82c3301",
|
||||
}
|
||||
}
|
||||
piexif.insert(piexif.dump(exif), str(path))
|
||||
|
||||
assert "xai_signature" in {s.name for s in identify(path, check_visible=False, check_invisible=False).signals}
|
||||
_assert_same_verdict(path)
|
||||
|
||||
|
||||
class TestRecordShape:
|
||||
def test_the_record_survives_json(self, tmp_path: Path):
|
||||
record = collect_metadata_record(_noise_png(tmp_path / "plain.png"))
|
||||
|
||||
assert json.loads(json.dumps(record))["container"] == "png"
|
||||
|
||||
def test_pixels_are_not_carried(self, tmp_path: Path):
|
||||
"""The reason the record walks regions instead of shipping the head: a record
|
||||
that carried the pixel stream would be the size of the image."""
|
||||
path = _noise_png(tmp_path / "big.png", size=(1200, 1200))
|
||||
record = collect_metadata_record(path)
|
||||
|
||||
idat = path.read_bytes()
|
||||
start = idat.index(b"IDAT") + 4
|
||||
assert idat[start : start + 512] not in json.dumps(record).encode()
|
||||
assert len(json.dumps(record)) < path.stat().st_size // 10
|
||||
|
||||
def test_an_unreadable_container_still_yields_a_record(self, tmp_path: Path):
|
||||
path = tmp_path / "junk.bin"
|
||||
path.write_bytes(b"\x00\x01\x02not an image at all" * 100)
|
||||
|
||||
record = collect_metadata_record(path)
|
||||
|
||||
assert record["container"] == "unknown"
|
||||
assert json.dumps(record) # serializable, and no exception on the way here
|
||||
_assert_same_verdict(path)
|
||||
|
||||
def test_a_missing_file_does_not_raise(self, tmp_path: Path):
|
||||
"""Collection runs over whatever a caller hands it, including a path that
|
||||
vanished between listing and reading."""
|
||||
record = collect_metadata_record(tmp_path / "gone.png")
|
||||
|
||||
assert record["container"] == "unknown"
|
||||
assert record["metadata_base64"] == ""
|
||||
|
||||
|
||||
def test_a_webp_record_matches(tmp_path: Path):
|
||||
"""RIFF has its own walk; a chunk kept or dropped wrongly shows up here."""
|
||||
path = tmp_path / "image.webp"
|
||||
Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "WEBP", xmp=b"<x:xmpmeta>plain</x:xmpmeta>")
|
||||
|
||||
assert collect_metadata_record(path)["container"] == "webp"
|
||||
_assert_same_verdict(path)
|
||||
|
||||
|
||||
def test_the_record_never_reopens_the_source(tmp_path: Path, monkeypatch):
|
||||
"""The reason the record exists: the verdict must run where the file is not."""
|
||||
path = FIXTURES / "chatgpt-1.png" if (FIXTURES / "chatgpt-1.png").exists() else _noise_png(tmp_path / "x.png")
|
||||
record = json.loads(json.dumps(collect_metadata_record(path)))
|
||||
|
||||
def fail_if_called(*args, **kwargs):
|
||||
raise AssertionError("the record path must not open the source file")
|
||||
|
||||
monkeypatch.setattr("builtins.open", fail_if_called)
|
||||
monkeypatch.setattr(Path, "open", fail_if_called)
|
||||
monkeypatch.setattr(Image, "open", fail_if_called)
|
||||
|
||||
report = identify_from_evidence(evidence_from_metadata_record(record, path=path))
|
||||
|
||||
assert isinstance(report, ProvenanceReport)
|
||||
Reference in New Issue
Block a user