""" Pre-compute Stage 3B frequency features for the dataset splits. Mirrors `scripts/precompute_embeddings.py` in shape (same split CSVs, same `--data-root` resolution, same label map, resumable shards) but emits the 15-dim `freq-v1` feature vector from `deepfake_scanner.detectors.frequency_features` instead of CLIP embeddings. Why a separate offline step? ---------------------------- Feature extraction is pure-numpy and CPU-cheap (~a few ms/image), but the full multi-generator dataset is ~100k+ images, so caching the features once lets the classifier (S3B-07) iterate in seconds without re-reading images. Output format ------------- One `{split}_freq.npz` per split in `--out-dir`: - features: float32 (N, 16) freq-v1 feature vectors - labels: int8 (N,) 0 = authentic, 1 = ai_generated - paths: unicode (N,) manifest-relative image paths - sources: unicode (N,) - generators: unicode (N,) - model_families: unicode (N,) - augmentations: unicode (N,) - feature_schema: scalar str FEATURE_SCHEMA_VERSION ("freq-v1") - feature_names: unicode (16,) FEATURE_NAMES, for validation Resumability ------------ Each split is encoded in shards (`{split}_freq_shards/shard_NNNNNNNN.npz`); on restart we resume from the first missing contiguous shard. Final files are concatenated at the end and the shard dir removed. Idempotent: an existing final `.npz` is skipped unless `--force`. Usage ----- python scripts/precompute_frequency_features.py \\ --splits-dir data/splits \\ --out-dir data/freq_features \\ --data-root data # Include the augmented splits used for robustness training/eval: python scripts/precompute_frequency_features.py \\ --splits-dir data/splits \\ --out-dir data/freq_features \\ --data-root data \\ --splits train val test heldout test_augmented train_augmented """ from __future__ import annotations import argparse import csv import shutil import time from dataclasses import dataclass from pathlib import Path import numpy as np from PIL import Image from deepfake_scanner.detectors.frequency_features import ( FEATURE_NAMES, FEATURE_SCHEMA_VERSION, N_FEATURES, extract_features, ) # Keep label mapping identical to precompute_embeddings.py / train_head.py: # logits[..., 1] = ai_generated. _LABEL_MAP = {"authentic": 0, "ai_generated": 1} # Images per shard. Feature extraction is cheap, so use a larger shard than # the embeddings script; a crash loses at most a few seconds of work. SHARD_SIZE = 2000 @dataclass(frozen=True) class FeatureItem: label: int manifest_path: str image_path: Path source: str generator: str model_family: str augmentation: str def _read_split(csv_path: Path, data_root: Path) -> list[FeatureItem]: items: list[FeatureItem] = [] with csv_path.open(encoding="utf-8") as fh: reader = csv.DictReader(fh) for row in reader: label = _LABEL_MAP[row["class"]] manifest_path = row["path"] items.append( FeatureItem( label=label, manifest_path=manifest_path, image_path=data_root / manifest_path, source=row.get("source", ""), generator=row.get("generator", ""), model_family=row.get("model_family", ""), augmentation=row.get("augmentation", ""), ) ) return items def _resume_cursor(shard_dir: Path) -> int: """Count the contiguous valid shard prefix; delete a trailing corrupt one.""" if not shard_dir.exists(): return 0 cursor = 0 for p in sorted(shard_dir.glob("shard_*.npz")): try: with np.load(p) as z: _ = z["features"], z["labels"] cursor += int(np.load(p)["features"].shape[0]) except Exception: print(f" removing corrupt shard {p.name}") p.unlink() break return cursor def _save_shard_atomic(shard_path: Path, **arrays) -> None: tmp = shard_path.with_suffix(".tmp.npz") np.savez_compressed(tmp, **arrays) tmp.replace(shard_path) def _concat_shards(shard_dir: Path, out_path: Path, expected_total: int) -> None: paths = sorted(shard_dir.glob("shard_*.npz")) keys = [ "features", "labels", "paths", "sources", "generators", "model_families", "augmentations", ] acc: dict[str, list[np.ndarray]] = {k: [] for k in keys} for p in paths: with np.load(p) as z: for k in keys: acc[k].append(z[k].copy()) merged = {k: np.concatenate(acc[k], axis=0) for k in keys} if len(merged["labels"]) != expected_total: raise RuntimeError( f"Concatenated {len(merged['labels'])} rows but expected " f"{expected_total}. Refusing to write {out_path}." ) np.savez_compressed( out_path, feature_schema=np.asarray(FEATURE_SCHEMA_VERSION), feature_names=np.asarray(FEATURE_NAMES), **merged, ) shutil.rmtree(shard_dir) def _encode_split( split: str, splits_dir: Path, out_dir: Path, data_root: Path, force: bool, ) -> None: csv_path = splits_dir / f"{split}.csv" if not csv_path.exists(): print(f"Skipping {split} — {csv_path} not found.") return out_path = out_dir / f"{split}_freq.npz" shard_dir = out_dir / f"{split}_freq_shards" if force: if out_path.exists(): out_path.unlink() if shard_dir.exists(): shutil.rmtree(shard_dir) if out_path.exists(): print(f"Skipping {split} — {out_path} already exists. Use --force.") return items = _read_split(csv_path, data_root) n = len(items) print(f"{split}: {n} images") cursor = _resume_cursor(shard_dir) if cursor: print(f" resuming from {cursor}/{n}") shard_dir.mkdir(parents=True, exist_ok=True) t0 = time.time() while cursor < n: shard_end = min(cursor + SHARD_SIZE, n) chunk = items[cursor:shard_end] feats = np.empty((len(chunk), N_FEATURES), dtype=np.float32) for i, item in enumerate(chunk): with Image.open(item.image_path) as img: feats[i] = extract_features(img) _save_shard_atomic( shard_dir / f"shard_{cursor:08d}.npz", features=feats, labels=np.asarray([it.label for it in chunk], dtype=np.int8), paths=np.asarray([it.manifest_path for it in chunk]), sources=np.asarray([it.source for it in chunk]), generators=np.asarray([it.generator for it in chunk]), model_families=np.asarray([it.model_family for it in chunk]), augmentations=np.asarray([it.augmentation for it in chunk]), ) cursor = shard_end elapsed = time.time() - t0 rate = cursor / elapsed if elapsed else 0.0 eta = (n - cursor) / rate if rate else 0.0 print( f" {cursor}/{n} ({rate:.0f} img/s, eta {eta / 60:.1f} min, " f"shard saved)" ) print(f" concatenating shards into {out_path}...") _concat_shards(shard_dir, out_path, expected_total=n) print(f" wrote {out_path} ({out_path.stat().st_size / 1e6:.1f} MB)") def main() -> None: parser = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--splits-dir", type=Path, required=True) parser.add_argument("--out-dir", type=Path, required=True) parser.add_argument("--data-root", type=Path, required=True) parser.add_argument( "--splits", nargs="+", default=["train", "val", "test", "heldout", "test_augmented"], help="Split names (without .csv) to encode.", ) parser.add_argument( "--force", action="store_true", help="Recompute even if an output .npz already exists.", ) args = parser.parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) print(f"Feature schema: {FEATURE_SCHEMA_VERSION} ({N_FEATURES} features)") for split in args.splits: _encode_split( split=split, splits_dir=args.splits_dir, out_dir=args.out_dir, data_root=args.data_root, force=args.force, ) if __name__ == "__main__": main()