mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-07 06:28:36 +02:00
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 <noreply@anthropic.com>
This commit is contained in:
Victor Kuznetsov
co-authored by
Claude
Claude
parent
11afe68411
commit
cdec348349
+275
-63
@@ -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 <scan_glob> <model.pkl>
|
||||
uv run --with scikit-learn python scripts/ai_score.py train <scan_glob> <model.pkl> [v1|v2]
|
||||
uv run --with scikit-learn python scripts/ai_score.py score <scan_glob> <model.pkl> <out.jsonl>
|
||||
|
||||
<scan_glob> is a glob of scan_dataset shards, e.g. 'data/scan/part_*.jsonl'.
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user