mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-07 06:28:36 +02:00
Add AI-generation scorer trained on metadata-labeled corpus
scripts/ai_score.py trains a gradient-boosted classifier on labels derived from the scan_dataset metadata (C2PA AI generators, TC260, local pipelines vs camera/screenshot/editor output) and scores every file from pixel and container statistics alone, so metadata-stripped files still get a score. Temporal holdout on the production corpus: AUC 0.966 / AP 0.988. CLIP ViT-L/14 was evaluated as an alternative and lost on both accuracy (0.82) and cost, so the model uses the structural features the scanner already collects; scoring is CPU-only and I/O-bound.
This commit is contained in:
Victor Kuznetsov
@@ -0,0 +1,224 @@
|
||||
"""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.
|
||||
|
||||
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.
|
||||
|
||||
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 score <scan_glob> <model.pkl> <out.jsonl>
|
||||
|
||||
<scan_glob> is a glob of scan_dataset shards, e.g. 'data/scan/part_*.jsonl'.
|
||||
score output: one JSON line per file with file, sha256, ai_score, and the
|
||||
label evidence when a metadata label existed (for monitoring drift).
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import pickle
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
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"),
|
||||
]
|
||||
|
||||
|
||||
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]) -> 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))
|
||||
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]
|
||||
return v
|
||||
|
||||
|
||||
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) -> 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 = [], [], []
|
||||
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]})")
|
||||
|
||||
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}")
|
||||
|
||||
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)
|
||||
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)
|
||||
with open(model_path, "wb") as f:
|
||||
pickle.dump({"imputer": imp, "clf": clf}, 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"]
|
||||
n = 0
|
||||
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)
|
||||
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) == 4:
|
||||
cmd_train(sys.argv[2], sys.argv[3])
|
||||
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()
|
||||
Reference in New Issue
Block a user