From 9a29dcac8a030c244a66226f6f73091cdbaef837 Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Wed, 5 Aug 2026 17:19:47 -0700 Subject: [PATCH] Keep the pixel forensics in the library, drop the ai-score tooling `scripts/ai_score.py` and the dataset scanner that fed it are gone: the detector they trained is not something this project runs, and the corpus lived outside the repository anyway. Nothing else referenced them. The scanner's pixel layer was worth keeping, so it moves into the package as `pixel_evidence.py` -- six families of scale-robust statistics (block-DCT histograms and Benford deviation, FFT band energies and CFA peaks, high-pass residual, error level, gradient, colour) measured in a single shared decode. The arithmetic was verified against the scanner over 60 corpus images, families and artifacts alike, before the scanner was removed; that comparison is no longer possible, which is why the tests now pin behavior instead: determinism, empty-not-wrong on images too small for a family, and one failing family not taking the others with it. It has no consumer. Nothing in the package reads it, and the module says so. `artifacts=True` returns the spatial layer -- perceptual hash, 128px thumbnail, coarse ELA/residual/phase maps. Those identify the source image rather than describe it, so they are opt-in and separate: everything else is a scalar or a fixed-length histogram nothing can be reconstructed from. Co-Authored-By: Claude Opus 5 --- docs/module-internals.md | 15 + scripts/ai_score.py | 436 -------- scripts/scan_dataset.py | 1111 -------------------- src/remove_ai_watermarks/pixel_evidence.py | 400 +++++++ tests/test_ai_score.py | 126 --- tests/test_pixel_evidence.py | 156 +++ 6 files changed, 571 insertions(+), 1673 deletions(-) delete mode 100644 scripts/ai_score.py delete mode 100644 scripts/scan_dataset.py create mode 100644 src/remove_ai_watermarks/pixel_evidence.py delete mode 100644 tests/test_ai_score.py create mode 100644 tests/test_pixel_evidence.py diff --git a/docs/module-internals.md b/docs/module-internals.md index 980045d..3f8ecfe 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -425,6 +425,21 @@ Pixel forensics are deliberately absent: nothing in the provenance path reads th Verdict equivalence is checked over tracked fixtures and a separate local evaluation corpus. +### Experimental pixel forensics + +[`pixel_evidence.py`](../src/remove_ai_watermarks/pixel_evidence.py) measures six +families of scale-robust pixel statistics (block-DCT histograms and Benford +deviation, FFT band energies and CFA peaks, high-pass residual, error level, +gradient, colour) in a single decode, sharing the intermediate maps between them. + +It has no consumer. Nothing in the package reads it -- not the verdict, not removal, +not the CLI -- and the shape is unstable until something does. + +`artifacts=True` additionally returns the spatial layer: a perceptual hash, a 128px +JPEG thumbnail, and coarse ELA, residual and phase maps. Those identify the source +image rather than describe it, which is why they are opt-in and a separate field: a +caller storing them is handling image content, not statistics about it. + 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 diff --git a/scripts/ai_score.py b/scripts/ai_score.py deleted file mode 100644 index d2f06c2..0000000 --- a/scripts/ai_score.py +++ /dev/null @@ -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 [v1|v2] - uv run --with scikit-learn python scripts/ai_score.py score - - 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() diff --git a/scripts/scan_dataset.py b/scripts/scan_dataset.py deleted file mode 100644 index 0fe4b8d..0000000 --- a/scripts/scan_dataset.py +++ /dev/null @@ -1,1111 +0,0 @@ -"""Collect all raw metadata over a dataset, for later offline analysis. -Read-only; analysis is NOT this script's job. For every image it writes -one JSONL record containing: - -- file basics: name, path, size in bytes, container format (content-sniffed), - pixel dimensions, dpi, color mode, sha256, mtime/birthtime -- raw metadata, so nothing has to be re-scanned later: - full EXIF (all IFDs, decoded tag names, MakerNote in full hex), all XMP - packets, all PNG chunks (text chunks decoded and inflated, binary chunks - as base64), all JPEG APP segments (marker + full base64), WebP RIFF - chunks, ISOBMFF box/item inventory for HEIC/AVIF/MOV, IPTC-IIM dataset, - ICC profile (full base64), C2PA manifest store JSON, post-EOI/IEND - trailers, EXIF thumbnail (bytes + its own JPEG structure), JPEG encoder - structure (quant tables, progressive scan script, Huffman tables, - subsampling, JFIF/Adobe markers), file hashes and timestamps, macOS - download provenance xattrs, Live Photo content identifier. Embedded - binary blobs are base64'd in full with a 1 MB ceiling per blob. - -By default everything collected is raw bytes or a mechanical container -decode; no verdicts, no estimates, no edit-trail interpretation. The -optional pixel modes add a clearly-separated derived layer (aggregate -statistics only, still no verdicts). - -Usage: - python scan_dataset.py /path/to/dataset out_prefix [--pixels|--pixels-full] - -Writes out_prefix.jsonl (one record per file) and out_prefix.csv (flat -summary for quick sorting). If the prefix ends with .gz, the JSONL is -gzip-compressed on the fly (3-5x smaller, still streamable line by line). - ---pixels adds aggregate pixel forensics to each record (DCT histograms, -noise/FFT/ELA/gradient/color statistics) from which the image CANNOT be -reconstructed, plus per-section timing_ms for pipeline latency planning. ---pixels-full additionally stores the privacy-lifting artifacts (phash, -128px thumbnail, coarse ELA/noise/FFT-phase maps). Pixel modes require -numpy; the metadata-only mode needs no numpy. - -For a large dataset, parallelize by sharding the input into folders and -running one process per shard (a multiprocessing pool hangs on macOS once -cv2/PIL are loaded; independent processes do not): - - for d in /dataset/*/; do - python scan_dataset.py "$d" "out_$(basename "$d")" & - done - wait - cat out_*.jsonl > dataset.jsonl # or just read the shards lazily - -Rerunning with the same prefix resumes: files already present in the -output are skipped, new files are appended. - -Reading big results: never load the JSONL whole. Stream it: -polars.scan_ndjson (lazy), pandas.read_json(lines=True, chunksize=...), -or a plain line loop. One line = one self-contained JSON record. - -Dependencies: pip install pillow piexif c2pa-python pillow-heif -(pillow-heif is only needed for HEIC/AVIF inputs). -""" - -from __future__ import annotations - -import base64 -import contextlib -import csv -import io -import json -import struct -import sys -import time -from pathlib import Path -from typing import Any - -from PIL import Image -from PIL.IptcImagePlugin import getiptcinfo - -# Pillow cannot open HEIC/HEIF without this opener, and it does not auto-register. -# Skipping it does not fail loudly: HEIC records lose EXIF and most pixel features -# while the scan continues as if the container were merely unreadable. Deliberately a -# copy of `image_io._register_heif` rather than an import: this script keeps the -# minimal dependency set its docstring advertises and never imports the package. The -# suppression is as wide as the original's, so a broken libheif degrades the scan -# instead of killing it at import. -with contextlib.suppress(Exception): - import pillow_heif - - pillow_heif.register_heif_opener() - -SUPPORTED = { - ".png", - ".jpg", - ".jpeg", - ".webp", - ".heic", - ".heif", - ".avif", - ".tif", - ".tiff", - ".bmp", - ".gif", - # video/px containers: no pixel decode, but C2PA reads them (Sora/Veo - # carry C2PA manifests) and the byte scans still apply - ".mp4", - ".mov", - ".m4v", - ".jxl", -} - -_B64_CAP = 1 << 20 # 1 MB safety ceiling per embedded blob - - -def _safe_str(v: Any) -> str: - try: - return str(v) - except Exception: - return repr(v) - - -def _b64(b: bytes) -> str: - """Full base64 of a binary blob, with a 1 MB safety ceiling.""" - if len(b) > _B64_CAP: - return base64.b64encode(b[:_B64_CAP]).decode("ascii") + f"...TRUNCATED({len(b)} bytes total)" - return base64.b64encode(b).decode("ascii") - - -def _decode_exif_value(v: Any) -> Any: - """Make a piexif value JSON-safe; bytes are kept in full as hex.""" - if isinstance(v, bytes): - if len(v) <= 64: - try: - return v.decode("utf-8", "strict") - except (UnicodeDecodeError, ValueError): - return f"hex:{v.hex()}" - return f"hex:{v.hex()}" - if isinstance(v, (tuple, list)): - return [_decode_exif_value(x) for x in v] - return v - - -def read_full_exif( - path: Path, exif_blob: bytes | None = None, data: bytes | None = None -) -> tuple[dict[str, Any], bytes | None]: - """All EXIF IFDs with decoded tag names (piexif, no re-encode), plus the - raw embedded-thumbnail bytes for the caller's own thumbnail forensics. - - ``exif_blob`` is the PIL-exposed EXIF blob (PNG/WebP/HEIC path) so the - caller's single Image.open is not repeated here. ``data`` is the - already-read file bytes so piexif does not re-read the file.""" - import piexif - - try: - exif = piexif.load(data) if data is not None else piexif.load(str(path)) - except Exception: - if not exif_blob: - return {}, None - try: - exif = piexif.load(exif_blob) - except Exception as exc: - return {"error": _safe_str(exc)}, None - out: dict[str, Any] = {} - thumbnail: bytes | None = None - for ifd, tags in exif.items(): - if ifd == "thumbnail": - thumbnail = tags if isinstance(tags, bytes) else None - out["thumbnail"] = f"{len(tags)} bytes" if isinstance(tags, bytes) else None - continue - if not isinstance(tags, dict): - continue - tag_names = piexif.TAGS.get(ifd, {}) - decoded = {} - for t, v in tags.items(): - name = tag_names.get(t, {}).get("name", f"tag_{t}") - if name == "MakerNote" and isinstance(v, bytes): - # Full hex, no cap: some Apple manifests are several kilobytes. - # but Canon reaches 28 KB and Sony 38 KB (AF data, serials, - # embedded previews) -- a cap would silently drop exactly the - # camera-original evidence this scan exists to preserve - decoded[name] = f"hex:{v.hex()}" - else: - decoded[name] = _decode_exif_value(v) - out[ifd] = decoded - return out, thumbnail - - -def _png_text_decode(ctype: str, body: bytes) -> str: - """Decode a tEXt/zTXt/iTXt chunk, inflating zlib where used. - - The compressed forms are where ComfyUI / Automatic1111 hide the - generation workflow and prompt, so skipping the inflate would drop - the strongest AI-provenance text a PNG can carry.""" - import zlib - - if ctype == "tEXt": - return body.decode("utf-8", "replace") - if ctype == "zTXt": - nul = body.find(b"\x00") - if nul == -1: - return body.decode("utf-8", "replace") - keyword = body[:nul].decode("latin-1", "replace") - # body[nul+1] = compression method (0 = zlib) - try: - text = zlib.decompress(body[nul + 2 :]).decode("utf-8", "replace") - except zlib.error: - text = body.decode("utf-8", "replace") - return f"{keyword}\x00{text}" - # iTXt: keyword\0 compflag(1) compmethod(1) lang\0 translated\0 text - parts = body.split(b"\x00", 1) - if len(parts) < 2: - return body.decode("utf-8", "replace") - keyword = parts[0].decode("latin-1", "replace") - rest = parts[1] - if len(rest) < 2: - return body.decode("utf-8", "replace") - compflag = rest[0] - tail = rest[2:] - for _ in range(2): # skip language tag and translated keyword - nul = tail.find(b"\x00") - if nul == -1: - return body.decode("utf-8", "replace") - tail = tail[nul + 1 :] - if compflag: - with contextlib.suppress(zlib.error): - tail = zlib.decompress(tail) - return f"{keyword}\x00{tail.decode('utf-8', 'replace')}" - - -def read_png_chunks(data: bytes) -> tuple[list[dict[str, Any]], int]: - """Every PNG chunk in order (type, length; text chunks decoded and - inflated, binary chunks as base64) plus the post-IEND trailer size.""" - chunks: list[dict[str, Any]] = [] - post_iend = 0 - try: - pos = 8 - while pos + 12 <= len(data): - length = struct.unpack(">I", data[pos : pos + 4])[0] - ctype = data[pos + 4 : pos + 8].decode("latin-1") - body = data[pos + 8 : pos + 8 + length] - entry: dict[str, Any] = {"type": ctype, "length": length} - if ctype in ("tEXt", "zTXt", "iTXt"): - entry["text"] = _png_text_decode(ctype, body) - if entry["text"].startswith("XML:com.adobe.xmp"): - entry["kind"] = "xmp" - elif ctype == "tIME" and length == 7: - y, mo, d, h, mi, s = struct.unpack(">HBBBBB", body) - entry["time"] = f"{y:04d}-{mo:02d}-{d:02d}T{h:02d}:{mi:02d}:{s:02d}Z" - elif ctype == "gAMA" and length == 4: - entry["gamma"] = struct.unpack(">I", body)[0] / 100000 - elif ctype == "sRGB" and length == 1: - entry["rendering_intent"] = body[0] - elif ctype == "iCCP": - nul = body.find(b"\x00") - if nul > 0: - entry["profile_name"] = body[:nul].decode("latin-1", "replace") - entry["base64"] = _b64(body) - elif ctype == "iDOT": - # present in iOS/macOS screenshots - entry["apple_screenshot_marker"] = True - elif ctype in ("IHDR", "IDAT"): - pass # pixel-data / header chunks: length is signal enough - elif length: - entry["base64"] = _b64(body) - chunks.append(entry) - pos += 12 + length - if ctype == "IEND": - post_iend = len(data) - pos - break - except Exception as exc: - chunks.append({"error": _safe_str(exc)}) - return chunks, post_iend - - -def read_jpeg_segments(data: bytes) -> dict[str, Any]: - """Every JPEG APP segment in order, plus post-EOI trailer size. - - XMP APP1 segments are kept as full text; every other segment body is - kept as full base64 (1 MB ceiling per segment). - """ - result: dict[str, Any] = {"segments": [], "post_eoi_bytes": 0} - try: - pos = 2 - while pos + 4 <= len(data): - if data[pos] != 0xFF: - break - marker = data[pos + 1] - if marker == 0xD9: # EOI - result["post_eoi_bytes"] = len(data) - (pos + 2) - break - if marker == 0xDA: # SOS: entropy-coded data follows - eoi = data.rfind(b"\xff\xd9") - if eoi != -1: - result["post_eoi_bytes"] = len(data) - (eoi + 2) - break - if not (0xE0 <= marker <= 0xEF): - length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] - pos += 2 + length - continue - length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] - body = data[pos + 4 : pos + 2 + length] - name = f"APP{marker - 0xE0}" - entry: dict[str, Any] = {"marker": name, "length": length} - if body.startswith(b"http://ns.adobe.com/xap/1.0/\x00"): - entry["kind"] = "xmp" - entry["text"] = body[29:].decode("utf-8", "replace") - elif name == "APP2" and body.startswith(b"MPF\x00"): - # Multi-Picture Format: Ultra HDR gain map, Samsung dual shot - entry["kind"] = "mpf" - entry["base64"] = _b64(body) - elif name == "APP2" and body.startswith(b"ICC_PROFILE"): - entry["kind"] = "icc" - entry["base64"] = _b64(body) - elif name == "APP2" and body.startswith(b"FPXR"): - entry["kind"] = "flashpix" - entry["base64"] = _b64(body) - elif name == "APP11": - entry["kind"] = "c2pa_or_jumbf" - # the parsed manifest is in c2pa_store, but the raw JUMBF - # also carries assertion thumbnails the JSON may omit - entry["base64"] = _b64(body) - elif body.startswith(b"Exif\x00\x00"): - entry["kind"] = "exif" - entry["base64"] = _b64(body) - elif body.startswith(b"Photoshop 3.0\x00"): - entry["kind"] = "iptc_iim" - entry["base64"] = _b64(body) - else: - entry["base64"] = _b64(body) - result["segments"].append(entry) - pos += 2 + length - except Exception as exc: - result["error"] = _safe_str(exc) - return result - - -def read_pil_info(path: Path) -> tuple[dict[str, Any], dict[str, Any], bytes | None]: - """One Image.open serving all PIL-derived data: container basics, - img.info passthrough (XMP, comments), the IPTC-IIM dataset, and the - raw EXIF blob (for the caller's piexif parse on PNG/WebP/HEIC).""" - out: dict[str, Any] = {} - iptc: dict[str, Any] = {} - exif_blob: bytes | None = None - try: - with Image.open(path) as img: - out["format"] = img.format - out["mode"] = img.mode - out["width"], out["height"] = img.size - out["n_frames"] = getattr(img, "n_frames", 1) - dpi = img.info.get("dpi") - if dpi: - out["dpi"] = [round(float(d), 2) for d in dpi] - icc = img.info.get("icc_profile") - if icc: - out["icc_profile"] = { - "length": len(icc), - # header: profile class, color space, PCS (bytes 12-24) - "header_hex": icc[12:24].hex() if len(icc) >= 24 else "", - "base64": _b64(icc), - } - blob = img.info.get("exif") - if isinstance(blob, bytes): - exif_blob = blob - try: - info = getiptcinfo(img) - except Exception: - info = None - if info: - iptc = {f"{k[0]}:{k[1]}": _decode_exif_value(v) for k, v in info.items()} - for key, value in img.info.items(): - if key in ("icc_profile", "exif", "dpi"): - continue - if isinstance(value, bytes): - try: - out[f"info:{key}"] = value.decode("utf-8", "strict") - except (UnicodeDecodeError, ValueError): - out[f"info:{key}"] = f"base64:{_b64(value)}" - else: - out[f"info:{key}"] = _safe_str(value) - except Exception as exc: - out["error"] = _safe_str(exc) - return out, iptc, exif_blob - - -def read_c2pa_store(path: Path) -> dict[str, Any]: - """Full C2PA manifest store JSON via the official c2pa-python Reader.""" - try: - from c2pa import Reader - - with Reader(str(path)) as reader: - return json.loads(reader.json()) - except Exception as exc: - return {"error": _safe_str(exc)} - - -def sniff_format(head: bytes) -> str: - if head.startswith(b"\x89PNG"): - return "png" - if head.startswith(b"\xff\xd8"): - return "jpeg" - if head.startswith(b"RIFF") and head[8:12] == b"WEBP": - return "webp" - if head[:6] in (b"GIF87a", b"GIF89a"): - return "gif" - if head.startswith(b"BM"): - return "bmp" - if head.startswith((b"II*\x00", b"MM\x00*")): - return "tiff" - if head[4:8] == b"ftyp": - return f"isobmff:{head[8:12].decode('latin-1', 'replace')}" - return f"unknown:{head[:16].hex()}" - - -# --- forensic helpers: signals that a file is not an untouched original --- - - -def _jpeg_forensics_bytes(data: bytes) -> dict[str, Any]: - """Structure-level JPEG forensics: DQT tables (encoder fingerprint), - SOF type (baseline/progressive) + chroma subsampling, DHT Huffman - tables (custom = optimizing encoder), per-scan spectral selection - (progressive scan script), JFIF/Adobe app markers, COM, DRI.""" - out: dict[str, Any] = {} - try: - if not data.startswith(b"\xff\xd8"): - return out - pos = 2 - scans: list[dict[str, int]] = [] - dqt: dict[str, list[int]] = {} - dht: list[str] = [] - comments: list[str] = [] - while pos + 4 <= len(data): - if data[pos] != 0xFF: - break - marker = data[pos + 1] - if marker in (0xD8, 0x01) or 0xD0 <= marker <= 0xD7: - pos += 2 - continue - if marker == 0xD9: - break - length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] - body = data[pos + 4 : pos + 2 + length] - if marker == 0xDB: # DQT - off = 0 - while off < len(body): - tid = body[off] & 0x0F - prec = body[off] >> 4 - n = 128 if prec else 64 - vals = list(body[off + 1 : off + 1 + n]) - if prec: # 16-bit entries - vals = [struct.unpack(">H", bytes(vals[i : i + 2]))[0] for i in range(0, len(vals) - 1, 2)] - dqt[str(tid)] = vals[:64] - off += 1 + n - elif marker == 0xC4: # DHT: custom tables mean an optimizing encoder - dht.append(body.hex()) - elif marker == 0xDD and len(body) >= 2: # DRI - out["restart_interval"] = struct.unpack(">H", body[:2])[0] - elif marker == 0xE0 and body.startswith(b"JFIF\x00") and len(body) >= 12: - out["jfif"] = { - "version": f"{body[5]}.{body[6]}", - "density_units": body[7], - "x_density": struct.unpack(">H", body[8:10])[0], - "y_density": struct.unpack(">H", body[10:12])[0], - } - elif marker == 0xEE and body.startswith(b"Adobe") and len(body) >= 12: - out["adobe_transform"] = body[11] - elif marker in (0xC0, 0xC1, 0xC2) and len(body) >= 6: - out["progressive"] = marker == 0xC2 - out["precision_bits"] = body[0] - out["sof_height"] = struct.unpack(">H", body[1:3])[0] - out["sof_width"] = struct.unpack(">H", body[3:5])[0] - comps = [] - for i in range(body[5]): - c = body[6 + i * 3 : 9 + i * 3] - if len(c) == 3: - comps.append({"h": c[1] >> 4, "v": c[1] & 0x0F, "tq": c[2]}) - if len(comps) >= 3: - lum = comps[0] - subs = {1: "4:4:4", 2: "4:2:2"}.get(lum["h"] * lum["v"]) - out["subsampling"] = subs or f"{lum['h']}x{lum['v']}" - elif marker == 0xFE: # COM - comments.append(body.decode("utf-8", "replace")[:2000]) - elif marker == 0xDA: - # SOS spectral selection: the progressive scan script - # differs across libjpeg / mozjpeg / Photoshop - if len(body) >= 3: - ns = body[0] - tail = body[1 + ns * 2 :] - if len(tail) >= 3: - scans.append({"ss": tail[0], "se": tail[1], "ah": tail[2] >> 4, "al": tail[2] & 0x0F}) - # skip entropy-coded data to the next marker - end = data.find(b"\xff\xd9", pos) - nxt = data.find(b"\xff", pos + 2) - while nxt != -1 and nxt + 1 < len(data) and data[nxt + 1] == 0x00: - nxt = data.find(b"\xff", nxt + 2) - if nxt == -1 or (end != -1 and nxt >= end): - break - pos = nxt - continue - pos += 2 + length - if dqt: - out["quant_tables"] = dqt - if dht: - out["huffman_tables_hex"] = dht - if comments: - out["comments"] = comments - if scans: - out["scan_count"] = len(scans) - out["scan_script"] = scans - except Exception as exc: - out["error"] = _safe_str(exc) - return out - - -def read_webp_chunks(data: bytes) -> list[dict[str, Any]]: - """WebP RIFF chunk inventory (VP8X/VP8/VP8L/EXIF/XMP/ICCP/ANIM...).""" - chunks: list[dict[str, Any]] = [] - try: - pos = 12 - while pos + 8 <= len(data): - ctype = data[pos : pos + 4].decode("latin-1") - length = struct.unpack(" str: - import hashlib - - return hashlib.sha256(data).hexdigest() - - -def xattr_where_from(path: Path) -> list[str]: - """macOS download-source URLs (kMDItemWhereFroms), empty elsewhere.""" - import os - import plistlib - - try: - raw = os.getxattr(path, "com.apple.metadata:kMDItemWhereFroms") - value = plistlib.loads(raw) - return [str(v) for v in value] if isinstance(value, list) else [str(value)] - except (AttributeError, OSError, ValueError): - return [] - - -def xattr_quarantine(path: Path) -> str | None: - """macOS quarantine string: flags; timestamp; downloading agent (Safari, - Telegram, Chrome...). Presence alone means 'came from the internet'.""" - import os - - try: - return os.getxattr(path, "com.apple.quarantine").decode("utf-8", "replace")[:500] - except (AttributeError, OSError): - return None - - -def read_isobmff_inventory(data: bytes) -> dict[str, Any]: - """HEIC/AVIF/MOV box inventory: top-level boxes plus the meta item - types (Exif, mime=XMP, auxl depth/gain-map, aae Apple-edits plist, - irot derived images). Strong phone-provenance signal.""" - out: dict[str, Any] = {} - try: - - def boxes(start: int, end: int) -> list[tuple[str, int, int]]: - result = [] - pos = start - while pos + 8 <= end: - size, btype = struct.unpack(">I4s", data[pos : pos + 8]) - t = btype.decode("latin-1") - header = 8 - if size == 1: - size = struct.unpack(">Q", data[pos + 8 : pos + 16])[0] - header = 16 - elif size == 0: - size = end - pos - if size < header or pos + size > end: - break - result.append((t, pos + header, pos + size)) - pos += size - return result - - top = boxes(0, len(data)) - out["boxes"] = [t for t, _, _ in top] - for t, s, e in top: - if t == "moov": - for ct, cs, ce in boxes(s, e): - if ct == "mvhd" and ce - cs >= 24: - # full box + creation/modification times (1904 epoch) - version = data[cs] - base = cs + 4 - creation = struct.unpack(">I", data[base : base + 4])[0] if version == 0 else None - if creation: - out["mvhd_creation_time"] = creation - 2082844800 - elif t == "meta": - # full box: 4 bytes version/flags, then child boxes - for ct, cs, ce in boxes(s + 4, e): - if ct == "iinf": - # full box + entry count, then infe entries - count = struct.unpack(">H", data[cs + 4 : cs + 6])[0] - out["meta_item_count"] = count - item_types = [] - for it, is_, ie in boxes(cs + 6, ce): - if it == "infe" and ie - is_ >= 8: - # infe full box: version(1)+flags(3), then - # v2: item_ID(2)+protection(2)+item_type(4) - # v3: item_ID(4)+protection(2)+item_type(4) - version = data[is_] - off = is_ + 4 + (4 if version == 3 else 2) + 2 - if off + 4 <= ie: - item_types.append(data[off : off + 4].decode("latin-1", "replace")) - if item_types: - out["meta_item_types"] = sorted(set(item_types)) - elif ct == "iprp": - out["has_iprp"] = True - for pt, ps, pe in boxes(cs, ce): - if pt == "ipco": - props = [t for t, _, _ in boxes(ps, pe)] - out["ipco_properties"] = props - # auxC holds the auxiliary image type URN - for qt, qs, qe in boxes(ps, pe): - if qt == "auxC": - out["auxc_types"] = ( - data[qs + 4 : qe].split(b"\x00")[0].decode("latin-1", "replace") - ) - elif ct == "iref": - out["has_iref"] = True - # QuickTime metadata keys (©mak/©mod/©swr) for the MOV side of - # Live Photos: tolerant printable-string grab after each atom - import re - - qt: dict[str, str] = {} - for atom, key in ((b"\xa9mak", "make"), (b"\xa9mod", "model"), (b"\xa9swr", "software")): - idx = data.find(atom) - if idx != -1: - m = re.search(rb"[ -~]{4,80}", data[idx + 4 : idx + 200]) - if m: - qt[key] = m.group(0).decode("ascii", "replace") - if qt: - out["quicktime"] = qt - except Exception as exc: - out["error"] = _safe_str(exc) - return out - - -def apple_live_photo_id(head: bytes) -> str | None: - """Apple Live Photo content identifier (links the still to its MOV). - - The UUID sits in the Apple MakerNote (tag 17) of the still and in the - MOV metadata; a raw head scan finds it in either container.""" - import re - - # the UUID string sits next to "content.identifier" in the MOV, but in - # the STILL it is a bare UUID inside the Apple MakerNote (whose header - # is "Apple iOS"), so gate on either marker - if b"content.identifier" not in head and b"com.apple.quicktime" not in head and b"Apple iOS" not in head: - return None - m = re.search( - rb"[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}", - head, - ) - return m.group(0).decode("ascii") if m else None - - -_MAX_FULL_READ = 256 << 20 # files bigger than this are scanned head-only -_HEAD_READ = 4 << 20 - - -def _sha256_stream(path: Path) -> str: - import hashlib - - h = hashlib.sha256() - with open(path, "rb") as f: - for block in iter(lambda: f.read(1 << 20), b""): - h.update(block) - return h.hexdigest() - - -def scan_file(path: Path, *, pixel_mode: str | None = None) -> dict[str, Any]: - stat = path.stat() - oversized = stat.st_size > _MAX_FULL_READ - if oversized: - data = None - with open(path, "rb") as f: - head = f.read(_HEAD_READ) - else: - data = path.read_bytes() - head = data - record: dict[str, Any] = { - "file": str(path), - "name": path.name, - "extension": path.suffix.lower(), - "size_bytes": stat.st_size, - "mtime": stat.st_mtime, - "birthtime": getattr(stat, "st_birthtime", None), - "sha256": _sha256_stream(path) if data is None else sha256_of(data), - "content_format": sniff_format(head), - } - if oversized: - # too big to hold in memory: path-based readers (PIL, piexif, - # c2pa) still run; byte-level walkers are skipped - record["oversized"] = {"head_scanned_bytes": len(head)} - where_from = xattr_where_from(path) - if where_from: - record["download_source_urls"] = where_from - quarantine = xattr_quarantine(path) - if quarantine: - record["quarantine"] = quarantine - live_photo_id = apple_live_photo_id(head[: 2 << 20]) - if live_photo_id: - record["live_photo_content_id"] = live_photo_id - record["pil"], record["iptc"], exif_blob = read_pil_info(path) - record["exif"], thumbnail = read_full_exif(path, exif_blob, data) - record["c2pa_store"] = read_c2pa_store(path) - if data is not None: - fmt = record["content_format"] - if fmt == "png": - record["png_chunks"], post_iend = read_png_chunks(data) - if post_iend: - record["png_post_iend_bytes"] = post_iend - elif fmt == "jpeg": - record["jpeg"] = read_jpeg_segments(data) - record["jpeg_forensics"] = _jpeg_forensics_bytes(data) - elif fmt == "webp": - record["webp_chunks"] = read_webp_chunks(data) - elif fmt.startswith("isobmff"): - record["isobmff"] = read_isobmff_inventory(data) - if thumbnail: - record["has_exif_thumbnail"] = True - # the embedded thumbnail is its own JPEG; after an edit its encoder - # forensics commonly MISMATCH the main image (classic tamper tell) - thumb_forensics = _jpeg_forensics_bytes(thumbnail) - thumb_forensics["base64"] = _b64(thumbnail) - record["exif_thumbnail_forensics"] = thumb_forensics - if pixel_mode is not None: - if oversized: - # symmetric with the metadata oversized marker: distinguish - # "skipped because oversized" from "decode failed" downstream - record["pixel"] = {"skipped": "oversized"} - else: - record.update(scan_pixels_of(path, full=pixel_mode == "full")) - return record - - -def scan_pixels_of(path: Path, *, full: bool) -> dict[str, Any]: - """The optional pixel layer for one file (see the pixel-layer section - below). Returns {"pixel_error": ...} when numpy is unavailable.""" - if np is None: - return {"pixel_error": "numpy not installed"} - t0 = time.perf_counter() - timing: dict[str, float] = {} - record: dict[str, Any] = {} - - t = time.perf_counter() - gray, rgb, info = read_gray(path) - record["pixel"] = info - timing["decode"] = time.perf_counter() - t - if gray is not None: - # maps shared between the scalar features and the --pixels-full - # artifacts, computed once (the JPEG re-save and the sliding-window - # conv are the two most expensive features in the scan) - t = time.perf_counter() - residual = noise_residual_map(gray) - timing["noise"] = time.perf_counter() - t - if residual is not None: - record["noise"] = noise_features(residual) - - t = time.perf_counter() - spectrum = fft_decompose(gray) - timing["fft"] = time.perf_counter() - t - mag = phase = None - if spectrum is not None: - mag, phase = spectrum - record["fft"] = fft_features(mag) - - t = time.perf_counter() - ela = ela_map(rgb) - timing["ela"] = time.perf_counter() - t - if ela is not None: - record["ela"] = ela_features(ela) - - for name, fn, arg in ( - ("dct", dct_features, gray), - ("gradient", gradient_features, gray), - ("color", color_features, rgb), - ): - t = time.perf_counter() - try: - result = fn(arg) - except Exception as exc: - result = {"error": _safe_str(exc)} - timing[name] = time.perf_counter() - t - if result: - record[name] = result - if full: - t = time.perf_counter() - try: - record["full"] = full_artifacts(gray, rgb, ela=ela, residual=residual, phase=phase) - except Exception as exc: - record["full"] = {"error": _safe_str(exc)} - timing["full_artifacts"] = time.perf_counter() - t - timing["total"] = time.perf_counter() - t0 - record["timing_ms"] = {k: round(v * 1000, 1) for k, v in timing.items()} - return record - - -def summary_row(record: dict[str, Any]) -> dict[str, Any]: - pil = record.get("pil", {}) - return { - "file": record.get("file"), - "name": record.get("name"), - "size_bytes": record.get("size_bytes"), - "sha256": record.get("sha256"), - "content_format": record.get("content_format"), - "width": pil.get("width"), - "height": pil.get("height"), - "dpi": json.dumps(pil.get("dpi")), - } - - -def main() -> None: - flags, positional = set(), [] - for arg in sys.argv[1:]: - (flags.add if arg.startswith("--") else positional.append)(arg) - if len(positional) != 2 or flags - {"--pixels", "--pixels-full"}: - print(__doc__) - sys.exit(2) - pixel_mode = "full" if "--pixels-full" in flags else ("basic" if "--pixels" in flags else None) - if pixel_mode is not None and np is None: - print("pixel modes require numpy: pip install numpy") - sys.exit(2) - root, prefix = Path(positional[0]), positional[1] - files = sorted(str(p) for p in root.rglob("*") if p.is_file() and p.suffix.lower() in SUPPORTED) - # resume: skip files already present in the output of a previous - # (interrupted) run with the same prefix - gz = prefix.endswith(".gz") - jsonl_name = prefix if gz else f"{prefix}.jsonl" - csv_name = f"{prefix[:-3] if gz else prefix}.csv" - import gzip - - jsonl_path = Path(jsonl_name) - done: set[str] = set() - if jsonl_path.exists(): - opener = gzip.open if gz else open - with opener(jsonl_path, "rt") as existing: # type: ignore[arg-type] - for line in existing: - with contextlib.suppress(Exception): - done.add(json.loads(line).get("file", "")) - files = [f for f in files if f not in done] - print(f"scanning {len(files)} files under {root} ({len(done)} already done)") - n_done = 0 - text_opener = gzip.open if gz else open - csv_exists = Path(csv_name).exists() - with ( - text_opener(jsonl_path, "at") as jsonl, # type: ignore[arg-type] - open(csv_name, "a", newline="") as csvf, - ): - writer = csv.DictWriter(csvf, fieldnames=list(summary_row({}))) - if not csv_exists: - writer.writeheader() - for path_str in files: - try: - record = scan_file(Path(path_str), pixel_mode=pixel_mode) - except Exception as exc: # one corrupt file must not kill the scan - record = {"file": path_str, "error": _safe_str(exc)} - jsonl.write(json.dumps(record, default=str) + "\n") - writer.writerow(summary_row(record)) - n_done += 1 - if n_done % 500 == 0: - print(f" {n_done}/{len(files)}", flush=True) - print(f"done: {jsonl_name}, {csv_name}") - - -# --- optional pixel layer (--pixels / --pixels-full), requires numpy --- - -try: - import numpy as np - from numpy.lib.stride_tricks import sliding_window_view -except ImportError: - np = None # type: ignore[assignment] - -_MAX_SIDE = 2048 # downscale before analysis; stats are scale-robust -_DCT_BINS = np.linspace(-20.5, 20.5, 22) if np is not None else None -_AC_POSITIONS = [(0, 1), (1, 0), (1, 1), (0, 2), (2, 0), (2, 1), (1, 2), (0, 3)] -_FFT_BANDS = 8 -_BAYER_OFFSETS = [(1, 1), (1, -1)] # CFA diagonal periodicity candidates - - -def _arr_b64(arr: np.ndarray) -> dict[str, Any]: - """Compact array payload for the --full spatial maps.""" - return { - "shape": list(arr.shape), - "dtype": str(arr.dtype), - "base64": _b64(arr.tobytes()), - } - - -def _coarse(arr: np.ndarray, side: int = 64) -> np.ndarray: - """Downscale a 2D map to at most `side` on the long edge.""" - h, w = arr.shape - if max(h, w) <= side: - return arr - img = Image.fromarray(arr.astype(np.float32), mode="F") - img.thumbnail((side, side), Image.BILINEAR) - return np.asarray(img) - - -def phash(gray: np.ndarray) -> str: - """64-bit DCT perceptual hash (invertible to a rough layout; --full only).""" - img = Image.fromarray(gray.astype(np.float32), mode="F").resize((32, 32), Image.LANCZOS) - small = np.asarray(img) - m = _dct_matrix(32) - coeff = m @ small @ m.T - low = coeff[:8, :8].ravel()[1:] # drop DC - bits = low > np.median(low) - return f"{int(''.join('1' if b else '0' for b in bits), 2):016x}" - - -def full_artifacts( - gray: np.ndarray, - rgb: np.ndarray, - *, - ela: np.ndarray | None, - residual: np.ndarray | None, - phase: np.ndarray | None, -) -> dict[str, Any]: - """The privacy-lifting set: phash, thumbnail, ELA map, noise residual, - FFT phase. Maps are computed once by the caller and shared with the - scalar feature paths.""" - out: dict[str, Any] = {} - out["phash"] = phash(gray) - img = Image.fromarray(rgb.astype(np.uint8)) - img.thumbnail((128, 128), Image.LANCZOS) - buf = io.BytesIO() - img.save(buf, "JPEG", quality=70) - out["thumbnail_jpeg_b64"] = _b64(buf.getvalue()) - - if ela is not None: - out["ela_map"] = _arr_b64(_coarse(ela)) - if residual is not None: - clipped = np.clip(residual / 4.0, -1, 1) - out["noise_residual"] = _arr_b64(_coarse((clipped * 127).astype(np.int8))) - if phase is not None: - out["fft_phase"] = _arr_b64(_coarse(phase.astype(np.float32), 32)) - return out - - -def read_gray(path: Path) -> tuple[np.ndarray | None, np.ndarray | None, dict[str, Any]]: - """Decode to float32 grayscale (and RGB for color stats), downscaled.""" - try: - with Image.open(path) as img: - info: dict[str, Any] = {"width": img.width, "height": img.height} - if max(img.size) > _MAX_SIDE: - img.thumbnail((_MAX_SIDE, _MAX_SIDE), Image.LANCZOS) - rgb = np.asarray(img.convert("RGB"), dtype=np.float32) - gray = np.asarray(img.convert("L"), dtype=np.float32) - return gray, rgb, info - except Exception as exc: - return None, None, {"error": _safe_str(exc)} - - -def _dct_matrix(n: int = 8) -> np.ndarray: - """Orthonormal n x n DCT-II basis: M[i, j] = cos(pi (2j + 1) i / 2n).""" - i = np.arange(n)[:, None] - j = np.arange(n)[None, :] - m = np.cos(np.pi * (2 * j + 1) * i / (2 * n)) - m[0, :] *= 1 / np.sqrt(2) - return m * np.sqrt(2 / n) - - -_DCT_M = _dct_matrix() if np is not None else None - - -def dct_features(gray: np.ndarray) -> dict[str, Any]: - """AC coefficient histograms over 8x8 block DCT + Benford deviation.""" - h, w = gray.shape - h8, w8 = h // 8 * 8, w // 8 * 8 - if h8 < 8 or w8 < 8: - return {} - blocks = gray[:h8, :w8].reshape(h8 // 8, 8, w8 // 8, 8).swapaxes(1, 2) - coeff = np.einsum("ij,abjk,lk->abil", _DCT_M, blocks, _DCT_M) - hists = [] - lead_vals: list[np.ndarray] = [] - for dy, dx in _AC_POSITIONS: - vals = coeff[:, :, dy, dx].ravel() - hists.append(np.histogram(vals, bins=_DCT_BINS)[0].tolist()) - lead_vals.append(np.abs(vals)) - out: dict[str, Any] = {"dct_ac_hist": hists} - flat = np.abs(np.concatenate(lead_vals)) - flat = flat[flat >= 1] - if flat.size > 100: - leading = (flat / 10 ** np.floor(np.log10(flat))).astype(int) - leading = leading[(leading >= 1) & (leading <= 9)] - if leading.size > 100: - obs = np.bincount(leading, minlength=10)[1:10] / leading.size - ben = np.log10(1 + 1 / np.arange(1, 10)) - out["benford_mad"] = float(np.abs(obs - ben).mean()) - return out - - -def noise_residual_map(gray: np.ndarray) -> np.ndarray | None: - """High-pass residual. The map itself is spatial (shows edges), so it - is only stored under --full; the default mode keeps scalar stats.""" - if gray.shape[0] < 3 or gray.shape[1] < 3: - return None - k = np.array([[-1.0, -1.0, -1.0], [-1.0, 8.0, -1.0], [-1.0, -1.0, -1.0]]) - h, w = gray.shape - # k is float64, so the residual is float64 like the unchunked form - out = np.empty((h - 2, w - 2), dtype=np.float64) - # row-chunked: the (win * k) temp is ~150 MB at 2048px if materialized - # whole; per-element 9-tap sums are computed in the same order, so the - # result is bit-identical to the unchunked form - for y0 in range(0, h - 2, 256): - y1 = min(y0 + 256, h - 2) - win = sliding_window_view(gray[y0 : y1 + 2], (3, 3)) - out[y0:y1] = (win * k).sum(axis=(-1, -2)) - return out - - -def noise_features(residual: np.ndarray) -> dict[str, Any]: - """High-pass residual std/kurtosis; the residual map is NOT stored.""" - flat = residual.ravel() - std = float(flat.std()) - if std < 1e-9: - return {"noise_std": 0.0, "noise_kurtosis": 0.0} - z = (flat - flat.mean()) / std - return {"noise_std": std, "noise_kurtosis": float((z**4).mean() - 3.0)} - - -def fft_decompose(gray: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None: - """Log-magnitude (fftshifted) and phase of the image spectrum.""" - if min(gray.shape) < 32: - return None - spectrum = np.fft.fftshift(np.fft.fft2(gray - gray.mean())) - return np.log1p(np.abs(spectrum)), np.angle(spectrum) - - -def fft_features(mag: np.ndarray) -> dict[str, Any]: - """Radial magnitude band energies (no phase) + CFA periodicity peaks.""" - h, w = mag.shape - cy, cx = h // 2, w // 2 - # 1D broadcast instead of an mgrid: saves ~160 MB of int64 temporaries - # at 2048px; the squares are exact in float64 (values < 2^53), so band - # means are identical to the mgrid form - r2y = (np.arange(h, dtype=np.float64) - cy) ** 2 - r2x = (np.arange(w, dtype=np.float64) - cx) ** 2 - r = np.sqrt(r2y[:, None] + r2x[None, :]) - r_max = r.max() - bands = [] - for i in range(_FFT_BANDS): - mask = (r >= r_max * i / _FFT_BANDS) & (r < r_max * (i + 1) / _FFT_BANDS) - bands.append(float(mag[mask].mean()) if mask.any() else 0.0) - # Bayer CFA shows as symmetric peaks at half the Nyquist on diagonals - peaks = [] - for dy, dx in _BAYER_OFFSETS: - y, x = cy + dy * (h // 4), cx + dx * (w // 4) - neighborhood = mag[y - 2 : y + 3, x - 2 : x + 3] - peaks.append(float(neighborhood.max() - mag.mean())) - return {"fft_band_energy": bands, "cfa_peaks": peaks, "cfa_peak": max(peaks)} - - -def ela_map(rgb: np.ndarray) -> np.ndarray | None: - """Absolute per-pixel error after a q90 JPEG re-save.""" - try: - img = Image.fromarray(rgb.astype(np.uint8)) - buf = io.BytesIO() - img.save(buf, "JPEG", quality=90) - buf.seek(0) - resaved = np.asarray(Image.open(buf).convert("RGB"), dtype=np.float32) - except Exception: - return None - if resaved.shape != rgb.shape: - return None - return np.abs(rgb - resaved).mean(axis=-1) - - -def ela_features(err: np.ndarray) -> dict[str, Any]: - """Error-level stats after a q90 JPEG re-save; global only, no map.""" - return {"ela_mean": float(err.mean()), "ela_p95": float(np.percentile(err, 95))} - - -def gradient_features(gray: np.ndarray) -> dict[str, Any]: - gy, gx = np.gradient(gray) - mag = np.sqrt(gx**2 + gy**2) - hist = np.histogram(mag, bins=10, range=(0, 255))[0].tolist() - lap = np.gradient(gy, axis=0) + np.gradient(gx, axis=1) - return {"gradient_hist": hist, "laplacian_var": float(lap.var())} - - -def color_features(rgb: np.ndarray) -> dict[str, Any]: - small = rgb[::4, ::4] # decimate; histogram is position-blind anyway - bins = (small / 256 * 4).astype(int).clip(0, 3) - idx = bins[..., 0] * 16 + bins[..., 1] * 4 + bins[..., 2] - hist = np.bincount(idx.ravel(), minlength=64).tolist() - mx = small.max(axis=-1) - mn = small.min(axis=-1) - sat = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-6), 0) - return { - "color_hist_4x4x4": hist, - "saturation_mean": float(sat.mean()), - "value_mean": float(mx.mean() / 255), - } - - -if __name__ == "__main__": - main() diff --git a/src/remove_ai_watermarks/pixel_evidence.py b/src/remove_ai_watermarks/pixel_evidence.py new file mode 100644 index 0000000..077dfdc --- /dev/null +++ b/src/remove_ai_watermarks/pixel_evidence.py @@ -0,0 +1,400 @@ +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false +"""Experimental: the complete pixel-forensics layer for one image. + +STATUS + +Experimental and unused. Nothing in this package reads it -- not the provenance +verdict, not removal, not the CLI. It is here because the research scanner that +produced these measurements is gone, and the capability was worth keeping: whatever +asks for pixel forensics next starts from a tested implementation instead of +rebuilding one. Treat the shape as unstable until a caller exists. + +WHAT IS MEASURED + +One decode, then six families of scale-robust statistics over it: + +* ``dct`` -- AC coefficient histograms over the 8x8 block DCT, plus the deviation of + leading digits from Benford's law. +* ``fft`` -- radial band energies of the log-magnitude spectrum, plus the + colour-filter-array periodicity peaks a demosaiced camera capture leaves. +* ``noise`` -- standard deviation and kurtosis of a high-pass residual. +* ``ela`` -- error level after a quality-90 JPEG re-save. +* ``gradient`` -- gradient-magnitude histogram and Laplacian variance. +* ``color`` -- 4x4x4 RGB histogram, mean saturation, mean value. + +and, in ``artifacts``, the spatial layer those statistics are computed from: a +64-bit perceptual hash, a 128px JPEG thumbnail, and coarse ELA, noise-residual and +FFT-phase maps. + +THE ARTIFACTS ARE NOT AGGREGATES + +Everything above ``artifacts`` is a scalar or a fixed-length histogram, and an image +cannot be reconstructed from those. ``artifacts`` is different in kind: a thumbnail +is a picture, a perceptual hash identifies one, and the coarse maps carry layout. +Collecting them makes a record that identifies the source image, so a caller storing +or forwarding them is handling image content, not statistics about it. That is why +they are a separate field and not merged into the families. + +REQUIREMENTS + +Needs the ``pixels`` extra (numpy). Guard a call with :func:`is_available` when the +caller must not hard-depend on it. +""" + +from __future__ import annotations + +import base64 +import io +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from pathlib import Path + +logger = logging.getLogger(__name__) + +# Analysis resolution. Every statistic here is scale-robust, and a 2048px cap keeps +# the FFT and the sliding-window residual bounded on a 100 MP input. +MAX_SIDE = 2048 +# The eight lowest-frequency AC positions of the 8x8 block DCT, zig-zag order. +AC_POSITIONS = ((0, 1), (1, 0), (1, 1), (0, 2), (2, 0), (2, 1), (1, 2), (0, 3)) +FFT_BANDS = 8 +# A Bayer CFA shows as symmetric peaks at half the Nyquist on the diagonals. +BAYER_OFFSETS = ((1, 1), (1, -1)) +INSTALL_HINT = "install the pixel extra: uv add 'remove-ai-watermarks[pixels]'" + + +@dataclass(frozen=True) +class PixelEvidence: + """Pixel statistics for one image, and the spatial artifacts behind them. + + ``decode`` carries the source dimensions, or ``{"error": ...}`` when the image + could not be decoded -- in which case every other field is empty. A family is also + empty when the image is too small for it (the block DCT needs 8x8, the FFT 32x32, + the residual 3x3), so a caller must treat every field as optional rather than + assume a fixed feature width. + """ + + path: Path + decode: dict[str, Any] + dct: dict[str, Any] = field(default_factory=dict[str, Any]) + fft: dict[str, Any] = field(default_factory=dict[str, Any]) + noise: dict[str, Any] = field(default_factory=dict[str, Any]) + ela: dict[str, Any] = field(default_factory=dict[str, Any]) + gradient: dict[str, Any] = field(default_factory=dict[str, Any]) + color: dict[str, Any] = field(default_factory=dict[str, Any]) + # Identifies the source image; see the module note. Empty unless asked for. + artifacts: dict[str, Any] = field(default_factory=dict[str, Any]) + + @property + def decoded(self) -> bool: + """False when the source could not be decoded at all.""" + return "error" not in self.decode + + +def is_available() -> bool: + """True when the optional pixel dependencies are installed.""" + from remove_ai_watermarks.optional_deps import module_available + + return module_available("numpy") + + +def _numpy() -> Any: + from remove_ai_watermarks.optional_deps import module_available + + if not module_available("numpy"): + raise RuntimeError(f"Pixel evidence needs numpy -- {INSTALL_HINT}") + import numpy as np + + return np + + +def _dct_matrix(np: Any, n: int = 8) -> Any: + """Orthonormal n x n DCT-II basis: M[i, j] = cos(pi (2j + 1) i / 2n).""" + i = np.arange(n)[:, None] + j = np.arange(n)[None, :] + m = np.cos(np.pi * (2 * j + 1) * i / (2 * n)) + m[0, :] *= 1 / np.sqrt(2) + return m * np.sqrt(2 / n) + + +def read_gray(image_path: Path) -> tuple[Any, Any, dict[str, Any]]: + """Decode to float32 grayscale (and RGB for colour stats), downscaled. + + Pillow, not cv2, and the source dimensions are recorded BEFORE the downscale. + """ + np = _numpy() + from PIL import Image + + try: + with Image.open(image_path) as img: + info: dict[str, Any] = {"width": img.width, "height": img.height} + if max(img.size) > MAX_SIDE: + img.thumbnail((MAX_SIDE, MAX_SIDE), Image.Resampling.LANCZOS) + rgb = np.asarray(img.convert("RGB"), dtype=np.float32) + gray = np.asarray(img.convert("L"), dtype=np.float32) + except Exception as exc: + logger.debug("pixel decode failed for %s: %s", image_path, exc) + return None, None, {"error": f"{type(exc).__name__}: {exc}"} + return gray, rgb, info + + +def dct_features(gray: Any) -> dict[str, Any]: + """AC coefficient histograms over the 8x8 block DCT + Benford deviation.""" + np = _numpy() + height, width = gray.shape + h8, w8 = height // 8 * 8, width // 8 * 8 + if h8 < 8 or w8 < 8: + return {} + basis = _dct_matrix(np) + bins = np.linspace(-20.5, 20.5, 22) + blocks = gray[:h8, :w8].reshape(h8 // 8, 8, w8 // 8, 8).swapaxes(1, 2) + coeff = np.einsum("ij,abjk,lk->abil", basis, blocks, basis) + hists = [] + lead_vals: list[Any] = [] + for dy, dx in AC_POSITIONS: + values = coeff[:, :, dy, dx].ravel() + hists.append(np.histogram(values, bins=bins)[0].tolist()) + lead_vals.append(np.abs(values)) + out: dict[str, Any] = {"dct_ac_hist": hists} + flat = np.abs(np.concatenate(lead_vals)) + flat = flat[flat >= 1] + if flat.size > 100: + leading = (flat / 10 ** np.floor(np.log10(flat))).astype(int) + leading = leading[(leading >= 1) & (leading <= 9)] + if leading.size > 100: + observed = np.bincount(leading, minlength=10)[1:10] / leading.size + benford = np.log10(1 + 1 / np.arange(1, 10)) + out["benford_mad"] = float(np.abs(observed - benford).mean()) + return out + + +def noise_residual_map(gray: Any) -> Any: + """High-pass residual, the map the noise statistics are computed from.""" + np = _numpy() + from numpy.lib.stride_tricks import sliding_window_view + + if gray.shape[0] < 3 or gray.shape[1] < 3: + return None + kernel = np.array([[-1.0, -1.0, -1.0], [-1.0, 8.0, -1.0], [-1.0, -1.0, -1.0]]) + height, width = gray.shape + # kernel is float64, so the residual is float64 like the unchunked form + out = np.empty((height - 2, width - 2), dtype=np.float64) + # Row-chunked: the (window * kernel) temporary is ~150 MB at 2048px if + # materialized whole. Per-element 9-tap sums are computed in the same order, + # so the result is bit-identical to the unchunked form. + for y0 in range(0, height - 2, 256): + y1 = min(y0 + 256, height - 2) + window = sliding_window_view(gray[y0 : y1 + 2], (3, 3)) + out[y0:y1] = (window * kernel).sum(axis=(-1, -2)) + return out + + +def noise_features(residual: Any) -> dict[str, Any]: + """High-pass residual std and kurtosis.""" + flat = residual.ravel() + std = float(flat.std()) + if std < 1e-9: + return {"noise_std": 0.0, "noise_kurtosis": 0.0} + z = (flat - flat.mean()) / std + return {"noise_std": std, "noise_kurtosis": float((z**4).mean() - 3.0)} + + +def fft_decompose(gray: Any) -> tuple[Any, Any] | None: + """Log-magnitude (fftshifted) and phase of the image spectrum.""" + np = _numpy() + if min(gray.shape) < 32: + return None + spectrum = np.fft.fftshift(np.fft.fft2(gray - gray.mean())) + return np.log1p(np.abs(spectrum)), np.angle(spectrum) + + +def fft_features(mag: Any) -> dict[str, Any]: + """Radial magnitude band energies (no phase) + CFA periodicity peaks.""" + np = _numpy() + height, width = mag.shape + cy, cx = height // 2, width // 2 + # 1D broadcast instead of an mgrid: saves ~160 MB of int64 temporaries at + # 2048px. The squares are exact in float64 (values < 2^53), so band means + # are identical to the mgrid form. + r2y = (np.arange(height, dtype=np.float64) - cy) ** 2 + r2x = (np.arange(width, dtype=np.float64) - cx) ** 2 + radius = np.sqrt(r2y[:, None] + r2x[None, :]) + r_max = radius.max() + bands = [] + for index in range(FFT_BANDS): + mask = (radius >= r_max * index / FFT_BANDS) & (radius < r_max * (index + 1) / FFT_BANDS) + bands.append(float(mag[mask].mean()) if mask.any() else 0.0) + peaks = [] + for dy, dx in BAYER_OFFSETS: + y, x = cy + dy * (height // 4), cx + dx * (width // 4) + neighborhood = mag[y - 2 : y + 3, x - 2 : x + 3] + peaks.append(float(neighborhood.max() - mag.mean())) + return {"fft_band_energy": bands, "cfa_peaks": peaks, "cfa_peak": max(peaks)} + + +def ela_map(rgb: Any) -> Any: + """Absolute per-pixel error after a quality-90 JPEG re-save.""" + np = _numpy() + from PIL import Image + + try: + buffer = io.BytesIO() + Image.fromarray(rgb.astype(np.uint8)).save(buffer, "JPEG", quality=90) + buffer.seek(0) + resaved = np.asarray(Image.open(buffer).convert("RGB"), dtype=np.float32) + except Exception as exc: + logger.debug("ELA re-save failed: %s", exc) + return None + if resaved.shape != rgb.shape: + return None + return np.abs(rgb - resaved).mean(axis=-1) + + +def ela_features(err: Any) -> dict[str, Any]: + """Error-level stats after a quality-90 JPEG re-save.""" + np = _numpy() + return {"ela_mean": float(err.mean()), "ela_p95": float(np.percentile(err, 95))} + + +def gradient_features(gray: Any) -> dict[str, Any]: + np = _numpy() + gy, gx = np.gradient(gray) + mag = np.sqrt(gx**2 + gy**2) + hist = np.histogram(mag, bins=10, range=(0, 255))[0].tolist() + laplacian = np.gradient(gy, axis=0) + np.gradient(gx, axis=1) + return {"gradient_hist": hist, "laplacian_var": float(laplacian.var())} + + +def color_features(rgb: Any) -> dict[str, Any]: + np = _numpy() + small = rgb[::4, ::4] # decimate; the histogram is position-blind anyway + bins = (small / 256 * 4).astype(int).clip(0, 3) + index = bins[..., 0] * 16 + bins[..., 1] * 4 + bins[..., 2] + hist = np.bincount(index.ravel(), minlength=64).tolist() + mx = small.max(axis=-1) + mn = small.min(axis=-1) + saturation = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-6), 0) + return { + "color_hist_4x4x4": hist, + "saturation_mean": float(saturation.mean()), + "value_mean": float(mx.mean() / 255), + } + + +def perceptual_hash(gray: Any) -> str: + """64-bit DCT perceptual hash. Identifies an image; see the module note.""" + np = _numpy() + from PIL import Image + + small = np.asarray(Image.fromarray(gray.astype(np.float32), mode="F").resize((32, 32), Image.Resampling.LANCZOS)) + basis = _dct_matrix(np, 32) + coeff = basis @ small @ basis.T + low = coeff[:8, :8].ravel()[1:] # drop DC + bits = low > np.median(low) + return f"{int(''.join('1' if bit else '0' for bit in bits), 2):016x}" + + +def _coarse(np: Any, arr: Any, side: int = 64) -> Any: + """Downscale a 2D map to at most ``side`` on the long edge.""" + from PIL import Image + + height, width = arr.shape + if max(height, width) <= side: + return arr + img = Image.fromarray(arr.astype(np.float32), mode="F") + img.thumbnail((side, side), Image.Resampling.BILINEAR) + return np.asarray(img) + + +def _array_payload(arr: Any) -> dict[str, Any]: + return { + "shape": list(arr.shape), + "dtype": str(arr.dtype), + "base64": base64.b64encode(arr.tobytes()).decode("ascii"), + } + + +def spatial_artifacts(gray: Any, rgb: Any, *, ela: Any, residual: Any, phase: Any) -> dict[str, Any]: + """Perceptual hash, thumbnail, and coarse ELA / residual / phase maps. + + These identify the source image rather than describe it -- see the module note. + The maps are the ones the statistics were computed from, passed in rather than + recomputed. + """ + np = _numpy() + from PIL import Image + + out: dict[str, Any] = {"phash": perceptual_hash(gray)} + thumbnail = Image.fromarray(rgb.astype(np.uint8)) + thumbnail.thumbnail((128, 128), Image.Resampling.LANCZOS) + buffer = io.BytesIO() + thumbnail.save(buffer, "JPEG", quality=70) + out["thumbnail_jpeg_b64"] = base64.b64encode(buffer.getvalue()).decode("ascii") + + if ela is not None: + out["ela_map"] = _array_payload(_coarse(np, ela)) + if residual is not None: + clipped = np.clip(residual / 4.0, -1, 1) + out["noise_residual"] = _array_payload(_coarse(np, (clipped * 127).astype(np.int8))) + if phase is not None: + out["fft_phase"] = _array_payload(_coarse(np, phase.astype(np.float32), 32)) + return out + + +def extract_pixel_evidence(image_path: Path, *, artifacts: bool = False) -> PixelEvidence: + """Measure every pixel-statistic family for one image in a single decode. + + The image is decoded ONCE and the intermediate maps (high-pass residual, ELA + error, FFT magnitude and phase) are computed once and shared, because the + residual's sliding window and the ELA re-save are the two expensive steps and + each family would otherwise redo them. + + A family that fails or does not apply is left empty rather than raising: an + undecodable file, or one too small for the block DCT, still returns a + :class:`PixelEvidence` whose ``decoded`` / empty fields say so. Missing numpy is + the one hard error, since then nothing can be measured at all. + + Args: + image_path: Path to the image. Any container Pillow can open. + artifacts: Also return the spatial layer -- perceptual hash, thumbnail and + coarse maps. Off by default: those identify the source image, so asking + for them is a decision the caller makes explicitly. + + Returns: + A :class:`PixelEvidence`. + """ + gray, rgb, info = read_gray(image_path) + if gray is None or rgb is None: + return PixelEvidence(path=image_path, decode=info) + + residual = noise_residual_map(gray) + spectrum = fft_decompose(gray) + error = ela_map(rgb) + + families: dict[str, dict[str, Any]] = {} + for name, compute in ( + ("noise", lambda: noise_features(residual) if residual is not None else {}), + ("fft", lambda: fft_features(spectrum[0]) if spectrum is not None else {}), + ("ela", lambda: ela_features(error) if error is not None else {}), + ("dct", lambda: dct_features(gray)), + ("gradient", lambda: gradient_features(gray)), + ("color", lambda: color_features(rgb)), + ): + try: + families[name] = compute() + except Exception as exc: # one bad family must not lose the other five + logger.debug("pixel family %s failed for %s: %s", name, image_path, exc) + families[name] = {"error": f"{type(exc).__name__}: {exc}"} + + if artifacts: + try: + families["artifacts"] = spatial_artifacts( + gray, rgb, ela=error, residual=residual, phase=spectrum[1] if spectrum is not None else None + ) + except Exception as exc: + logger.debug("pixel artifacts failed for %s: %s", image_path, exc) + families["artifacts"] = {"error": f"{type(exc).__name__}: {exc}"} + + return PixelEvidence(path=image_path, decode=info, **families) diff --git a/tests/test_ai_score.py b/tests/test_ai_score.py deleted file mode 100644 index 96336e5..0000000 --- a/tests/test_ai_score.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Tests for the standalone structural AI-generation scorer.""" - -from __future__ import annotations - -import sys -from pathlib import Path -from typing import Any - -import numpy as np - -sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts")) - -import ai_score - - -def _complete_record() -> dict[str, Any]: - return { - "noise": {"noise_std": 1.0, "noise_kurtosis": 2.0}, - "fft": { - "cfa_peak": 3.0, - "cfa_peaks": [2.5, 3.0], - "fft_band_energy": list(range(8)), - }, - "ela": {"ela_mean": 4.0, "ela_p95": 5.0}, - "gradient": {"laplacian_var": 6.0, "gradient_hist": list(range(10))}, - "color": { - "saturation_mean": 0.5, - "value_mean": 0.75, - "color_hist_4x4x4": list(range(64)), - }, - "dct": { - "benford_mad": 0.1, - "dct_ac_hist": [[1] * 21 for _ in range(8)], - }, - "jpeg_forensics": { - "subsampling": "4:4:4", - "progressive": True, - "quant_tables": { - "0": list(range(1, 65)), - "1": list(range(65, 129)), - }, - "huffman_tables_hex": ["00ff", "abcd12"], - "scan_count": 10, - "restart_interval": 4, - "precision_bits": 8, - "adobe_transform": 1, - "jfif": {"version": "1.1"}, - }, - "pil": {"width": 2000, "height": 1000}, - "content_format": "jpeg", - } - - -def test_v1_feature_schema_is_fixed_for_sparse_records() -> None: - assert len(ai_score.feature_names("v1")) == 97 - assert len(ai_score.features_of({}, schema="v1")) == 97 - - -def test_v2_feature_schema_includes_existing_forensic_data() -> None: - record = _complete_record() - - names = ai_score.feature_names("v2") - values = ai_score.features_of(record, schema="v2") - by_name = dict(zip(names, values, strict=True)) - - assert len(names) == len(values) == 406 - assert by_name["cfa_peak_0"] == 2.5 - assert by_name["cfa_peak_1"] == 3.0 - assert by_name["dct_ac_0_0"] == 1 / 21 - assert by_name["dct_ac_7_20"] == 1 / 21 - assert by_name["jpeg_quant_0_0"] == 1.0 - assert by_name["jpeg_quant_1_63"] == 128.0 - assert by_name["jpeg_quant_table_count"] == 2.0 - assert by_name["jpeg_huffman_table_count"] == 2.0 - assert by_name["jpeg_huffman_total_bytes"] == 5.0 - assert by_name["jpeg_scan_count"] == 10.0 - assert by_name["jpeg_jfif_present"] == 1.0 - assert by_name["format_webp"] == 0.0 - assert by_name["format_isobmff"] == 0.0 - assert by_name["format_other"] == 0.0 - - -def test_v2_feature_schema_is_fixed_when_forensics_are_missing() -> None: - names = ai_score.feature_names("v2") - values = ai_score.features_of({}, schema="v2") - - assert len(names) == len(values) == 406 - assert np.isnan(values[names.index("dct_ac_0_0")]) - assert np.isnan(values[names.index("jpeg_quant_0_0")]) - assert np.isnan(values[names.index("jpeg_scan_count")]) - - -def test_grouped_stratified_split_keeps_hashes_on_one_side() -> None: - labels = np.asarray([1, 1, 1, 0, 0, 0, 1, 0]) - hashes = np.asarray(["a", "a", "b", "c", "c", "d", "e", "f"]) - - train, test = ai_score.grouped_stratified_split(labels, hashes, test_size=0.5, random_state=7) - - assert set(hashes[train]).isdisjoint(set(hashes[test])) - assert set(labels[train]) == {0, 1} - assert set(labels[test]) == {0, 1} - assert sorted(np.concatenate([train, test]).tolist()) == list(range(len(labels))) - - -def test_grouped_stratified_split_rejects_conflicting_labels() -> None: - labels = np.asarray([0, 1, 0, 1]) - hashes = np.asarray(["same", "same", "negative", "positive"]) - - with np.testing.assert_raises_regex(ValueError, "conflicting labels"): - ai_score.grouped_stratified_split(labels, hashes) - - -def test_temporal_holdout_excludes_hashes_seen_during_training() -> None: - dates = np.asarray(["2026-01-01", "2026-01-01", "2026-01-02", "2026-01-03", "2026-01-04", "2026-01-04"]) - hashes = np.asarray(["repeated", "old", "middle", "new-a", "repeated", "new-b"]) - - train, test, cutoff = ai_score.temporal_holdout_split(dates, hashes, train_fraction=0.5) - - assert cutoff == "2026-01-03" - assert set(hashes[train]).isdisjoint(set(hashes[test])) - assert set(hashes[test]) == {"new-a", "new-b"} - - -def test_legacy_model_bundle_defaults_to_v1_schema() -> None: - assert ai_score.model_schema({}) == "v1" - assert ai_score.model_schema({"feature_schema": "v2"}) == "v2" diff --git a/tests/test_pixel_evidence.py b/tests/test_pixel_evidence.py new file mode 100644 index 0000000..eb387f1 --- /dev/null +++ b/tests/test_pixel_evidence.py @@ -0,0 +1,156 @@ +"""Tests for the experimental pixel-forensics collector. + +It has no consumer, so there is no downstream behavior to pin. What these guard is +the part a future consumer would rely on and could not discover from the code: that +a family is empty rather than wrong when the image is too small for it, that one +failing family does not lose the other five, and that the artifacts which identify +the source image stay behind their opt-in. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import numpy as np +import pytest +from PIL import Image + +from remove_ai_watermarks.pixel_evidence import PixelEvidence, extract_pixel_evidence, is_available + +if TYPE_CHECKING: + from pathlib import Path + +FAMILIES = ("dct", "fft", "noise", "ela", "gradient", "color") + + +def _textured(path: Path, size: tuple[int, int] = (192, 160), *, seed: int = 0) -> Path: + """A textured image. Flat colour would make several families degenerate (zero + residual, empty gradient histogram) and hide a real break.""" + rng = np.random.default_rng(seed) + base = rng.integers(0, 255, (size[1], size[0], 3), dtype=np.uint8) + ramp = np.linspace(0, 255, size[0], dtype=np.uint8)[None, :, None] + Image.fromarray(np.clip(base // 2 + ramp // 2, 0, 255).astype(np.uint8)).save(path) + return path + + +class TestFamilies: + def test_every_family_is_measured_on_a_textured_image(self, tmp_path: Path): + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png")) + + assert evidence.decoded + for family in FAMILIES: + assert getattr(evidence, family), family + + def test_the_same_image_measures_the_same_twice(self, tmp_path: Path): + """Determinism is what makes these comparable across runs and machines; the + residual is computed in row chunks, which is exactly the kind of optimization + that can perturb the last bits.""" + path = _textured(tmp_path / "textured.png") + + first, second = extract_pixel_evidence(path), extract_pixel_evidence(path) + + assert dataclasses.asdict(first) == dataclasses.asdict(second) + + def test_source_dimensions_survive_the_downscale(self, tmp_path: Path): + """The recorded size is the SOURCE size, read before the 2048px cap. Recording + the analysed size instead would still pass every statistic check, because + those run on the downscaled array either way.""" + path = tmp_path / "oversize.png" + Image.fromarray(np.zeros((80, 3000, 3), dtype=np.uint8)).save(path) + + assert extract_pixel_evidence(path).decode == {"width": 3000, "height": 80} + + +class TestDegenerateInputs: + def test_undecodable_file_reports_the_error_and_stays_empty(self, tmp_path: Path): + path = tmp_path / "broken.png" + path.write_bytes(b"not an image") + + evidence = extract_pixel_evidence(path, artifacts=True) + + assert evidence.decoded is False + assert "error" in evidence.decode + assert all(getattr(evidence, family) == {} for family in FAMILIES) + assert evidence.artifacts == {} + + def test_image_too_small_for_a_family_leaves_it_empty(self, tmp_path: Path): + """8x8 is below the FFT's 32px floor but at the block DCT's. A consumer must + not assume a fixed feature width, so the narrow case is pinned.""" + path = tmp_path / "tiny.png" + Image.fromarray(np.arange(8 * 8 * 3, dtype=np.uint8).reshape(8, 8, 3)).save(path) + + evidence = extract_pixel_evidence(path) + + assert evidence.decoded is True + assert evidence.fft == {} + assert evidence.color != {} + + def test_a_failing_family_does_not_lose_the_others(self, tmp_path: Path, monkeypatch): + path = _textured(tmp_path / "textured.png") + + def boom(*args, **kwargs): + raise ValueError("family exploded") + + monkeypatch.setattr("remove_ai_watermarks.pixel_evidence.color_features", boom) + evidence = extract_pixel_evidence(path) + + assert "error" in evidence.color + assert evidence.dct != {} + assert evidence.gradient != {} + + +class TestArtifactsAreOptIn: + """The artifacts identify the source image -- a thumbnail is a picture, a + perceptual hash matches one. Everything else is a scalar or a fixed-length + histogram. That difference in kind is the reason for the flag, so the flag is + what these guard.""" + + def test_off_by_default(self, tmp_path: Path): + assert extract_pixel_evidence(_textured(tmp_path / "textured.png")).artifacts == {} + + def test_on_request_it_returns_the_spatial_layer(self, tmp_path: Path): + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png"), artifacts=True) + + assert set(evidence.artifacts) == {"phash", "thumbnail_jpeg_b64", "ela_map", "noise_residual", "fft_phase"} + assert len(evidence.artifacts["phash"]) == 16 + + def test_the_thumbnail_is_a_readable_image_of_the_source(self, tmp_path: Path): + """Stated plainly because it is the privacy claim: this field reconstructs + the picture, at 128px.""" + import base64 + import io + + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png"), artifacts=True) + + with Image.open(io.BytesIO(base64.b64decode(evidence.artifacts["thumbnail_jpeg_b64"]))) as thumb: + assert max(thumb.size) <= 128 + + def test_a_different_image_hashes_differently(self, tmp_path: Path): + one = extract_pixel_evidence(_textured(tmp_path / "a.png", seed=1), artifacts=True) + two = extract_pixel_evidence(_textured(tmp_path / "b.png", seed=2), artifacts=True) + + assert one.artifacts["phash"] != two.artifacts["phash"] + + def test_the_statistics_carry_no_array_payloads(self, tmp_path: Path): + """Without the flag nothing array-shaped may appear: that is what makes the + default set aggregates rather than content.""" + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png")) + + for family in FAMILIES: + for key, value in getattr(evidence, family).items(): + assert isinstance(value, (int, float, str, list)), (family, key) + if isinstance(value, list): + for item in value: + assert isinstance(item, (int, float, list)), (family, key) + + +def test_is_available_reports_the_optional_dependency(): + assert is_available() is True # the test environment installs the pixels extra + + +def test_evidence_is_frozen(tmp_path: Path): + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png")) + assert isinstance(evidence, PixelEvidence) + with pytest.raises(dataclasses.FrozenInstanceError): + evidence.decode = {} # type: ignore[misc]