From cdec3483495bd28c0c3fede5f70f09f117323f7f Mon Sep 17 00:00:00 2001 From: Victor Kuznetsov Date: Fri, 24 Jul 2026 19:24:39 -0700 Subject: [PATCH] Add v2 forensic features to the scorer, split without hash leakage Train now dedups by sha256 (earliest date wins), splits by hash group, and drops post-cutoff rows whose hash was seen in training, so neither holdout reports memorized duplicates. The feature schema is explicit (v1 = the original 97 structural features, v2 adds CFA peaks, DCT AC histograms and JPEG quant/Huffman/scan stats), stored in the bundle and read back at scoring time; legacy bundles default to v1. Vectors are fixed-width with NaN padding, so a sparse record no longer shifts every column. Scoring runs in batches instead of one predict_proba per record. Co-Authored-By: Claude --- scripts/ai_score.py | 338 +++++++++++++++++++++++++++++++++-------- tests/test_ai_score.py | 126 +++++++++++++++ 2 files changed, 401 insertions(+), 63 deletions(-) create mode 100644 tests/test_ai_score.py diff --git a/scripts/ai_score.py b/scripts/ai_score.py index dddedfc..539064b 100644 --- a/scripts/ai_score.py +++ b/scripts/ai_score.py @@ -4,17 +4,11 @@ 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. - -Measured on the production corpus (temporal holdout, honest for drift): -structural-feature GBM reaches AUC ~0.94-0.96 on-distribution. The model -is distribution-specific by design (it is trained on this service's -traffic); do not expect the same numbers on unrelated data. CLIP-ViT-L/14 -features were evaluated as an alternative and scored lower (0.82) while -being ~30x slower, so the shipped model uses the structural features the -scanner already collects. +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 + 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'. @@ -24,10 +18,12 @@ 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 +from typing import Any, Literal import numpy as np @@ -71,6 +67,78 @@ _FEAT_SCALAR = [ ("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: @@ -105,15 +173,15 @@ def label_of(record: dict[str, Any]) -> int | None: return None -def features_of(record: dict[str, Any]) -> list[float]: +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] - for b in (record.get("fft") or {}).get("fft_band_energy", []): - v.append(float(b)) - for h in (record.get("gradient") or {}).get("gradient_hist", []): - v.append(float(h)) - for h in (record.get("color") or {}).get("color_hist_4x4x4", []): - v.append(float(h)) + 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) @@ -122,9 +190,115 @@ def features_of(record: dict[str, Any]) -> list[float]: 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: @@ -132,49 +306,77 @@ def iter_records(pattern: str) -> Any: yield json.loads(line) -def cmd_train(pattern: str, model_path: str) -> None: +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 - from sklearn.model_selection import train_test_split - X, y, dates = [], [], [] + # keep the feature vector, not the record: a parsed scan record is an order of + # magnitude larger than the row it collapses to, and the corpus is ~5-digit. + 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 - X.append(features_of(r)) - y.append(lab) - dates.append(Path(r["file"]).parent.name) - X, y = np.array(X), np.array(y) - print(f"labeled: {len(y)} (pos {int(y.sum())}, neg {int((1 - y).sum())}, features {X.shape[1]})") + 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) - Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y) - imp = SimpleImputer(strategy="median").fit(Xtr) - clf = HistGradientBoostingClassifier(random_state=0, max_iter=300) - clf.fit(imp.transform(Xtr), ytr) - p = clf.predict_proba(imp.transform(Xte))[:, 1] - print(f"holdout: AUC {roc_auc_score(yte, p):.4f} | AP {average_precision_score(yte, p):.4f}") + 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})" + ) - uniq = sorted(set(dates)) - cutoff = uniq[int(len(uniq) * 0.7)] - tr = np.array([d < cutoff for d in dates]) - if tr.sum() > 100 and (~tr).sum() > 100: - clf_t = HistGradientBoostingClassifier(random_state=0, max_iter=300) - imp_t = SimpleImputer(strategy="median").fit(X[tr]) - clf_t.fit(imp_t.transform(X[tr]), y[tr]) - pt = clf_t.predict_proba(imp_t.transform(X[~tr]))[:, 1] - auc_t = roc_auc_score(y[~tr], pt) - ap_t = average_precision_score(y[~tr], pt) + 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}") - imp = SimpleImputer(strategy="median").fit(X) - clf = HistGradientBoostingClassifier(random_state=0, max_iter=300) - clf.fit(imp.transform(X), y) + 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({"imputer": imp, "clf": clf}, f) + pickle.dump(bundle, f) print(f"model written: {model_path}") @@ -183,27 +385,36 @@ def cmd_score(pattern: str, model_path: str, out_path: str) -> None: 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 - score = float(clf.predict_proba(imp.transform([features_of(r)]))[0, 1]) - lab = label_of(r) - out.write( - json.dumps( - { - "file": r["file"], - "sha256": r.get("sha256"), - "ai_score": round(score, 4), - "metadata_label": lab, - } - ) - + "\n" - ) - n += 1 - if n % 5000 == 0: - print(f" {n}", flush=True) + 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}") @@ -211,8 +422,9 @@ 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) == 4: - cmd_train(sys.argv[2], sys.argv[3]) + 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: diff --git a/tests/test_ai_score.py b/tests/test_ai_score.py new file mode 100644 index 0000000..96336e5 --- /dev/null +++ b/tests/test_ai_score.py @@ -0,0 +1,126 @@ +"""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"