mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
Merge remote-tracking branch 'origin/main' into research/video-synthid-quality-groundwork
This commit is contained in:
Victor Kuznetsov
@@ -1,436 +0,0 @@
|
||||
"""AI-generation scorer for scanned datasets.
|
||||
|
||||
Trains a gradient-boosted classifier on the metadata-derived labels of a
|
||||
scan_dataset.py output and scores every file, WITHOUT needing the
|
||||
metadata to be present at scoring time: the features are pixel and
|
||||
container statistics, so a metadata-stripped file still gets a score.
|
||||
The model is distribution-specific by design; evaluate it again before
|
||||
using it on an unrelated target distribution.
|
||||
|
||||
Modes:
|
||||
uv run --with scikit-learn python scripts/ai_score.py train <scan_glob> <model.pkl> [v1|v2]
|
||||
uv run --with scikit-learn python scripts/ai_score.py score <scan_glob> <model.pkl> <out.jsonl>
|
||||
|
||||
<scan_glob> is a glob of scan_dataset shards, e.g. 'data/scan/part_*.jsonl'.
|
||||
score output: one JSON line per file with file, sha256, ai_score, and the
|
||||
label evidence when a metadata label existed (for monitoring drift).
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import math
|
||||
import pickle
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import numpy as np
|
||||
|
||||
# AI-generator names in C2PA claim_generator_info (metadata-level truth).
|
||||
AI_GENERATORS = (
|
||||
"openai",
|
||||
"adobe firefly",
|
||||
"adobe_firefly",
|
||||
"microsoft responsible ai",
|
||||
"microsoft_designer",
|
||||
"black forest labs",
|
||||
"fal-ai",
|
||||
"bria",
|
||||
"chatgpt",
|
||||
"stability",
|
||||
"dreamina",
|
||||
"canva",
|
||||
)
|
||||
# Human-origin software (negative evidence, NOT proof by itself).
|
||||
HUMAN_SOFTWARE = (
|
||||
"photoshop",
|
||||
"lightroom",
|
||||
"picsart",
|
||||
"snapseed",
|
||||
"paint.net",
|
||||
"gimp",
|
||||
"capture one",
|
||||
"meitu",
|
||||
"xingtu",
|
||||
"snow",
|
||||
)
|
||||
|
||||
_FEAT_SCALAR = [
|
||||
("noise", "noise_std"),
|
||||
("noise", "noise_kurtosis"),
|
||||
("fft", "cfa_peak"),
|
||||
("ela", "ela_mean"),
|
||||
("ela", "ela_p95"),
|
||||
("gradient", "laplacian_var"),
|
||||
("color", "saturation_mean"),
|
||||
("color", "value_mean"),
|
||||
("dct", "benford_mad"),
|
||||
]
|
||||
FeatureSchema = Literal["v1", "v2"]
|
||||
_FFT_BANDS = 8
|
||||
_GRADIENT_BINS = 10
|
||||
_COLOR_BINS = 64
|
||||
_DCT_AC_POSITIONS = 8
|
||||
_DCT_AC_BINS = 21
|
||||
_JPEG_QUANT_TABLES = 2
|
||||
_JPEG_QUANT_VALUES = 64
|
||||
_SCORE_BATCH = 4096
|
||||
|
||||
_V1_FEATURE_NAMES = (
|
||||
*(key for _, key in _FEAT_SCALAR),
|
||||
*(f"fft_band_energy_{index}" for index in range(_FFT_BANDS)),
|
||||
*(f"gradient_hist_{index}" for index in range(_GRADIENT_BINS)),
|
||||
*(f"color_hist_4x4x4_{index}" for index in range(_COLOR_BINS)),
|
||||
"jpeg_444",
|
||||
"jpeg_progressive",
|
||||
"megapixels",
|
||||
"aspect_ratio",
|
||||
"format_jpeg",
|
||||
"format_png",
|
||||
)
|
||||
_V2_EXTRA_FEATURE_NAMES = (
|
||||
"cfa_peak_0",
|
||||
"cfa_peak_1",
|
||||
*(f"dct_ac_{position}_{bin_index}" for position in range(_DCT_AC_POSITIONS) for bin_index in range(_DCT_AC_BINS)),
|
||||
*(f"jpeg_quant_{table}_{index}" for table in range(_JPEG_QUANT_TABLES) for index in range(_JPEG_QUANT_VALUES)),
|
||||
"jpeg_quant_table_count",
|
||||
"jpeg_huffman_table_count",
|
||||
"jpeg_huffman_total_bytes",
|
||||
"jpeg_scan_count",
|
||||
"jpeg_restart_interval",
|
||||
"jpeg_precision_bits",
|
||||
"jpeg_adobe_transform",
|
||||
"jpeg_jfif_present",
|
||||
"format_webp",
|
||||
"format_isobmff",
|
||||
"format_other",
|
||||
)
|
||||
_SCHEMA_FEATURE_NAMES: dict[str, tuple[str, ...]] = {
|
||||
"v1": _V1_FEATURE_NAMES,
|
||||
"v2": _V1_FEATURE_NAMES + _V2_EXTRA_FEATURE_NAMES,
|
||||
}
|
||||
|
||||
|
||||
def parse_schema(value: Any) -> FeatureSchema:
|
||||
"""Validate a feature-schema token (CLI argument or model-bundle field)."""
|
||||
if value == "v1" or value == "v2":
|
||||
return value
|
||||
raise ValueError(f"Unknown feature schema: {value!r}")
|
||||
|
||||
|
||||
def feature_names(schema: FeatureSchema = "v1") -> tuple[str, ...]:
|
||||
"""Return the stable ordered feature names for a model schema."""
|
||||
return _SCHEMA_FEATURE_NAMES[parse_schema(schema)]
|
||||
|
||||
|
||||
def _fixed_values(values: Any, count: int, *, normalize: bool = False) -> list[float]:
|
||||
"""Return a fixed-width float vector, padding absent values with NaN."""
|
||||
if not isinstance(values, (list, tuple)):
|
||||
return [float("nan")] * count
|
||||
result = [float(value) for value in values[:count]]
|
||||
if normalize and result:
|
||||
total = sum(value for value in result if math.isfinite(value))
|
||||
if total > 0:
|
||||
result = [value / total for value in result]
|
||||
return result + [float("nan")] * (count - len(result))
|
||||
|
||||
|
||||
def _optional_float(mapping: dict[str, Any], key: str) -> float:
|
||||
value = mapping.get(key)
|
||||
return float(value) if isinstance(value, (int, float)) else float("nan")
|
||||
|
||||
|
||||
def label_of(record: dict[str, Any]) -> int | None:
|
||||
"""1 = strong metadata AI, 0 = human-origin, None = unlabeled."""
|
||||
store = record.get("c2pa_store") or {}
|
||||
for m in (store.get("manifests") or {}).values():
|
||||
for cgi in m.get("claim_generator_info") or []:
|
||||
if any(g in str(cgi.get("name", "")).lower() for g in AI_GENERATORS):
|
||||
return 1
|
||||
for a in m.get("assertions") or []:
|
||||
d = a.get("data") if isinstance(a.get("data"), dict) else {}
|
||||
if "trainedAlgorithmicMedia" in str(d.get("digitalSourceType", "")):
|
||||
return 1
|
||||
for c in record.get("png_chunks", []):
|
||||
t = c.get("text", "")
|
||||
kw = t.split("\x00")[0]
|
||||
if kw in ("prompt", "workflow", "parameters") or "AIGC" in t[:80]:
|
||||
return 1
|
||||
exif = record.get("exif", {})
|
||||
z = exif.get("0th") or {}
|
||||
e = exif.get("Exif") or {}
|
||||
if z.get("Make") and z.get("Model") and (e.get("MakerNote") or e.get("LensModel") or e.get("LensMake")):
|
||||
return 0
|
||||
for c in record.get("png_chunks", []):
|
||||
if c.get("apple_screenshot_marker"):
|
||||
return 0
|
||||
for v in (record.get("iptc") or {}).values():
|
||||
if str(v).strip() == "Screenshot":
|
||||
return 0
|
||||
if any(t in str(z.get("Software", "")).lower() for t in HUMAN_SOFTWARE):
|
||||
return 0
|
||||
return None
|
||||
|
||||
|
||||
def features_of(record: dict[str, Any], *, schema: FeatureSchema = "v1") -> list[float]:
|
||||
"""The structural feature vector (pixel + container stats, no metadata)."""
|
||||
v = [float((record.get(s) or {}).get(k, np.nan)) for s, k in _FEAT_SCALAR]
|
||||
fft = record.get("fft") or {}
|
||||
gradient = record.get("gradient") or {}
|
||||
color = record.get("color") or {}
|
||||
v.extend(_fixed_values(fft.get("fft_band_energy"), _FFT_BANDS))
|
||||
v.extend(_fixed_values(gradient.get("gradient_hist"), _GRADIENT_BINS))
|
||||
v.extend(_fixed_values(color.get("color_hist_4x4x4"), _COLOR_BINS))
|
||||
jf = record.get("jpeg_forensics") or {}
|
||||
v.append(1.0 if jf.get("subsampling") == "4:4:4" else 0.0)
|
||||
v.append(1.0 if jf.get("progressive") else 0.0)
|
||||
pil = record.get("pil") or {}
|
||||
w, h = pil.get("width") or 0, pil.get("height") or 0
|
||||
v += [float(w * h) / 1e6, float(w) / max(h, 1)]
|
||||
fmt = record.get("content_format")
|
||||
v += [1.0 if fmt == "jpeg" else 0.0, 1.0 if fmt == "png" else 0.0]
|
||||
if schema == "v2":
|
||||
v.extend(_fixed_values(fft.get("cfa_peaks"), 2))
|
||||
dct_hist = (record.get("dct") or {}).get("dct_ac_hist")
|
||||
for position in range(_DCT_AC_POSITIONS):
|
||||
values = dct_hist[position] if isinstance(dct_hist, list) and position < len(dct_hist) else None
|
||||
v.extend(_fixed_values(values, _DCT_AC_BINS, normalize=True))
|
||||
quant_tables = jf.get("quant_tables")
|
||||
for table in range(_JPEG_QUANT_TABLES):
|
||||
values = quant_tables.get(str(table)) if isinstance(quant_tables, dict) else None
|
||||
v.extend(_fixed_values(values, _JPEG_QUANT_VALUES))
|
||||
is_jpeg = fmt == "jpeg"
|
||||
huffman_tables = jf.get("huffman_tables_hex")
|
||||
huffman_values = huffman_tables if isinstance(huffman_tables, list) else []
|
||||
v.extend(
|
||||
[
|
||||
float(len(quant_tables)) if isinstance(quant_tables, dict) else (0.0 if is_jpeg else float("nan")),
|
||||
float(len(huffman_values)) if is_jpeg else float("nan"),
|
||||
(
|
||||
float(sum(len(value) // 2 for value in huffman_values if isinstance(value, str)))
|
||||
if is_jpeg
|
||||
else float("nan")
|
||||
),
|
||||
_optional_float(jf, "scan_count"),
|
||||
_optional_float(jf, "restart_interval"),
|
||||
_optional_float(jf, "precision_bits"),
|
||||
_optional_float(jf, "adobe_transform"),
|
||||
1.0 if jf.get("jfif") else (0.0 if is_jpeg else float("nan")),
|
||||
1.0 if fmt == "webp" else 0.0,
|
||||
1.0 if isinstance(fmt, str) and fmt.startswith("isobmff:") else 0.0,
|
||||
(
|
||||
1.0
|
||||
if fmt is not None and fmt not in {"jpeg", "png", "webp"} and not str(fmt).startswith("isobmff:")
|
||||
else 0.0
|
||||
),
|
||||
]
|
||||
)
|
||||
elif schema != "v1":
|
||||
raise ValueError(f"Unknown feature schema: {schema}")
|
||||
expected = len(feature_names(schema))
|
||||
if len(v) != expected:
|
||||
raise ValueError(f"Feature schema {schema} produced {len(v)} values; expected {expected}")
|
||||
return v
|
||||
|
||||
|
||||
def grouped_stratified_split(
|
||||
labels: np.ndarray,
|
||||
hashes: np.ndarray,
|
||||
*,
|
||||
test_size: float = 0.3,
|
||||
random_state: int = 0,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Split records by hash while preserving label balance across hash groups."""
|
||||
if len(labels) != len(hashes):
|
||||
raise ValueError("labels and hashes must have equal length")
|
||||
if not 0 < test_size < 1:
|
||||
raise ValueError("test_size must be between 0 and 1")
|
||||
groups: dict[str, list[int]] = defaultdict(list)
|
||||
group_labels: dict[str, int] = {}
|
||||
for index, (label, digest) in enumerate(zip(labels, hashes, strict=True)):
|
||||
key = str(digest) if digest else f"__row_{index}"
|
||||
int_label = int(label)
|
||||
if key in group_labels and group_labels[key] != int_label:
|
||||
raise ValueError(f"Hash {key!r} has conflicting labels")
|
||||
group_labels[key] = int_label
|
||||
groups[key].append(index)
|
||||
keys_by_label: dict[int, list[str]] = defaultdict(list)
|
||||
for key, label in group_labels.items():
|
||||
keys_by_label[label].append(key)
|
||||
rng = np.random.default_rng(random_state)
|
||||
test_set: set[str] = set()
|
||||
for class_keys in keys_by_label.values():
|
||||
if len(class_keys) < 2:
|
||||
raise ValueError("Each label needs at least two distinct hash groups")
|
||||
shuffled = np.asarray(class_keys)
|
||||
rng.shuffle(shuffled)
|
||||
test_count = min(max(round(len(shuffled) * test_size), 1), len(shuffled) - 1)
|
||||
test_set.update(shuffled[:test_count].tolist())
|
||||
train_set = set(groups) - test_set
|
||||
train = np.asarray([index for key, indices in groups.items() if key in train_set for index in indices])
|
||||
test = np.asarray([index for key, indices in groups.items() if key in test_set for index in indices])
|
||||
return np.sort(train), np.sort(test)
|
||||
|
||||
|
||||
def temporal_holdout_split(
|
||||
dates: np.ndarray,
|
||||
hashes: np.ndarray,
|
||||
*,
|
||||
train_fraction: float = 0.7,
|
||||
) -> tuple[np.ndarray, np.ndarray, str]:
|
||||
"""Split chronologically and remove later copies of training hashes."""
|
||||
if len(dates) != len(hashes):
|
||||
raise ValueError("dates and hashes must have equal length")
|
||||
if not 0 < train_fraction < 1:
|
||||
raise ValueError("train_fraction must be between 0 and 1")
|
||||
unique_dates = sorted({str(date) for date in dates})
|
||||
if len(unique_dates) < 2:
|
||||
raise ValueError("Temporal holdout needs at least two distinct dates")
|
||||
cutoff_index = min(int(len(unique_dates) * train_fraction), len(unique_dates) - 1)
|
||||
cutoff = unique_dates[cutoff_index]
|
||||
train = np.flatnonzero(dates < cutoff)
|
||||
test = np.flatnonzero((dates >= cutoff) & ~np.isin(hashes, hashes[train]))
|
||||
return train, test, cutoff
|
||||
|
||||
|
||||
def model_schema(bundle: dict[str, Any]) -> FeatureSchema:
|
||||
"""Read a model's feature schema, defaulting legacy bundles to v1."""
|
||||
return parse_schema(bundle.get("feature_schema", "v1"))
|
||||
|
||||
|
||||
def iter_records(pattern: str) -> Any:
|
||||
for path in sorted(glob.glob(pattern)):
|
||||
with open(path) as fh:
|
||||
for line in fh:
|
||||
yield json.loads(line)
|
||||
|
||||
|
||||
def cmd_train(pattern: str, model_path: str, schema: FeatureSchema = "v2") -> None:
|
||||
from sklearn.ensemble import HistGradientBoostingClassifier
|
||||
from sklearn.impute import SimpleImputer
|
||||
from sklearn.metrics import average_precision_score, roc_auc_score
|
||||
|
||||
# keep the feature vector, not the record: a parsed scan record is an order of
|
||||
# magnitude larger than the row it collapses to.
|
||||
labeled: dict[str, tuple[list[float], int, str]] = {}
|
||||
missing_hash_index = 0
|
||||
for r in iter_records(pattern):
|
||||
if "noise" not in r:
|
||||
continue
|
||||
lab = label_of(r)
|
||||
if lab is None:
|
||||
continue
|
||||
digest = r.get("sha256")
|
||||
if not digest:
|
||||
digest = f"__missing_{missing_hash_index}"
|
||||
missing_hash_index += 1
|
||||
date = Path(r["file"]).parent.name
|
||||
previous = labeled.get(str(digest))
|
||||
if previous is not None and previous[1] != lab:
|
||||
raise ValueError(f"Hash {digest!r} has conflicting labels")
|
||||
if previous is None or date < previous[2]:
|
||||
labeled[str(digest)] = (features_of(r, schema=schema), lab, date)
|
||||
|
||||
rows, labels, row_dates = zip(*labeled.values(), strict=True)
|
||||
X = np.asarray(rows)
|
||||
y = np.asarray(labels)
|
||||
dates = np.asarray(row_dates)
|
||||
hashes = np.asarray(list(labeled))
|
||||
print(
|
||||
f"unique labeled: {len(y)} "
|
||||
f"(pos {int(y.sum())}, neg {int((1 - y).sum())}, features {X.shape[1]}, schema {schema})"
|
||||
)
|
||||
|
||||
def fit(train: np.ndarray) -> tuple[Any, Any]:
|
||||
rows = X[train]
|
||||
imputer = SimpleImputer(strategy="median")
|
||||
classifier = HistGradientBoostingClassifier(
|
||||
random_state=0,
|
||||
max_iter=300,
|
||||
l2_regularization=5.0,
|
||||
early_stopping=False,
|
||||
)
|
||||
classifier.fit(imputer.fit_transform(rows), y[train])
|
||||
return imputer, classifier
|
||||
|
||||
train, test = grouped_stratified_split(y, hashes, test_size=0.3, random_state=0)
|
||||
imp, clf = fit(train)
|
||||
p = clf.predict_proba(imp.transform(X[test]))[:, 1]
|
||||
print(f"SHA-grouped holdout: AUC {roc_auc_score(y[test], p):.4f} | AP {average_precision_score(y[test], p):.4f}")
|
||||
|
||||
temporal_train, temporal_test, cutoff = temporal_holdout_split(dates, hashes)
|
||||
if len(temporal_train) > 100 and len(temporal_test) > 100:
|
||||
imp_t, clf_t = fit(temporal_train)
|
||||
pt = clf_t.predict_proba(imp_t.transform(X[temporal_test]))[:, 1]
|
||||
auc_t = roc_auc_score(y[temporal_test], pt)
|
||||
ap_t = average_precision_score(y[temporal_test], pt)
|
||||
print(f"temporal (<{cutoff}): AUC {auc_t:.4f} | AP {ap_t:.4f}")
|
||||
|
||||
all_indices = np.arange(len(y))
|
||||
imp, clf = fit(all_indices)
|
||||
bundle = {
|
||||
"imputer": imp,
|
||||
"clf": clf,
|
||||
"feature_schema": schema,
|
||||
"feature_names": feature_names(schema),
|
||||
}
|
||||
with open(model_path, "wb") as f:
|
||||
pickle.dump(bundle, f)
|
||||
print(f"model written: {model_path}")
|
||||
|
||||
|
||||
def cmd_score(pattern: str, model_path: str, out_path: str) -> None:
|
||||
# the model file is produced locally by `train`; do not load third-party pickles
|
||||
with open(model_path, "rb") as f:
|
||||
bundle = pickle.load(f) # noqa: S301
|
||||
imp, clf = bundle["imputer"], bundle["clf"]
|
||||
schema = model_schema(bundle)
|
||||
stored_names = bundle.get("feature_names")
|
||||
if stored_names is not None and tuple(stored_names) != feature_names(schema):
|
||||
raise ValueError(f"Model feature names do not match schema {schema}")
|
||||
n = 0
|
||||
batch_rows: list[list[float]] = []
|
||||
batch_meta: list[dict[str, Any]] = []
|
||||
|
||||
def flush(out: Any) -> None:
|
||||
"""Score one batch: per-record predict_proba is dominated by call overhead."""
|
||||
nonlocal n
|
||||
if not batch_rows:
|
||||
return
|
||||
scores = clf.predict_proba(imp.transform(np.asarray(batch_rows)))[:, 1]
|
||||
for meta, score in zip(batch_meta, scores, strict=True):
|
||||
out.write(json.dumps({**meta, "ai_score": round(float(score), 4)}) + "\n")
|
||||
n += len(batch_rows)
|
||||
batch_rows.clear()
|
||||
batch_meta.clear()
|
||||
print(f" {n}", flush=True)
|
||||
|
||||
with open(out_path, "w") as out:
|
||||
for r in iter_records(pattern):
|
||||
if "noise" not in r:
|
||||
continue
|
||||
batch_rows.append(features_of(r, schema=schema))
|
||||
batch_meta.append({"file": r["file"], "sha256": r.get("sha256"), "metadata_label": label_of(r)})
|
||||
if len(batch_rows) >= _SCORE_BATCH:
|
||||
flush(out)
|
||||
flush(out)
|
||||
print(f"scored {n} -> {out_path}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if len(sys.argv) < 2 or sys.argv[1] not in ("train", "score"):
|
||||
print(__doc__)
|
||||
sys.exit(2)
|
||||
if sys.argv[1] == "train" and len(sys.argv) in (4, 5):
|
||||
schema = parse_schema(sys.argv[4]) if len(sys.argv) == 5 else "v2"
|
||||
cmd_train(sys.argv[2], sys.argv[3], schema)
|
||||
elif sys.argv[1] == "score" and len(sys.argv) == 5:
|
||||
cmd_score(sys.argv[2], sys.argv[3], sys.argv[4])
|
||||
else:
|
||||
print(__doc__)
|
||||
sys.exit(2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,241 @@
|
||||
"""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 <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
|
||||
|
||||
# The package's OWN tree, not the repository root: from a worktree, an editable
|
||||
# install resolves `remove_ai_watermarks` to the MAIN checkout, so a script measuring
|
||||
# this tree would silently import a different one.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
||||
|
||||
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())
|
||||
@@ -0,0 +1,224 @@
|
||||
"""Audit the record path against the file path over a whole dataset.
|
||||
|
||||
WHY THIS EXISTS
|
||||
|
||||
Two seams reach the same provenance verdict:
|
||||
|
||||
identify(path, check_visible=False, check_invisible=False)
|
||||
|
||||
identify_metadata_record(collect_metadata_record(path), path=path)
|
||||
|
||||
Their equality is the record's entire contract, 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; those cover the signal families we already know about. This covers the ones
|
||||
we do not: every real placement in a real corpus, which is where all three defects
|
||||
found so far actually came from.
|
||||
|
||||
The record is round-tripped through ``json.dumps``/``loads`` before it is judged, so
|
||||
a value that only survives in memory fails here rather than at a customer.
|
||||
|
||||
WHAT IT REPORTS
|
||||
|
||||
One JSONL row per image: both verdicts, whether they agree, the record size, and any
|
||||
exception from either side. The summary counts disagreements by field and by signal,
|
||||
so "the record lost samsung_genai on 13 files" reads directly off the output instead
|
||||
of being reconstructed.
|
||||
|
||||
Pass ``--baseline`` with an earlier run to also diff against it. That answers the
|
||||
other question a detection change raises: which files changed verdict, and are they
|
||||
exactly the ones that were meant to.
|
||||
|
||||
DATA SAFETY
|
||||
|
||||
Read-only over a local dataset. Writes only the given output path, which belongs
|
||||
outside the repository. Resumable: rerunning skips files already recorded.
|
||||
|
||||
uv run python scripts/record_parity_audit.py data/spaces/originals .local-eval/parity.jsonl
|
||||
uv run python scripts/record_parity_audit.py <dataset> <out> --baseline .local-eval/previous.jsonl
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import collections
|
||||
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 Iterator
|
||||
|
||||
# The package's OWN tree, not the repository root: from a worktree, an editable
|
||||
# install resolves `remove_ai_watermarks` to the MAIN checkout, so a script measuring
|
||||
# this tree would silently import a different one.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
|
||||
|
||||
from remove_ai_watermarks.identify import identify, identify_metadata_record
|
||||
from remove_ai_watermarks.metadata_record import collect_metadata_record
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SUPPORTED = frozenset({".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif", ".avif"})
|
||||
# Every field of the verdict a caller can act on. `path` is excluded: it is extraction
|
||||
# context, and the two paths are handed the same one by construction.
|
||||
COMPARED = ("is_ai_generated", "platform", "confidence", "ai_source_kind", "ai_from_metadata")
|
||||
|
||||
|
||||
def _verdict(report: Any) -> dict[str, Any]:
|
||||
return {
|
||||
**{field: getattr(report, field) for field in COMPARED},
|
||||
"signals": sorted(signal.name for signal in report.signals),
|
||||
"watermarks": sorted(report.watermarks),
|
||||
}
|
||||
|
||||
|
||||
def _audit(path: Path) -> dict[str, Any]:
|
||||
row: dict[str, Any] = {"path": str(path)}
|
||||
try:
|
||||
row["bytes"] = path.stat().st_size
|
||||
except OSError as exc:
|
||||
return {**row, "error": f"stat: {exc}"}
|
||||
|
||||
try:
|
||||
started = time.perf_counter()
|
||||
record = json.loads(json.dumps(collect_metadata_record(path)))
|
||||
row["collect_ms"] = (time.perf_counter() - started) * 1000
|
||||
row["record_bytes"] = len(json.dumps(record))
|
||||
row["container"] = record.get("container")
|
||||
via_record = _verdict(identify_metadata_record(record, path=path))
|
||||
except Exception as exc:
|
||||
return {**row, "error": f"record path: {type(exc).__name__}: {exc}"}
|
||||
|
||||
try:
|
||||
via_file = _verdict(identify(path, check_visible=False, check_invisible=False))
|
||||
except Exception as exc:
|
||||
return {**row, "error": f"file path: {type(exc).__name__}: {exc}"}
|
||||
|
||||
row["record"] = via_record
|
||||
row["file"] = via_file
|
||||
row["agree"] = via_record == via_file
|
||||
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(out_path: Path) -> set[str]:
|
||||
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 _summarize(rows: list[dict[str, Any]], baseline: Path | None) -> None:
|
||||
failed = [r for r in rows if "error" in r]
|
||||
usable = [r for r in rows if "error" not in r]
|
||||
disagreed = [r for r in usable if not r["agree"]]
|
||||
|
||||
print(f"\nimages: {len(rows)} errors: {len(failed)} compared: {len(usable)}")
|
||||
print(f"record path disagrees with file path: {len(disagreed)}")
|
||||
for row in failed[:10]:
|
||||
print(f" ERROR {Path(row['path']).name}: {row['error']}")
|
||||
|
||||
fields: collections.Counter[str] = collections.Counter()
|
||||
for row in disagreed:
|
||||
fields.update(field for field in COMPARED if row["record"][field] != row["file"][field])
|
||||
for name in set(row["file"]["signals"]) - set(row["record"]["signals"]):
|
||||
fields[f"signal missing from record: {name}"] += 1
|
||||
for name in set(row["record"]["signals"]) - set(row["file"]["signals"]):
|
||||
fields[f"signal only in record: {name}"] += 1
|
||||
for label, count in fields.most_common():
|
||||
print(f" {label}: {count}")
|
||||
for row in disagreed[:10]:
|
||||
print(f" {Path(row['path']).name}\n record: {row['record']}\n file: {row['file']}")
|
||||
|
||||
if baseline is None:
|
||||
return
|
||||
previous = {}
|
||||
with baseline.open(encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except ValueError:
|
||||
continue
|
||||
if "error" not in item:
|
||||
previous[Path(item["path"]).name] = item
|
||||
|
||||
changed = []
|
||||
for row in usable:
|
||||
was = previous.get(Path(row["path"]).name)
|
||||
if was is None:
|
||||
continue
|
||||
# The WHOLE verdict, not a chosen subset. A first version compared confidence
|
||||
# and signals only and reported "0 changed" for a run whose single intended
|
||||
# correction was a watermark line -- the change it existed to show.
|
||||
before = was.get("file") or {}
|
||||
if before and before != row["file"]:
|
||||
changed.append((row["path"], before, row["file"]))
|
||||
print(f"\nverdicts changed against the baseline: {len(changed)}")
|
||||
moved: collections.Counter[str] = collections.Counter()
|
||||
for _, before, after in changed:
|
||||
for name in set(after["signals"]) - set(before.get("signals") or []):
|
||||
moved[f"gained signal {name}"] += 1
|
||||
for name in set(before.get("signals") or []) - set(after["signals"]):
|
||||
moved[f"LOST signal {name}"] += 1
|
||||
for field in (*COMPARED, "watermarks"):
|
||||
if before.get(field) != after.get(field):
|
||||
moved[f"{field} changed"] += 1
|
||||
for label, count in moved.most_common():
|
||||
print(f" {label}: {count}")
|
||||
for path, before, after in changed[:10]:
|
||||
print(f" {Path(path).name}\n before: {before}\n after: {after}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("dataset", type=Path)
|
||||
parser.add_argument("out", type=Path)
|
||||
parser.add_argument("--baseline", type=Path, default=None, help="an earlier run to diff verdicts against")
|
||||
parser.add_argument("--limit", type=int, default=0)
|
||||
parser.add_argument("--progress-every", type=int, default=2000)
|
||||
args = parser.parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
|
||||
args.out.parent.mkdir(parents=True, exist_ok=True)
|
||||
done = _done(args.out)
|
||||
if done:
|
||||
log.info("resuming: %d images already audited", len(done))
|
||||
|
||||
processed = 0
|
||||
started = time.monotonic()
|
||||
with args.out.open("a", encoding="utf-8") as handle:
|
||||
for path in _iter_images(args.dataset):
|
||||
if str(path) in done:
|
||||
continue
|
||||
handle.write(json.dumps(_audit(path), ensure_ascii=False, default=str) + "\n")
|
||||
handle.flush()
|
||||
processed += 1
|
||||
if processed % args.progress_every == 0:
|
||||
log.info("%d images, %.1f/s", processed, processed / (time.monotonic() - started))
|
||||
if args.limit and processed >= args.limit:
|
||||
break
|
||||
|
||||
log.info("audited %d images in %.1f s", processed, time.monotonic() - started)
|
||||
with args.out.open(encoding="utf-8") as handle:
|
||||
rows = [json.loads(line) for line in handle]
|
||||
_summarize(rows, args.baseline)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user