Spaces:

Veridicate
/

scanner

Running

scanner / scripts /precompute_frequency_features.py

github-actions[bot]

Deploy from GitHub 1af737237f5d10dfa9cf12849d98f47c3b30ce37 (filtered: code + Dockerfile + README + NOTICES only)

a2f55b2 about 2 months ago

8.62 kB

"""
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()

	"""
	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()