diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/aigc_detector.egg-info/PKG-INFO b/src/aigc_detector.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..562515bdaf5908c114be158b1ec9f1c7d1d66bde --- /dev/null +++ b/src/aigc_detector.egg-info/PKG-INFO @@ -0,0 +1,20 @@ +Metadata-Version: 2.4 +Name: aigc-detector +Version: 0.1.0 +Requires-Python: >=3.11 +Requires-Dist: torch>=2.2 +Requires-Dist: torchvision>=0.17 +Requires-Dist: numpy>=1.26 +Requires-Dist: pandas>=2.2 +Requires-Dist: pyarrow>=15 +Requires-Dist: pillow>=10 +Requires-Dist: scikit-learn>=1.4 +Requires-Dist: wandb>=0.16 +Requires-Dist: psutil>=5.9 +Requires-Dist: opencv-python-headless>=4.9 +Provides-Extra: dev +Requires-Dist: pytest>=8; extra == "dev" +Provides-Extra: data +Requires-Dist: datasets>=2.18; extra == "data" +Requires-Dist: huggingface_hub>=0.23; extra == "data" +Requires-Dist: modelscope>=1.13; extra == "data" diff --git a/src/aigc_detector.egg-info/SOURCES.txt b/src/aigc_detector.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..4cbe95385b70eed6165d6e8c5d835e9352eb9a0b --- /dev/null +++ b/src/aigc_detector.egg-info/SOURCES.txt @@ -0,0 +1,80 @@ +README.md +pyproject.toml +src/__init__.py +src/config.py +src/aigc_detector.egg-info/PKG-INFO +src/aigc_detector.egg-info/SOURCES.txt +src/aigc_detector.egg-info/dependency_links.txt +src/aigc_detector.egg-info/requires.txt +src/aigc_detector.egg-info/top_level.txt +src/analysis/__init__.py +src/analysis/format_probe.py +src/analysis/subset_weakness.py +src/data/__init__.py +src/data/dataset.py +src/data/preprocess.py +src/data/shard_reader.py +src/data/aug/__init__.py +src/data/aug/augment.py +src/data/aug/augment_impls.py +src/data/aug/deshortcut.py +src/data/aug/random_resolution.py +src/data/aug/symmetry_audit.py +src/data/aug/warp.py +src/data/provenance/__init__.py +src/data/provenance/audit.py +src/data/provenance/cleanup.py +src/data/provenance/dda_audit.py +src/data/provenance/dedup_audit.py +src/data/provenance/diskcheck.py +src/data/provenance/families.py +src/data/provenance/feature_cache.py +src/data/provenance/finalize.py +src/data/provenance/journal.py +src/data/provenance/manifest_build.py +src/data/provenance/orchestrate.py +src/data/provenance/push.py +src/data/provenance/rangezip.py +src/data/provenance/resolver.py +src/data/provenance/run_curator.py +src/data/provenance/selection.py +src/data/provenance/sources.py +src/data/provenance/splitter.py +src/data/provenance/subsample.py +src/data/provenance/val_holdout.py +src/data/provenance/wildfake.py +src/data/provenance/writer.py +src/eval/__init__.py +src/eval/error_gallery.py +src/eval/explain.py +src/eval/harness.py +src/eval/metrics.py +src/eval/predictions.py +src/eval/stream.py +src/inference/__init__.py +src/inference/predict.py +src/inference/server.py +src/model/__init__.py +src/model/_dinov2.py +src/model/backbone.py +src/model/detector.py +src/model/freq.py +src/model/fusion.py +src/model/head.py +src/train/__init__.py +src/train/checkpoint.py +src/train/compile_backbone.py +src/train/driver.py +src/train/logging_utils.py +src/train/loop.py +src/train/loss.py +src/train/preflight.py +src/train/watchdog.py +tests/test_aug.py +tests/test_evaluation.py +tests/test_explainability.py +tests/test_inference.py +tests/test_model.py +tests/test_preprocessing.py +tests/test_provenance.py +tests/test_training.py \ No newline at end of file diff --git a/src/aigc_detector.egg-info/dependency_links.txt b/src/aigc_detector.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/aigc_detector.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/aigc_detector.egg-info/requires.txt b/src/aigc_detector.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..6bd91b4d0eb50d4488bc35ca1365e4dc5c7f3cfa --- /dev/null +++ b/src/aigc_detector.egg-info/requires.txt @@ -0,0 +1,18 @@ +torch>=2.2 +torchvision>=0.17 +numpy>=1.26 +pandas>=2.2 +pyarrow>=15 +pillow>=10 +scikit-learn>=1.4 +wandb>=0.16 +psutil>=5.9 +opencv-python-headless>=4.9 + +[data] +datasets>=2.18 +huggingface_hub>=0.23 +modelscope>=1.13 + +[dev] +pytest>=8 diff --git a/src/aigc_detector.egg-info/top_level.txt b/src/aigc_detector.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..82eae626e32259e92b50414632e3318def29ff74 --- /dev/null +++ b/src/aigc_detector.egg-info/top_level.txt @@ -0,0 +1,8 @@ +__init__ +analysis +config +data +eval +inference +model +train diff --git a/src/analysis/__init__.py b/src/analysis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/analysis/format_probe.py b/src/analysis/format_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..8c662d2e40482878fb0d51330c81e9104c46e010 --- /dev/null +++ b/src/analysis/format_probe.py @@ -0,0 +1,115 @@ +# src/analysis/format_probe.py +"""Format/codec-shortcut probe — validates that the holdout AUC is content, not codec. + +A format-only rule ("PNG = fake, JPEG = real") scores ~0.845 on our +holdout — essentially identical to the model's 0.84. DINOv2 never sees filenames, but it DOES +see JPEG quantization/blocking/chroma-subsampling residuals, and published detectors have been +shown to treat JPEG compression as evidence of authenticity ("Fake or JPEG?", arXiv 2403.17608). +Our resolution defenses do nothing for this axis. + +The decisive test: restrict to a SINGLE codec (JPEG) present in BOTH classes, size-match, and +stratify by generator. If the AUC collapses toward 0.5 inside that controlled slice, most of the +headline is codec/provenance, not detection. If it stays strong, the shortcut is not what's +driving the score. This module computes exactly that controlled AUC alongside the raw AUC. + +`format` is derived from the image path extension (verbatim from the sniffed source format in +save_source_image), so no manifest schema change is needed — .jpg/.jpeg -> JPEG, .png -> PNG. +""" +import os + +import numpy as np +import pandas as pd + +from eval.metrics import roc_auc +from analysis.subset_weakness import resolution_band + + +def format_of(path_or_ext) -> str: + """Map a path or extension to a coarse codec label: JPEG / PNG / WEBP / OTHER.""" + ext = os.path.splitext(str(path_or_ext))[1].lower().lstrip(".") or str(path_or_ext).lower().lstrip(".") + if ext in ("jpg", "jpeg"): + return "JPEG" + if ext == "png": + return "PNG" + if ext == "webp": + return "WEBP" + return "OTHER" + + +def _format_col(df) -> pd.Series: + """Return a 'format' Series for df, using an existing column or deriving from image_path.""" + if "format" in df.columns and df["format"].notna().any(): + return df["format"].astype(str) + src = "image_path" if "image_path" in df.columns else "path" + return df[src].map(format_of) + + +def format_leakage_auc(pred_df, condition="clean") -> dict: + """How well does codec ALONE separate the classes? Scores format as an ordinal signal + (JPEG=0, PNG=1) against the label — the 'format-only classifier' baseline. + A value near the model's AUC means codec could account for most of the headline.""" + df = pred_df[pred_df["condition"] == condition] if "condition" in pred_df else pred_df + fmt = _format_col(df) + # PNG=1 leans fake, JPEG=0 leans real (our corpus polarity); WEBP/OTHER map to 0.5. + score = fmt.map({"PNG": 1.0, "JPEG": 0.0}).fillna(0.5).to_numpy() + y = df["label"].to_numpy() + auc = roc_auc(y, score) if len(np.unique(y)) == 2 else float("nan") + counts = {f"{lab}:{f}": int(n) for (lab, f), n in df.assign(_f=fmt).groupby(["label", "_f"]).size().items()} + return {"format_only_auc": float(auc), "counts": counts} + + +def format_controlled_auc(pred_df, keep_format="JPEG", condition="clean", size_match=True, + min_n=15) -> dict: + """Model AUC restricted to a single codec present in BOTH classes (default JPEG), optionally + size-band-matched, so codec cannot be the discriminator. Returns the controlled AUC, the raw + AUC on the same condition, the drop, and per-generator controlled AUCs. + + A large drop (controlled << raw, controlled -> ~0.5) is the smoking gun: the headline was + codec/provenance, not detection. A small drop means the score survives codec control.""" + df = pred_df[pred_df["condition"] == condition].copy() if "condition" in pred_df else pred_df.copy() + df["_fmt"] = _format_col(df).values + y_all, p_all = df["label"].to_numpy(), df["prob"].to_numpy() + raw_auc = roc_auc(y_all, p_all) if len(np.unique(y_all)) == 2 else float("nan") + + sub = df[df["_fmt"] == keep_format] + # size-match: keep only resolution bands where BOTH classes are present, so within the kept + # codec the classes also share the resolution support (removes the residual size confound). + matched_bands = None + if size_match and "native_w" in sub.columns and sub["native_w"].notna().any(): + sub = sub.assign(_band=sub["native_w"].map(resolution_band)) + bands_real = set(sub[sub["label"] == 0]["_band"]) + bands_fake = set(sub[sub["label"] == 1]["_band"]) + matched_bands = sorted(bands_real & bands_fake) + sub = sub[sub["_band"].isin(matched_bands)] + + y, p = sub["label"].to_numpy(), sub["prob"].to_numpy() + n_real = int((sub["label"] == 0).sum()) + n_fake = int((sub["label"] == 1).sum()) + ctrl_auc = (roc_auc(y, p) if len(np.unique(y)) == 2 and min(n_real, n_fake) >= min_n + else float("nan")) + + per_gen = [] + if "generator_id" in sub.columns: + reals = sub[sub["label"] == 0]["prob"].to_numpy() + for gen, g in sub[sub["label"] == 1].groupby("generator_id"): + fp = g["prob"].to_numpy() + yy = np.r_[np.zeros(len(reals)), np.ones(len(fp))] + pp = np.r_[reals, fp] + a = roc_auc(yy, pp) if len(fp) and len(reals) else float("nan") + per_gen.append({"generator_id": str(gen), "auc": a, "n_fake": int(len(fp)), + "reliable": len(fp) >= min_n}) + per_gen_df = pd.DataFrame(per_gen, columns=["generator_id", "auc", "n_fake", "reliable"]) + if not per_gen_df.empty: + per_gen_df = per_gen_df.sort_values("auc", na_position="last").reset_index(drop=True) + + return { + "keep_format": keep_format, + "raw_auc": float(raw_auc), + "controlled_auc": float(ctrl_auc), + "drop": float(raw_auc - ctrl_auc) if not (np.isnan(raw_auc) or np.isnan(ctrl_auc)) else float("nan"), + "n_real": n_real, "n_fake": n_fake, + "matched_bands": matched_bands, + "size_matched": bool(size_match and matched_bands is not None), + "per_generator": per_gen_df, + "reliable": min(n_real, n_fake) >= min_n, + } diff --git a/src/analysis/subset_weakness.py b/src/analysis/subset_weakness.py new file mode 100644 index 0000000000000000000000000000000000000000..40b6751663fe507ab7017a343401a054d3c19b04 --- /dev/null +++ b/src/analysis/subset_weakness.py @@ -0,0 +1,95 @@ +# src/analysis/subset_weakness.py +"""Subset-weakness analysis — WHICH slices does the detector fail on? + +The aggregate AUC hides generalization failures. A slice with AUC < 0.5 (below chance) means +the model ranks that slice's fakes as MORE real than reals — a genuine failure mode, not just +noise. This module scores each subset (generator, real_source, resolution band, format) against +the shared real pool and surfaces the weak ones, always with the sample count N so a real +failure is distinguished from a small-sample artifact. + +Motivated by the observed styleGAN ≈ 0.5 slice: GANs are a systematically weaker subset for a +DINOv2-semantic-feature detector, which averaging into a 0.9+ headline would hide. +""" +import numpy as np +import pandas as pd + +from eval.metrics import roc_auc + + +def resolution_band(w): + """Bucket a width into a coarse resolution band (for slicing by resolution).""" + w = float(w) + if w < 256: + return "<256" + if w < 512: + return "256-512" + if w < 1024: + return "512-1024" + return ">=1024" + + +def subset_weakness(pred_df, by="generator_id", condition="clean", min_n=15): + """Per-subset AUC worst-first, with N and a reliability flag. Auto-detects which class the + `by` column actually varies on (`by='real_source'` previously grouped FAKE rows silently, + which all carry real_source='none', so it never measured per-real-source weakness): + + - FAKE-side attribute (e.g. generator_id): each fake subset vs the SHARED real pool. + - REAL-side attribute (e.g. real_source): each real subset vs the SHARED fake pool. A LOW AUC + here means the detector ranks that real source as more fake than average (a real-content + weakness) — the diagnostic the report intends. + + Returns a DataFrame [subset, auc, n_fake, n_real, reliable].""" + df = pred_df[pred_df["condition"] == condition] if "condition" in pred_df else pred_df + if by not in df.columns: + return pd.DataFrame(columns=["subset", "auc", "n_fake", "n_real", "reliable"]) + reals = df[df["label"] == 0] + fakes = df[df["label"] == 1] + # which class does `by` actually distinguish? Use the side with >1 distinct non-null value. + real_varies = reals[by].nunique(dropna=True) > 1 + fake_varies = fakes[by].nunique(dropna=True) > 1 + real_side = real_varies and not fake_varies # real_source-type attribute + rows = [] + if real_side: + # each REAL subset vs the shared FAKE pool + fake_probs = fakes["prob"].to_numpy() + for subset, g in reals.groupby(by): + rp = g["prob"].to_numpy() + y = np.r_[np.zeros(len(rp)), np.ones(len(fake_probs))] + p = np.r_[rp, fake_probs] + auc = roc_auc(y, p) if len(rp) and len(fake_probs) else float("nan") + rows.append({"subset": str(subset), "auc": auc, "n_fake": int(len(fake_probs)), + "n_real": int(len(rp)), "reliable": len(rp) >= min_n}) + else: + # each FAKE subset vs the shared REAL pool + real_probs = reals["prob"].to_numpy() + for subset, g in fakes.groupby(by): + fp = g["prob"].to_numpy() + y = np.r_[np.zeros(len(real_probs)), np.ones(len(fp))] + p = np.r_[real_probs, fp] + auc = roc_auc(y, p) if len(fp) and len(real_probs) else float("nan") + rows.append({"subset": str(subset), "auc": auc, "n_fake": int(len(fp)), + "n_real": int(len(real_probs)), "reliable": len(fp) >= min_n}) + out = pd.DataFrame(rows, columns=["subset", "auc", "n_fake", "n_real", "reliable"]) + return out.sort_values("auc", na_position="last").reset_index(drop=True) + + +def find_weak_subsets(pred_df, dims=("generator_id", "real_source"), auc_threshold=0.7, + min_n=15, condition="clean"): + """Scan multiple slice dimensions and return every subset below `auc_threshold` that has + enough samples to be reliable. Adds a resolution-band dimension if native_w is present. + Returns a list of dicts {dim, subset, auc, n_fake} worst-first — the failure-mode report.""" + df = pred_df.copy() + dims = list(dims) + if "native_w" in df.columns and "res_band" not in df.columns: + df["res_band"] = df["native_w"].map(resolution_band) + dims.append("res_band") + weak = [] + for dim in dims: + if dim not in df.columns: + continue + rep = subset_weakness(df, by=dim, condition=condition, min_n=min_n) + for _, r in rep.iterrows(): + if r["reliable"] and r["auc"] < auc_threshold: + weak.append({"dim": dim, "subset": r["subset"], "auc": float(r["auc"]), + "n_fake": int(r["n_fake"])}) + return sorted(weak, key=lambda x: x["auc"]) diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000000000000000000000000000000000000..977dac6d5ebe7d227872fee88c88dca86610cffc --- /dev/null +++ b/src/config.py @@ -0,0 +1,81 @@ +# src/config.py +from dataclasses import dataclass + +@dataclass +class Config: + resolution: int = 518 + patch_grid: int = 37 + n_patch: int = 1369 + embed_dim: int = 768 + backbone: str = "stub" # "stub" (local CI) | "dinov2_vitb14_reg" (GPU host) + # REJECTED (kept for reference): RINE-style last-N fusion never beat single-layer readout. + # n_intermediate_layers: int = 1 + extract_layer: int = 8 # SHIP DEFAULT (train.py trains 8 and 10 for the ensemble) — >0 = read patch tokens after DINOv2 block L (1-indexed) and TRUNCATE the stack to L blocks (real FLOP saving); tests whether localized signal survives earlier. 0 = full depth. Mutually exclusive with n_intermediate_layers>1. + # REJECTED: MIL top-k image logits (k=16/64 and DC-centered) gave no transfer gain — the + # head's self-attention broadcasts the verdict into every patch, so top-k of a broadcast + # is the broadcast. Left commented for the record. + # mil_top_k: int = 0 + # mil_centered: bool = False + label_smooth: float = 0.05 # BCE target smoothing eps: targets eps/1−eps instead of 0/1 (bounds optimal logits, calibrates probs). 0 = hard labels. + # REJECTED: warp minimal-pair augmentation taught the warp artifact, not anatomy + # (unseen-GAN AUC −0.027). Implementation retained in data/aug/warp.py for reference. + # warp_strength: float = 1.0 + # p_global_warp: float = 0.0 + corruption_scale: float = 1.0 # train-only: >1 stretches corruption severity ranges toward their harsh end (clamped). Ship default 1.0 — scaling alone measured negative; keep at 1.0 unless re-validated together with lambda_feat. + lambda_feat: float = 1.5 # feature-consistency weight: cosine drift of the pooled trainable representation across the clean/corrupted pair (label-blind, both classes) — forces the decision path onto corruption-stable semantic/geometric directions. + fusion_project: bool = True # fusion projection: True=dense DxD (~590K); False=scalar layer-mix (K params, preferred at low data) + train_impls: dict | None = None # {family:[impl,...]} operator subset for TRAINING — set BEFORE training so held-out impls stay genuinely unseen (the impl-holdout gap is only valid when training is restricted to this subset) + data_parallel: bool = False # DataParallel across all visible GPUs (e.g. dual T4): splits the heavy frozen-backbone forward for ~2x + halves per-GPU batch memory. No idle GPU regardless of arm count. + head_hidden: int = 0 # 0 = linear-probe head; >0 = MLP head (Linear-GELU-Linear) + head_pool: str = "attn" # "attn" (baseline MIL) | "dist" (+mean/std/top-k patch-logit stats) | "nbr" (dist + zero-init 3x3 depthwise conv over the patch grid = local adjacency; density-gate redirect, orthogonal to global-realism signal) | "enc" (dist + zero-init self-attention encoder w/ 2D pos-embed = configural "sequence-not-bag" test) + # NULLED (never shipped): opt-in spectral branch — JPEG/blur corruption destroys the + # frequency signal it reads; only viable as a fusion member, which we did not ship. + # use_freq_branch: bool = False + enc_layers: int = 2 # EncoderPoolHead: transformer-encoder depth (head_pool="enc") + enc_heads: int = 8 # EncoderPoolHead: self-attention heads (head_pool="enc") + use_pos_embed: bool = True # EncoderPoolHead: add fixed 2D sin/cos positional embedding before the encoder + # SUPERSEDED by deshortcut=True (the full pipeline); kept for ablation only. + # defeat_resolution: bool = False + deshortcut: bool = False # train: thumbnail-degrade+contrast+sharpness+codec aug (stronger shortcut defeat) + p_codec: float = 0.5 # deshortcut: P(random JPEG/PNG re-encode) — kills the codec-history shortcut + random_resolution: bool = False # train: per-batch random resolution (varied token grid, no res anchor) + res_lo: int = 392 # random_resolution range lo (mult of 14) + res_hi: int = 644 # random_resolution range hi (mult of 14) + fair_eval: bool = False # eval: normalize real+fake to common size+format (no res/format tell) + feature_mode: str = "jit" # Phase-2: "jit"|"cache" — cache path not yet implemented (no consumer in Phase 1) + compile: bool = False # torch.compile toggle (GPU-host only) + amp: bool = False # fp16 autocast on the FROZEN backbone only (T4 tensor cores ~2x); head stays fp32 + gpu_preprocess: bool = False # workers emit uint8 (resize only); float/normalize/res-resize happen ON GPU (halves worker CPU, 4x smaller host transfers). Bilinear resize commutes with per-channel normalize, so numerics match the CPU path. + lambda_aug: float = 0.8 + lambda_cons_max: float = 0.3 + # NULLED (never shipped): teacher distillation from a pretrained forensics model was cut + # early on licensing + complexity grounds; the frozen-backbone head needs no teacher. + # lambda_kd: float = 0.0 + pos_weight: float = 0.0 # 0 = unweighted BCE; >0 scales the fake(1) term (calibration) + entropy_gate: float = 0.5 + stack_prob: float = 0.1 + seed: int = 0 + epochs: int = 10 + epoch_frac: float = 1.0 # train on this fraction of the loader per epoch (0.5 = half an epoch); shuffle=True makes the half-pass a random subsample + batch_size: int = 32 + eval_batch_size: int = 32 # eval forward-pass chunk (predict_fn); bounds VRAM on large val sets + eval_pool_max: int = 64 # max decoded val images resident at once (bounds HOST RAM in streaming eval) + eval_workers: int = 0 # >0 = thread-pool the eval corruption prep (PIL releases the GIL); 0 = serial (legacy) + holdout_eval_cap: int = 0 # 0=eval whole holdout; >0=stratified subsample cap (speed) + val_eval_cap: int = 0 # 0=eval whole val EVERY epoch; >0=stratified per-epoch subsample (a 7.5k val x 19 conds = hours of CPU-bound corruption PER EPOCH on the 47k corpus) + lr: float = 1e-3 + lr_schedule: str = "none" # "none" | "onecycle" (per-step mid-epoch warmup+cosine) + num_workers: int = 2 + pin_memory: bool = True + max_images: int = 0 # 0 = all + data_root: str = "data" + cache_dir: str = "cache" + # Weights & Biases logging (OFF by default; training hosts flip use_wandb=True) + use_wandb: bool = False + wandb_project: str = "track5-aigc" + wandb_run_name: str = "" # empty -> wandb autogenerates + log_every: int = 10 # log train metrics every N steps + + @classmethod + def smoke(cls) -> "Config": + return cls(max_images=200, epochs=1, compile=False, batch_size=8) diff --git a/src/data/__init__.py b/src/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/data/aug/__init__.py b/src/data/aug/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/data/aug/augment.py b/src/data/aug/augment.py new file mode 100644 index 0000000000000000000000000000000000000000..ef9a7204d7afdc77eb8d6d1c62519f872d3cb14e --- /dev/null +++ b/src/data/aug/augment.py @@ -0,0 +1,77 @@ +# src/data/augment.py +import io +import numpy as np +from PIL import Image, ImageFilter, ImageEnhance + +CORRUPTIONS = ["jpeg", "blur", "resize", "noise", "jitter", "crop"] + +EVAL_GRID = { + "jpeg": [90, 70, 50, 30], + "blur": [0.5, 1.0, 2.0], + "resize": [0.5, 0.25], + "noise": [0.02, 0.05, 0.10], + "jitter": [0.2], # ±20% + "crop": [0.8], # keep 80% +} + +# Continuous training ranges (min,max) spanning the eval grid. +_TRAIN_RANGE = { + "jpeg": (30.0, 95.0), + "blur": (0.3, 2.2), + "resize": (0.25, 0.9), + "noise": (0.01, 0.12), + "jitter": (0.05, 0.25), + "crop": (0.7, 0.95), +} + +def apply_corruption(img: Image.Image, name: str, severity: float) -> Image.Image: + img = img.convert("RGB") + if name == "jpeg": + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=int(round(severity))) + buf.seek(0) + return Image.open(buf).convert("RGB") + if name == "blur": + return img.filter(ImageFilter.GaussianBlur(radius=float(severity))) + if name == "resize": + w, h = img.size + small = img.resize((max(1, int(w*severity)), max(1, int(h*severity))), Image.BILINEAR) + return small.resize((w, h), Image.BILINEAR) + if name == "noise": + arr = np.asarray(img).astype(np.float32) / 255.0 + rng = np.random.default_rng(int(severity * 1e6)) # deterministic per severity + arr = arr + rng.normal(0, float(severity), arr.shape) + return Image.fromarray((np.clip(arr, 0, 1) * 255).astype("uint8")) + if name == "jitter": + out = ImageEnhance.Brightness(img).enhance(1 + severity) + out = ImageEnhance.Contrast(out).enhance(1 + severity) + out = ImageEnhance.Color(out).enhance(1 + severity) + return out + if name == "crop": + w, h = img.size + cw, ch = int(w*severity), int(h*severity) + left, top = (w - cw)//2, (h - ch)//2 + return img.crop((left, top, left+cw, top+ch)) + raise ValueError(f"unknown corruption {name}") + +# Direction of "harsher" per corruption: -1 = lower value is harsher (jpeg quality, +# resize/crop keep-fraction), +1 = higher value is harsher (blur radius, noise sigma, jitter). +_HARSH_DIR = {"jpeg": -1, "blur": +1, "resize": -1, "noise": +1, "jitter": +1, "crop": -1} +# Hard floors/ceilings for scaled ranges (beyond these, reals become unrecognizable and +# clean-AUC pays for it). +_HARSH_LIMIT = {"jpeg": 12.0, "blur": 4.5, "resize": 0.12, "noise": 0.25, "jitter": 0.5, + "crop": 0.5} + + +def sample_train_corruption(rng: np.random.Generator, scale: float = 1.0) -> tuple[str, float]: + """scale>1 stretches each range toward its HARSH end by (scale-1) x range-width, clamped + at _HARSH_LIMIT — 'drown the residual texture signal' training. scale=1 = original ranges.""" + name = CORRUPTIONS[rng.integers(len(CORRUPTIONS))] + lo, hi = _TRAIN_RANGE[name] + if scale > 1.0: + stretch = (scale - 1.0) * (hi - lo) + if _HARSH_DIR[name] < 0: + lo = max(_HARSH_LIMIT[name], lo - stretch) + else: + hi = min(_HARSH_LIMIT[name], hi + stretch) + return name, float(rng.uniform(lo, hi)) diff --git a/src/data/aug/augment_impls.py b/src/data/aug/augment_impls.py new file mode 100644 index 0000000000000000000000000000000000000000..497ee974694d9b1a09dbace530980874a6c2db70 --- /dev/null +++ b/src/data/aug/augment_impls.py @@ -0,0 +1,146 @@ +# src/data/augment_impls.py +"""Multi-IMPLEMENTATION corruption bank for the operator-realization holdout. + +We are scored on 6 corruptions but do NOT control the grader's exact operator (their +JPEG quantization/subsampling, their resize kernel). Training a detector to be robust to +ONE implementation risks overfitting that operator's artifact. This bank exposes several +GENUINELY DIVERGENT implementations per family so we can train on one subset and evaluate +on a held-out subset — measuring (and defending against) operator overfit. + +The divergence axes are empirically chosen: naive "PIL-jpeg vs cv2-jpeg" is a FAKE +holdout (both wrap libjpeg → byte-identical). What actually changes the artifact is the +PARAMETERIZATION — JPEG chroma subsampling (4:4:4 vs 4:2:0), resize KERNEL (bilinear / +bicubic / lanczos / area), blur border handling. Those are the implementations here. + +cv2 is an OPTIONAL backend (may be absent on minimal images); every family has at +least two PIL-only implementations so the holdout works with the standard library alone. +""" +import io +import numpy as np +from PIL import Image, ImageFilter, ImageEnhance + +try: + import cv2 as _cv2 +except Exception: # cv2 optional — PIL-only impls still give a holdout + _cv2 = None + + +# ---- JPEG: the artifact axis is chroma subsampling + codec, not the wrapper library ---- +def _jpeg_pil(img, q, subsampling): + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=int(round(q)), subsampling=subsampling) + return Image.open(io.BytesIO(buf.getvalue())).convert("RGB") + +def _jpeg_444(img, q): return _jpeg_pil(img, q, 0) # no chroma subsampling +def _jpeg_420(img, q): return _jpeg_pil(img, q, 2) # 2x2 chroma subsampling (very common) + +def _jpeg_cv2(img, q): + a = _cv2.cvtColor(np.asarray(img), _cv2.COLOR_RGB2BGR) + ok, enc = _cv2.imencode(".jpg", a, [_cv2.IMWRITE_JPEG_QUALITY, int(round(q))]) + dec = _cv2.imdecode(enc, _cv2.IMREAD_COLOR) + return Image.fromarray(_cv2.cvtColor(dec, _cv2.COLOR_BGR2RGB)) + + +# ---- RESIZE: down-then-up; the artifact axis is the KERNEL ---- +def _resize_pil(img, s, flt): + w, h = img.size + small = img.resize((max(1, int(w * s)), max(1, int(h * s))), flt) + return small.resize((w, h), flt).convert("RGB") + +def _resize_bilinear(img, s): return _resize_pil(img, s, Image.BILINEAR) +def _resize_bicubic(img, s): return _resize_pil(img, s, Image.BICUBIC) +def _resize_lanczos(img, s): return _resize_pil(img, s, Image.LANCZOS) + +def _resize_cv2_area(img, s): + a = np.asarray(img); h, w = a.shape[:2] + small = _cv2.resize(a, (max(1, int(w * s)), max(1, int(h * s))), interpolation=_cv2.INTER_AREA) + up = _cv2.resize(small, (w, h), interpolation=_cv2.INTER_LINEAR) + return Image.fromarray(up) + + +# ---- BLUR: PIL Gaussian vs a box-blur approximation vs cv2 (border handling differs) ---- +def _blur_gauss(img, sigma): return img.filter(ImageFilter.GaussianBlur(float(sigma))).convert("RGB") +def _blur_box(img, sigma): + # box radius chosen so variance ≈ Gaussian sigma² (uniform kernel, different tail) + r = max(1, int(round(float(sigma) * 1.7))) + return img.filter(ImageFilter.BoxBlur(r)).convert("RGB") +def _blur_cv2(img, sigma): + a = _cv2.GaussianBlur(np.asarray(img), (0, 0), float(sigma)) + return Image.fromarray(a) + + +# ---- NOISE: additive gaussian in float vs uint8 domain (clipping order differs) ---- +def _noise_float(img, sigma): + arr = np.asarray(img).astype(np.float32) / 255.0 + rng = np.random.default_rng(int(sigma * 1e6)) + arr = np.clip(arr + rng.normal(0, float(sigma), arr.shape), 0, 1) + return Image.fromarray((arr * 255).astype("uint8")) +def _noise_uint8(img, sigma): + arr = np.asarray(img).astype(np.int16) + rng = np.random.default_rng(int(sigma * 1e6) + 1) + arr = arr + rng.normal(0, float(sigma) * 255.0, arr.shape).astype(np.int16) + return Image.fromarray(np.clip(arr, 0, 255).astype("uint8")) + + +# ---- JITTER: order of brightness/contrast/color ops differs the result ---- +def _jitter_bcs(img, s): + out = ImageEnhance.Brightness(img).enhance(1 + s) + out = ImageEnhance.Contrast(out).enhance(1 + s) + return ImageEnhance.Color(out).enhance(1 + s).convert("RGB") +def _jitter_csb(img, s): + out = ImageEnhance.Color(img).enhance(1 + s) + out = ImageEnhance.Contrast(out).enhance(1 + s) + return ImageEnhance.Brightness(out).enhance(1 + s).convert("RGB") + + +# ---- CROP: center vs corner (both keep `s` fraction; different pixels retained) ---- +def _crop_center(img, s): + w, h = img.size; cw, ch = int(w * s), int(h * s) + l, t = (w - cw) // 2, (h - ch) // 2 + return img.crop((l, t, l + cw, t + ch)).convert("RGB") +def _crop_corner(img, s): + w, h = img.size; cw, ch = int(w * s), int(h * s) + return img.crop((0, 0, cw, ch)).convert("RGB") + + +# Registry: family -> {impl_name: fn}. cv2 impls added only when cv2 is importable. +IMPL_BANK = { + "jpeg": {"pil444": _jpeg_444, "pil420": _jpeg_420}, + "resize": {"bilinear": _resize_bilinear, "bicubic": _resize_bicubic, "lanczos": _resize_lanczos}, + "blur": {"gauss": _blur_gauss, "box": _blur_box}, + "noise": {"float": _noise_float, "uint8": _noise_uint8}, + "jitter": {"bcs": _jitter_bcs, "csb": _jitter_csb}, + "crop": {"center": _crop_center, "corner": _crop_corner}, +} +if _cv2 is not None: + IMPL_BANK["jpeg"]["cv2"] = _jpeg_cv2 + IMPL_BANK["resize"]["cv2area"] = _resize_cv2_area + IMPL_BANK["blur"]["cv2"] = _blur_cv2 + + +def list_impls(family): + return list(IMPL_BANK[family]) + + +def apply_impl(img, family, impl, severity): + return IMPL_BANK[family][impl](img.convert("RGB"), severity) + + +def split_impls_train_test(seed=0): + """Partition each family's implementations into disjoint train/test sets (the + implementation holdout). Deterministic; every family with >=2 impls contributes at + least one impl to each split. The TEST impls are never seen during training, so the + train→test AUC gap measures operator-overfit.""" + import random + rng = random.Random(seed) + train, test = {}, {} + for fam, impls in IMPL_BANK.items(): + names = sorted(impls) + rng.shuffle(names) + if len(names) == 1: # can't hold out; use in both (documented) + train[fam] = test[fam] = list(names) + continue + n_test = max(1, len(names) // 2) + test[fam] = sorted(names[:n_test]) + train[fam] = sorted(names[n_test:]) + return train, test diff --git a/src/data/aug/deshortcut.py b/src/data/aug/deshortcut.py new file mode 100644 index 0000000000000000000000000000000000000000..6817f6dd71cb6a0f6aedd84cb0f265a13e7023dc --- /dev/null +++ b/src/data/aug/deshortcut.py @@ -0,0 +1,85 @@ +# src/data/deshortcut.py +"""Shortcut-defeating augmentations. + +The eval exposed low-level tells: the COCO reals are 200px JPEG thumbnails while fakes are +large/sharp, so absolute resolution/format/sharpness separate the classes perfectly (AUC 1.0) +without any AI-detection. These augmentations push BOTH classes, label-blindly, through the +SAME low-level degradations so those statistics carry no label signal and the model must use +content: + +- thumbnail_degrade: downsize to a small size -> JPEG re-encode -> upsize back to original. + This REPRODUCES exactly what a 200px COCO thumbnail underwent, applied to every image, so a + large sharp fake sometimes looks thumbnail-degraded too. The single most targeted fix. +- randomize_contrast / randomize_sharpness: jitter the contrast and sharpness so those global + statistics (which also differ between thumbnail-reals and full-res-fakes) are not a tell. + +All are seeded by a caller-supplied numpy Generator (the dataset's label-blind rng) so the +degradation is class-independent and reproducible. +""" +import io +import numpy as np +from PIL import Image, ImageEnhance + + +def thumbnail_degrade(img, rng, min_side=96, max_side=384, jpeg_q=(40, 95)): + """Downsize -> JPEG -> upsize back to the original size, reproducing the thumbnail pipeline + the small COCO reals went through. Random intermediate size + JPEG quality per call.""" + img = img.convert("RGB") + w, h = img.size + side = int(rng.integers(min_side, max_side + 1)) + # preserve aspect: scale the SHORT side to `side` + if w <= h: + nw, nh = side, max(1, int(round(h * side / w))) + else: + nw, nh = max(1, int(round(w * side / h))), side + small = img.resize((nw, nh), Image.BILINEAR) + q = int(rng.integers(jpeg_q[0], jpeg_q[1] + 1)) + buf = io.BytesIO(); small.save(buf, format="JPEG", quality=q) + small = Image.open(io.BytesIO(buf.getvalue())).convert("RGB") + return small.resize((w, h), Image.BILINEAR) + + +def randomize_contrast(img, rng, lo=0.7, hi=1.3): + """Jitter contrast so global contrast is not a real/fake tell.""" + return ImageEnhance.Contrast(img.convert("RGB")).enhance(float(rng.uniform(lo, hi))) + + +def randomize_sharpness(img, rng, lo=0.4, hi=1.8): + """Jitter sharpness (blur<1, sharpen>1) so the sharp-fake / soft-thumbnail-real tell is + randomized out.""" + return ImageEnhance.Sharpness(img.convert("RGB")).enhance(float(rng.uniform(lo, hi))) + + +def randomize_codec(img, rng, jpeg_q=(45, 95)): + """Re-encode through a RANDOM codec so compression history carries no label signal. + + The #1 measured shortcut: our reals are all JPEG, GAN/Imagen fakes all PNG, so 'JPEG⟺real' holds + on the real side no matter what — and DINOv2 reads JPEG quantization/blocking residuals. This + picks JPEG (random quality) OR lossless PNG per call, label-blindly, so a real is sometimes + seen as clean-PNG and a fake sometimes as JPEG-compressed. Deliberately: we do + NOT blanket-JPEG (that double-compresses the already-JPEG reals while single-compressing PNG + fakes — a NEW tell); randomizing the codec per image within BOTH classes avoids it.""" + img = img.convert("RGB") + if rng.random() < 0.5: + q = int(rng.integers(jpeg_q[0], jpeg_q[1] + 1)) + buf = io.BytesIO(); img.save(buf, format="JPEG", quality=q) + return Image.open(io.BytesIO(buf.getvalue())).convert("RGB") + buf = io.BytesIO(); img.save(buf, format="PNG") # lossless re-encode + return Image.open(io.BytesIO(buf.getvalue())).convert("RGB") + + +def apply_deshortcut(img, rng, p_thumbnail=0.5, p_contrast=0.5, p_sharpness=0.5, p_codec=0.5): + """Compose the shortcut-defeating augs, each applied with its probability. Label-blind: + the caller passes the same rng distribution for real and fake, so none of these low-level + statistics (resolution, sharpness, contrast, AND codec history) can encode the label. + Returns an RGB image of the original size.""" + out = img.convert("RGB") + if rng.random() < p_thumbnail: + out = thumbnail_degrade(out, rng) + if rng.random() < p_contrast: + out = randomize_contrast(out, rng) + if rng.random() < p_sharpness: + out = randomize_sharpness(out, rng) + if rng.random() < p_codec: + out = randomize_codec(out, rng) + return out diff --git a/src/data/aug/random_resolution.py b/src/data/aug/random_resolution.py new file mode 100644 index 0000000000000000000000000000000000000000..c9344c9c26eabc159fae70e1c2920d175baf16ff --- /dev/null +++ b/src/data/aug/random_resolution.py @@ -0,0 +1,61 @@ +# src/data/random_resolution.py +"""Per-batch random resolution (user-directed). + +Forcing every image to a fixed 518×518 bakes the resolution difference in (thumbnails always +upscaled, large images always downscaled) and lets the model memorize a fixed 37×37 token +layout. Instead we resize each BATCH to a random valid resolution (a multiple of the patch +size), so: + - the same content is seen at varied token-grid densities (a general robustness aug), and + - absolute source resolution is randomized across batches, removing an anchor for the + resolution shortcut. + +DINOv2 interpolates its position embeddings for arbitrary resolution natively, and the MIL +head is patch-count-agnostic, so no model change is needed. Per-BATCH (not per-image) because +a batch is one tensor and must share a size anyway — which is also the cheap, worker-safe +implementation the user suggested ("the entire batch can resize"). +""" +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data.dataloader import default_collate + + +def valid_resolutions(patch=14, lo=392, hi=644): + """Resolutions that are multiples of `patch` in [lo, hi] — the sizes DINOv2 accepts.""" + k0 = (lo + patch - 1) // patch + k1 = hi // patch + return [k * patch for k in range(k0, k1 + 1)] + + +class _RandomResolutionCollate: + """Picklable collate callable (a plain closure isn't picklable, so DataLoader workers with + num_workers>0 can't use it). Stacks a batch (default_collate) then resizes the image tensors + to ONE random valid resolution for that batch. Non-image fields pass through unchanged. + lo==hi disables randomization (deterministic base size). Each worker seeds independently.""" + def __init__(self, patch=14, lo=392, hi=644, seed=0, image_keys=("clean_img", "aug_img")): + self.sizes = valid_resolutions(patch, lo, hi) + self.seed = seed + self.image_keys = image_keys + self._rng = None # lazily created per-process (post-fork/pickle) + + def __call__(self, samples): + if self._rng is None: + import os + self._rng = np.random.default_rng((self.seed, os.getpid())) + batch = default_collate(samples) + target = int(self.sizes[self._rng.integers(len(self.sizes))]) if len(self.sizes) > 1 else self.sizes[0] + for k in self.image_keys: + if k in batch and torch.is_tensor(batch[k]) and batch[k].shape[-1] != target: + if batch[k].dtype == torch.uint8: + # gpu_preprocess path: DEFER the resize to the GPU (gpu_finish) — + # interpolating a big fp32 batch here was pure worker-CPU burn. + batch["target_res"] = target + else: + batch[k] = F.interpolate(batch[k], size=(target, target), + mode="bilinear", align_corners=False) + return batch + + +def random_resolution_collate(patch=14, lo=392, hi=644, seed=0, image_keys=("clean_img", "aug_img")): + """Build a picklable per-batch random-resolution collate (see _RandomResolutionCollate).""" + return _RandomResolutionCollate(patch, lo, hi, seed, image_keys) diff --git a/src/data/aug/symmetry_audit.py b/src/data/aug/symmetry_audit.py new file mode 100644 index 0000000000000000000000000000000000000000..7f3562d30bbd33b7fff5667929df11c09d64b39f --- /dev/null +++ b/src/data/aug/symmetry_audit.py @@ -0,0 +1,53 @@ +# src/data/symmetry_audit.py +"""Empirical class-symmetry audit. + +Class-INDEPENDENT augmentation sampling is necessary but not sufficient for "no shortcut": +real and fake images can still receive different REALIZED corruptions if the sources differ +in resolution, format, or transform failure rate. The decisive test is empirical — fit a +cheap classifier on the augmentation METADATA alone (which corruption, severity, realized +dims) and check its AUC at predicting the label. If aug metadata predicts the label, the +detector could learn "was-corrupted-this-way" instead of "is-synthetic" — the shortcut that +"Fake or JPEG?" exposed. A passing audit shows AUC ≈ 0.5. +""" +import numpy as np +import pandas as pd + + +def aug_metadata_leakage_auc(df, seed=0): + """AUC of a logistic model predicting `label` from augmentation metadata only + (aug_id one-hot + severity + any realized-dim columns present). ~0.5 = no leakage.""" + from sklearn.linear_model import LogisticRegression + from sklearn.model_selection import cross_val_predict + from sklearn.metrics import roc_auc_score + + y = df["label"].to_numpy().astype(int) + if len(np.unique(y)) < 2: + return 0.5 + feats = [] + if "aug_id" in df: + feats.append(pd.get_dummies(df["aug_id"], prefix="aug")) + for col in ("severity", "out_w", "out_h", "out_bytes"): + if col in df: + feats.append(df[[col]].astype(float)) + if not feats: + return 0.5 + X = pd.concat(feats, axis=1).to_numpy(dtype=float) + # cross-validated out-of-fold probabilities -> honest AUC (no train-on-test optimism) + clf = LogisticRegression(max_iter=500) + try: + proba = cross_val_predict(clf, X, y, cv=5, method="predict_proba")[:, 1] + except Exception: + clf.fit(X, y); proba = clf.predict_proba(X)[:, 1] + return float(roc_auc_score(y, proba)) + + +def audit_class_symmetry(df, threshold=0.6, seed=0): + """Report whether augmentation metadata leaks the label. passed = AUC <= threshold.""" + auc = aug_metadata_leakage_auc(df, seed=seed) + return { + "leakage_auc": auc, + "passed": auc <= threshold, + "threshold": threshold, + "note": ("aug metadata should NOT predict the label; AUC>threshold means real/fake " + "receive systematically different corruptions -> a 'was-corrupted' shortcut."), + } diff --git a/src/data/aug/warp.py b/src/data/aug/warp.py new file mode 100644 index 0000000000000000000000000000000000000000..64b85b9a33dfb296a5360e04394416d28d64035c --- /dev/null +++ b/src/data/aug/warp.py @@ -0,0 +1,120 @@ +# src/data/warp.py +"""Anatomical-plausibility augmentation: local discontinuous warps. + +Motivation (run-4..8 + XAI): every reader of DINOv2-B features hits the same GigaGAN +ceiling — texture separates train classes so cheaply that anatomy/arrangement is never +learned. A real photo and its LOCALLY-warped twin form a minimal pair: identical camera, +JPEG history, and texture statistics, differing ONLY in local geometric plausibility. With +the warped twin labeled FAKE, texture is uninformative on the pair by construction and the +gradient's only descent direction is the geometry axis (the six-fingers class of wrongness). + +Two warp families with OPPOSITE roles: + local_deform -> label-flipping (real -> warped twin labeled fake): 1-3 localized mesh + displacements — a region's geometry disagrees with its surroundings. + global_warp -> class-SYMMETRIC aug (both classes, label-blind): mild affine/perspective, + breaks no anatomy. Its role is to put warp-resampling smear on BOTH + classes so "was resampled" cannot become the shortcut for "fake". + +Both re-quantize through JPEG so interpolation smoothness is not the tell either. +""" +import io + +import numpy as np +from PIL import Image + + +def _jpeg(img: Image.Image, quality: int) -> Image.Image: + buf = io.BytesIO() + img.save(buf, "JPEG", quality=int(quality)) + buf.seek(0) + return Image.open(buf).convert("RGB") + + +def _remap_mesh(arr: np.ndarray, dx: np.ndarray, dy: np.ndarray) -> np.ndarray: + """Bilinear remap of HxWx3 uint8 by per-pixel displacement fields (float32 HxW).""" + h, w = arr.shape[:2] + yy, xx = np.meshgrid(np.arange(h, dtype=np.float32), np.arange(w, dtype=np.float32), + indexing="ij") + sx = np.clip(xx + dx, 0, w - 1) + sy = np.clip(yy + dy, 0, h - 1) + x0 = np.floor(sx).astype(np.int32); y0 = np.floor(sy).astype(np.int32) + x1 = np.minimum(x0 + 1, w - 1); y1 = np.minimum(y0 + 1, h - 1) + fx = (sx - x0)[..., None]; fy = (sy - y0)[..., None] + a = arr[y0, x0].astype(np.float32); b = arr[y0, x1].astype(np.float32) + c = arr[y1, x0].astype(np.float32); d = arr[y1, x1].astype(np.float32) + out = (a * (1 - fx) + b * fx) * (1 - fy) + (c * (1 - fx) + d * fx) * fy + return np.clip(out, 0, 255).astype(np.uint8) + + +def local_deform(img: Image.Image, rng: np.random.Generator, + strength: float = 1.0, jpeg_q: (int, int) = (70, 92), + max_side: int = 640) -> Image.Image: + """1-3 localized smooth-bump displacements — geometry locally inconsistent with the + surround, texture untouched. strength scales displacement amplitude. + + CPU note: the numpy remap is O(H*W); native-res warping of a 1500px image is ~6x the + work of 640px for zero semantic gain (deshortcut crushes to ~200px downstream anyway). + Warp at <=max_side and restore the ORIGINAL size afterward, so the twin's resolution + statistics stay identical to its real original (no size leak).""" + orig_size = img.size + img = img.convert("RGB") + if max(img.size) > max_side: + s = max_side / max(img.size) + img = img.resize((max(1, int(img.size[0] * s)), max(1, int(img.size[1] * s))), + Image.BILINEAR) + arr = np.asarray(img) + h, w = arr.shape[:2] + dx = np.zeros((h, w), np.float32); dy = np.zeros((h, w), np.float32) + n = int(rng.integers(1, 4)) + for _ in range(n): + # bump center anywhere; radius 8-25% of min side; amplitude 4-12% of radius x strength + cy, cx = rng.uniform(0.15, 0.85) * h, rng.uniform(0.15, 0.85) * w + rad = rng.uniform(0.08, 0.25) * min(h, w) + amp = rng.uniform(0.15, 0.45) * rad * strength + ang = rng.uniform(0, 2 * np.pi) + yy, xx = np.meshgrid(np.arange(h, dtype=np.float32), np.arange(w, dtype=np.float32), + indexing="ij") + d2 = ((yy - cy) ** 2 + (xx - cx) ** 2) / (rad ** 2) + bump = np.exp(-d2 * 2.0).astype(np.float32) # smooth, compact support + dx += np.cos(ang) * amp * bump + dy += np.sin(ang) * amp * bump + out_img = Image.fromarray(_remap_mesh(arr, dx, dy)) + if out_img.size != orig_size: + out_img = out_img.resize(orig_size, Image.BILINEAR) # restore native size (no leak) + q = int(rng.integers(jpeg_q[0], jpeg_q[1] + 1)) + return _jpeg(out_img, q) + + +def global_warp(img: Image.Image, rng: np.random.Generator, + jpeg_q: (int, int) = (70, 92)) -> Image.Image: + """Mild global perspective/affine jitter (anatomy-preserving), then JPEG. Class-symmetric: + exists so BOTH classes carry warp-resampling residue.""" + w, h = img.size + j = 0.04 # corner jitter, fraction of side + src = [(0, 0), (w, 0), (w, h), (0, h)] + dst = [(x + rng.uniform(-j, j) * w, y + rng.uniform(-j, j) * h) for x, y in src] + # solve the 8-dof perspective transform mapping dst -> src (PIL wants inverse coeffs) + A = [] + for (x, y), (u, v) in zip(dst, src): + A.append([x, y, 1, 0, 0, 0, -u * x, -u * y]) + A.append([0, 0, 0, x, y, 1, -v * x, -v * y]) + B = np.array([c for (u, v) in src for c in (u, v)], dtype=np.float64) + coeffs = np.linalg.solve(np.array(A, dtype=np.float64), B) + out = img.transform((w, h), Image.PERSPECTIVE, tuple(coeffs), resample=Image.BILINEAR) + q = int(rng.integers(jpeg_q[0], jpeg_q[1] + 1)) + return _jpeg(out, q) + + +def add_warp_fakes(train_df, frac: float, seed: int = 0): + """Duplicate a `frac` sample of REAL train rows as warped fakes: label=1, warp_fake=1, + generator_id='warp_real'. The original real rows stay — the minimal pair is explicit + (same content under both labels; only the warp separates them).""" + reals = train_df[train_df.label == 0] + dup = reals.sample(int(len(reals) * frac), random_state=seed).copy() + dup["label"] = 1 + dup["warp_fake"] = 1 + dup["generator_id"] = "warp_real" + out = train_df.copy() + out["warp_fake"] = 0 + import pandas as pd + return pd.concat([out, dup], ignore_index=True) diff --git a/src/data/dataset.py b/src/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..c8a82070195cac6fed9a35c62ef63c50e9abf620 --- /dev/null +++ b/src/data/dataset.py @@ -0,0 +1,108 @@ +# src/data/dataset.py +import numpy as np +from pandas import notna as pd_notna +from PIL import Image +from torch.utils.data import Dataset +from data.aug.augment import apply_corruption, sample_train_corruption +from data.preprocess import preprocess_image + +class AigcDataset(Dataset): + def __init__(self, df, cfg, train=True, resolver=None): + self.df = df.reset_index(drop=True) + self.cfg = cfg + self.train = train + self.res = cfg.resolution + self.resolver = resolver + # Epoch folded into the aug seed so each image gets a DIFFERENT corruption realization + # every epoch (a seed of only (cfg.seed, img_id) froze the aug across epochs, + # killing augmentation diversity). Epoch is class-INDEPENDENT, so label-blindness holds. + # The driver calls set_epoch(e) before each epoch's loop. + self.epoch = 0 + + def set_epoch(self, epoch): + self.epoch = int(epoch) + + def __len__(self): + return len(self.df) + + def _prep(self, img): + if getattr(self.cfg, "gpu_preprocess", False): + from data.preprocess import preprocess_image_uint8 + return preprocess_image_uint8(img, self.res) # float+normalize done on GPU + return preprocess_image(img, self.res) + + def __getitem__(self, i): + row = self.df.iloc[i] + if self.resolver is not None and "ref_type" in row: + img = self.resolver.resolve(row) + else: + with Image.open(row["path"]) as im: + img = im.convert("RGB") + # LABEL-BLIND seed: derive the aug RNG from the IMMUTABLE img_id, not the row index + # i (which correlates with class when the manifest is label-sorted). This guarantees + # the realized corruption cannot carry label signal via seeding. + _idkey = row.get("img_id") if "img_id" in row and pd_notna(row.get("img_id")) else \ + row.get("rel_path", row.get("path", i)) + # epoch is part of the seed (train only) so the realized corruption VARIES across epochs + # while staying label-blind (epoch does not depend on the label). + _epoch = self.epoch if self.train else 0 + seed_key = abs(hash((self.cfg.seed, _epoch, str(_idkey)))) % (2**32) + rng = np.random.default_rng(seed_key) + # DEFEAT THE RESOLUTION/FORMAT/SHARPNESS SHORTCUT (label-blind, both classes): the eval + # showed reals (200px JPEG thumbnails) vs fakes (large/sharp) separate by low-level + # statistics ALONE (resolution-only AUC 1.0). apply_deshortcut pushes every image through + # the SAME degradations — thumbnail(downsize->JPEG->upsize, reproducing what the COCO + # reals underwent) + contrast + sharpness jitter — so those stats carry no label signal. + # Applied BEFORE clean/aug so both views inherit it. (defeat_resolution kept as the + # lighter resolution-only variant for ablation.) + # WARP-FAKE rows (label-flipped duplicates of reals, add_warp_fakes): apply the local + # deform FIRST so the twin then rides the identical deshortcut+corruption pipeline as + # its real original — the minimal pair differs ONLY in local geometry. Seed offset by a + # constant so the twin's downstream aug draw still matches nothing label-correlated. + if row.get("warp_fake", 0) == 1: + from data.aug.warp import local_deform + img = local_deform(img, np.random.default_rng((seed_key + 777) % (2**32)), + strength=getattr(self.cfg, "warp_strength", 1.0)) + # class-SYMMETRIC global warp (both classes, label-blind): puts warp-resampling residue + # on reals and fakes alike so "was resampled" cannot proxy for "warp_fake". + if self.train and rng.random() < getattr(self.cfg, "p_global_warp", 0.0): + from data.aug.warp import global_warp + img = global_warp(img, rng) + if self.train and getattr(self.cfg, "deshortcut", False): + from data.aug.deshortcut import apply_deshortcut + img = apply_deshortcut(img, rng, p_codec=getattr(self.cfg, "p_codec", 0.5)) + elif self.train and getattr(self.cfg, "defeat_resolution", False): + from data.preprocess import defeat_resolution_shortcut + img = defeat_resolution_shortcut(img, self.res, rng) + name, sev = sample_train_corruption( + rng, scale=getattr(self.cfg, "corruption_scale", 1.0) if self.train else 1.0) + # CLASS-SYMMETRIC by construction: corruption is sampled from a class-INDEPENDENT + # distribution and applied to every image, so "was-corrupted" carries no label + # signal (prevents the JPEG/resize dataset-shortcut that "Fake or JPEG?" exposed). + # When train_impls is set, draw from the TRAINING implementation subset only, so the + # held-out implementations stay genuinely unseen for the invariance measurement. + train_impls = getattr(self.cfg, "train_impls", None) + if train_impls: + from data.aug.augment_impls import apply_impl + allowed = train_impls.get(name) + impl = allowed[int(rng.integers(len(allowed)))] if allowed else None + aug = apply_impl(img, name, impl, sev) if impl else apply_corruption(img, name, sev) + else: + aug = apply_corruption(img, name, sev) + # img_id: the curator manifest identifies images by rel_path, not img_id — fall back + # to rel_path (then path, then row index) so the dataset reads any manifest schema. + img_id = row.get("img_id") if "img_id" in row and pd_notna(row.get("img_id")) else \ + row.get("rel_path", row.get("path", i)) + out = { + "clean_img": self._prep(img), + "aug_img": self._prep(aug), + "aug_id": name, + "severity": float(sev), + "label": float(row["label"]), + "source": row["source_dataset"], + "generator_id": row["generator_id"], + "img_id": str(img_id), + } + if "teacher_prob" in row and row["teacher_prob"] == row["teacher_prob"]: # not NaN + out["teacher_prob"] = float(row["teacher_prob"]) + return out diff --git a/src/data/preprocess.py b/src/data/preprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..d2d9fd3e2fcb0e747bbf4afdfc818f0e3d68915d --- /dev/null +++ b/src/data/preprocess.py @@ -0,0 +1,76 @@ +# src/data/preprocess.py +import io +from PIL import Image +from torchvision import transforms + +_MEAN = [0.485, 0.456, 0.406] +_STD = [0.229, 0.224, 0.225] +_NORM = transforms.Normalize(mean=_MEAN, std=_STD) +_TO = transforms.ToTensor() + + +def preprocess_image_uint8(img, resolution): + """CPU-cheap half of preprocessing: resize + HWC->CHW uint8 tensor. The float + conversion + normalization (the expensive per-pixel CPU work) moves to the GPU + via gpu_finish — workers then ship 4x fewer bytes and burn ~half the CPU.""" + import numpy as np + import torch + img = img.convert("RGB").resize((resolution, resolution), Image.BILINEAR) + arr = np.asarray(img, dtype=np.uint8) # HWC + return torch.from_numpy(arr).permute(2, 0, 1).contiguous() + + +def gpu_finish(x, target_res=None): + """Finish preprocessing ON DEVICE for a uint8 [B,C,H,W] batch: float/255, + ImageNet-normalize, optional resize to the batch's random resolution. Exactly + equivalent to ToTensor+Normalize (ToTensor is uint8/255) + collate interpolate, + minus the worker CPU. No-op passthrough for float inputs (already normalized).""" + import torch + import torch.nn.functional as F + if x.dtype == torch.uint8: + x = x.float().div_(255.0) + mean = torch.tensor(_MEAN, device=x.device).view(1, 3, 1, 1) + std = torch.tensor(_STD, device=x.device).view(1, 3, 1, 1) + x = x.sub_(mean).div_(std) + if target_res is not None and x.shape[-1] != int(target_res): + x = F.interpolate(x, size=(int(target_res), int(target_res)), + mode="bilinear", align_corners=False) + return x + +def preprocess_image(img, resolution): + img = img.convert("RGB").resize((resolution, resolution), Image.BILINEAR) + return _NORM(_TO(img)) + + +def defeat_resolution_shortcut(img, resolution, rng): + """Randomize the resolution PATH so absolute source resolution carries no label signal. + + The eval exposed a catastrophic shortcut: reals were all 200x200 thumbnails and fakes were + large, so a plain resize-to-518 leaves reals upscale-blurred and fakes downscale-sharp — a + 'was-this-upscaled' tell (resolution-only AUC = 1.0). Here every image is first resized to a + RANDOM intermediate size (thumbnailing large images, and sometimes down-then-up), then to the + model resolution, with a random interpolation. Applied to BOTH classes, this randomizes out + the upscale/downscale signature so the model must use content, not resolution.""" + img = img.convert("RGB") + # random intermediate short-side in [96, 1.4x resolution]; random interp both ways + inter = int(rng.integers(96, int(resolution * 1.4))) + filt_down = rng.choice([Image.BILINEAR, Image.BICUBIC, Image.LANCZOS, Image.NEAREST]) + filt_up = rng.choice([Image.BILINEAR, Image.BICUBIC, Image.NEAREST]) + small = img.resize((inter, inter), filt_down) + return small.resize((resolution, resolution), filt_up) + + +def fair_eval_normalize(img, resolution, jpeg_quality=90, common_small=200): + """Fair-eval normalizer: equalize real and fake so absolute resolution/format is not the + discriminator. Push every image DOWN to a small common size (`common_small`, matching the + holdout reals' 200px thumbnails) + JPEG re-encode, THEN up to the model resolution. + + Downsizing to the common-small size is load-bearing: an earlier version that resized straight + to `resolution` (518) UPSCALED the small reals and AMPLIFIED the sharpness gap (a small real + upscaled gets even smoother; a large fake downscaled stays sharp) — sharpness-oracle AUC rose + 0.63→0.77. Matching everything DOWN to the real's information level destroys the fake's detail + advantage instead (oracle AUC → ~0.4), so any residual AUC is content, not resolution.""" + img = img.convert("RGB").resize((common_small, common_small), Image.BILINEAR) + buf = io.BytesIO(); img.save(buf, format="JPEG", quality=jpeg_quality) # common format + img = Image.open(io.BytesIO(buf.getvalue())).convert("RGB") + return img.resize((resolution, resolution), Image.BILINEAR) # up to model res diff --git a/src/data/provenance/__init__.py b/src/data/provenance/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/data/provenance/audit.py b/src/data/provenance/audit.py new file mode 100644 index 0000000000000000000000000000000000000000..0006c69e68733b8e29935f984b8bdfbd97e8266a --- /dev/null +++ b/src/data/provenance/audit.py @@ -0,0 +1,17 @@ +# src/curator/audit.py +from data.provenance.dda_audit import audit_class_separability + +def audit_corpus(df) -> dict: + dda = audit_class_separability(df) + per_family = df.groupby("family").size().to_dict() + n_real = int((df["label"] == 0).sum()) + n_fake = int((df["label"] == 1).sum()) + if dda["flag"]: + report = (f"SHORTCUT FLAGGED: resolution_auc={dda['resolution_auc']:.3f} " + f"compression_auc={dda['compression_auc']:.3f} (>0.65). " + f"Mitigate: resolution/JPEG-normalize BOTH classes or rebalance sources, then re-audit.") + else: + report = (f"clean: resolution_auc={dda['resolution_auc']:.3f} " + f"compression_auc={dda['compression_auc']:.3f} (both <=0.65).") + return {**dda, "per_family_counts": per_family, + "class_balance": {"n_real": n_real, "n_fake": n_fake}, "report": report} diff --git a/src/data/provenance/cleanup.py b/src/data/provenance/cleanup.py new file mode 100644 index 0000000000000000000000000000000000000000..84711816f9435b1ee04325495b1fdec6c5150a53 --- /dev/null +++ b/src/data/provenance/cleanup.py @@ -0,0 +1,60 @@ +# src/shards/cleanup.py +"""Marker-guarded cleanup. Every recursive delete must target a directory that the +pipeline itself created (contains .track5_run_root) — filesystem roots, homes, the +workspace, and anything unmarked are refused. All operations report item count and +bytes reclaimed so GC is observable, not best-effort.""" +import os +import shutil +import tempfile + +MARKER = ".track5_run_root" + + +def make_run_root(parent, prefix="run_"): + root = tempfile.mkdtemp(prefix=prefix, dir=parent) + open(os.path.join(root, MARKER), "w").close() + return root + + +def _tree_stats(target): + n, b = 0, 0 + for dp, _dirs, files in os.walk(target): + for f in files: + n += 1 + try: + b += os.path.getsize(os.path.join(dp, f)) + except OSError: + pass + return n, b + + +def safe_rmtree(target): + resolved = os.path.realpath(target) + home = os.path.realpath(os.path.expanduser("~")) + cwd = os.path.realpath(os.getcwd()) + if (not resolved or resolved == os.path.sep or os.path.dirname(resolved) == resolved + or resolved in (home, cwd)): + raise ValueError(f"refusing to delete unsafe target: {resolved}") + if not os.path.isfile(os.path.join(resolved, MARKER)): + raise ValueError(f"refusing: {resolved} has no {MARKER} marker") + n, b = _tree_stats(resolved) + shutil.rmtree(resolved) + return n, b + + +def reclaim_unreferenced(raw_root, referenced_relpaths): + """Delete every file under raw_root whose posix relpath is NOT in + referenced_relpaths; prune empty dirs. Returns (n_deleted, bytes_reclaimed).""" + referenced = set(referenced_relpaths) + n, b = 0, 0 + for dp, _dirs, files in os.walk(raw_root, topdown=False): + for f in files: + p = os.path.join(dp, f) + rel = os.path.relpath(p, raw_root).replace(os.sep, "/") + if rel not in referenced: + b += os.path.getsize(p) + os.remove(p) + n += 1 + if dp != raw_root and not os.listdir(dp): + os.rmdir(dp) + return n, b diff --git a/src/data/provenance/dda_audit.py b/src/data/provenance/dda_audit.py new file mode 100644 index 0000000000000000000000000000000000000000..4d893aa210c96f85d0b08be3afbd5979eb67b080 --- /dev/null +++ b/src/data/provenance/dda_audit.py @@ -0,0 +1,16 @@ +# src/manifest/dda_audit.py +import numpy as np +from eval.metrics import roc_auc + +def audit_class_separability(df) -> dict: + labels = df["label"].to_numpy() + res = (df["native_w"] * df["native_h"]).to_numpy().astype(float) + comp = df["jpeg_bytes_per_pixel"].to_numpy().astype(float) + # AUC of a single scalar predicting label; flip so AUC>=0.5 + def _one_sided(x): + a = roc_auc(labels, x) + return max(a, 1 - a) if not np.isnan(a) else float("nan") + res_auc = _one_sided(res) + comp_auc = _one_sided(comp) + flag = bool((res_auc > 0.65) or (comp_auc > 0.65)) + return {"resolution_auc": res_auc, "compression_auc": comp_auc, "flag": flag} diff --git a/src/data/provenance/dedup_audit.py b/src/data/provenance/dedup_audit.py new file mode 100644 index 0000000000000000000000000000000000000000..93753927e1ceb70ae896ca349a0e799f6d544e83 --- /dev/null +++ b/src/data/provenance/dedup_audit.py @@ -0,0 +1,184 @@ +# src/curator/dedup_audit.py +"""Near-duplicate audit + dedup for the corpus manifest. OmniFake (the aggregator) +is known to partially CONTAIN WildFake+GenImage images, so sources overlap; this +module finds near-identical images across sources via a DCT pHash and collapses +each duplicate cluster down to one row, per a source-priority policy.""" +import numpy as np +import pandas as pd +from PIL import Image + +from data.provenance.manifest_build import MANIFEST_COLUMNS + +DEFAULT_PRIORITY = ["wildfake_valset", "wildfake", "sid_set", "commforensics", + "deepface", "omnifake", "cifake"] + +def _dct_matrix(n: int) -> np.ndarray: + """(n,n) DCT-II basis: D[k,n] = cos(pi*(2n+1)*k/(2N)). dct2(img) = D @ img @ D.T.""" + k = np.arange(n).reshape(-1, 1) + x = np.arange(n).reshape(1, -1) + return np.cos(np.pi * (2 * x + 1) * k / (2 * n)) + +def phash(image_path, hash_size=8) -> int: + """DCT-based pHash, numpy-only (no imagehash/scipy dependency). Deterministic: + same file -> same int every call. Robust to JPEG re-encode (low-freq DCT + coefficients barely move under quantization noise).""" + n = hash_size * 4 + with Image.open(image_path) as im: + img = np.asarray(im.convert("L").resize((n, n), Image.Resampling.LANCZOS), dtype=np.float64) + d = _dct_matrix(n) + dct2 = d @ img @ d.T + block = dct2[:hash_size, :hash_size] + # Exclude the [0,0] DC term from the median (it's just average brightness, + # not structure) so the threshold reflects the actual frequency content. + flat = block.flatten() + median = np.median(np.delete(flat, 0)) + bits = (flat > median) + h = 0 + for b in bits: + h = (h << 1) | int(b) + return h + +def hamming(a: int, b: int) -> int: + return (a ^ b).bit_count() + +_POP_C1 = np.uint64(0x5555555555555555) +_POP_C2 = np.uint64(0x3333333333333333) +_POP_C4 = np.uint64(0x0f0f0f0f0f0f0f0f) +_POP_H01 = np.uint64(0x0101010101010101) + +def _popcount64(x): + """Vectorized 64-bit popcount over a uint64 numpy array (SWAR / bit-twiddle). ~50x faster + than a Python-per-element loop, so the Hamming all-pairs scan is dominated by SIMD, not + interpreter overhead.""" + x = x - ((x >> np.uint64(1)) & _POP_C1) + x = (x & _POP_C2) + ((x >> np.uint64(2)) & _POP_C2) + x = (x + (x >> np.uint64(4))) & _POP_C4 + return (x * _POP_H01) >> np.uint64(56) + +def _within_threshold_neighbors(hashes_u64, threshold): + """Yield (i, j>i) index pairs with Hamming(phash) <= threshold, computed row-at-a-time + with a VECTORIZED popcount (bounded memory O(n), one XOR-broadcast per row). This is the + O(n^2)-comparisons scan but with a tiny constant factor — measured faster than LSH banding + on mostly-unique real-image hashes (banding degrades to millions of chance-collision + candidates when bands are narrow). Exact: no false negatives, no tuning.""" + n = len(hashes_u64) + for i in range(n): + if i + 1 >= n: + break + d = _popcount64(hashes_u64[i] ^ hashes_u64[i + 1:]) + for off in np.nonzero(d <= threshold)[0]: + yield i, i + 1 + int(off) + + +def _phash_many(paths, workers=8): + """Compute phash for a list of paths, in PARALLEL over `workers` threads, preserving + input order. pHash is decode(disk-I/O) + DCT(numpy, releases the GIL), so threads overlap + the I/O and the released-GIL compute — measured ~20-27x over sequential even on 2 cores. + workers<=1 runs sequentially.""" + if workers is None or workers <= 1 or len(paths) < 32: + return [phash(p) for p in paths] + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=workers) as ex: + return list(ex.map(phash, paths)) # ex.map preserves order + + +def add_phash(df, workers=8): + """Copy of df with a uint64 `phash` column (64-bit DCT pHash of each row's file).""" + out = df.copy() + hashes = _phash_many(out["path"].tolist(), workers=workers) + out["phash"] = np.asarray([int(h) & ((1 << 64) - 1) for h in hashes], dtype=np.uint64) + return out + + +def cross_leak_mask(cand_hashes, hold_hashes, threshold=6): + """Boolean mask over cand_hashes: True where Hamming distance to ANY holdout + hash is <= threshold. Vectorized row-at-a-time over the (smaller) holdout side.""" + cand = np.asarray(cand_hashes, dtype=np.uint64) + hold = np.asarray(hold_hashes, dtype=np.uint64) + mask = np.zeros(len(cand), dtype=bool) + if len(cand) == 0 or len(hold) == 0: + return mask + for h in hold: + d = _popcount64(cand ^ h) + mask |= d <= threshold + return mask + + +def dedup_manifest(df, threshold=6, priority=None, phash_workers=8): + """Collapse near-duplicate rows (hamming(phash) <= threshold) to one survivor + per cluster, keeping the row whose source ranks best in `priority`. + + Complexity: O(n^2) comparisons but VECTORIZED (numpy SWAR popcount, row-at-a-time), so the + constant factor is ~50x smaller than a Python-per-pair loop — 30k rows in ~10s vs minutes, + 100k in ~2min. (LSH banding was tried and REJECTED: on mostly-unique real-image hashes, + narrow bands collide by chance and generate millions of spurious candidates, making it + SLOWER than the vectorized scan while adding tuning + false-negative risk.) Output is + identical to an exhaustive scan. + """ + priority = list(priority) if priority is not None else DEFAULT_PRIORITY + rank = {src: i for i, src in enumerate(priority)} + n_sources = len(priority) + + work = df.copy() + work["_phash"] = _phash_many(work["path"].tolist(), workers=phash_workers) + work["_rank"] = work["source_dataset"].map(lambda s: rank.get(s, n_sources)) + + n = len(work) + hashes = work["_phash"].to_numpy() + hashes_u64 = np.asarray([int(h) & ((1 << 64) - 1) for h in hashes], dtype=np.uint64) + ranks = work["_rank"].to_numpy() + dropped = np.zeros(n, dtype=bool) + + # Gather within-threshold neighbours per row (vectorized), then the SAME union-find-free + # clustering: for each surviving row i, its within-threshold matches form a cluster; keep + # the best-priority survivor, drop the rest. Ties broken by stable first-seen (ascending i). + from collections import defaultdict + neighbors = defaultdict(list) + for i, j in _within_threshold_neighbors(hashes_u64, threshold): + neighbors[i].append(j) + for i in range(n): + if dropped[i]: + continue + cluster = [i] + for j in neighbors.get(i, ()): + if not dropped[j]: + cluster.append(j) + if len(cluster) == 1: + continue + best = min(cluster, key=lambda idx: (ranks[idx], idx)) + for idx in cluster: + if idx != best: + dropped[idx] = True + + kept_mask = ~dropped + cols = [c for c in MANIFEST_COLUMNS if c in df.columns] + deduped_df = df.loc[kept_mask, cols].reset_index(drop=True) + # Persist the survivor's OWN pHash so downstream cross-artifact leak subtraction + # (candidate vs official holdout) is manifest-only — no re-reading 45k images. + deduped_df["phash"] = hashes_u64[kept_mask] + + # --- provenance table --------------------------------------------------- + sources_present = list(dict.fromkeys(df["source_dataset"])) # first-seen order + sampled = df["source_dataset"].value_counts() + after = df.loc[kept_mask, "source_dataset"].value_counts() + + rows = [] + for src in sources_present: + s = int(sampled.get(src, 0)) + a = int(after.get(src, 0)) + rows.append({"source_dataset": src, "total_available": pd.NA, + "sampled": s, "after_dedup": a, "dropped_as_dupe": s - a}) + + # Sort by priority order (unlisted sources after all listed, in first-seen order). + rows.sort(key=lambda r: rank.get(r["source_dataset"], n_sources)) + + total = {"source_dataset": "TOTAL", "total_available": pd.NA, + "sampled": sum(r["sampled"] for r in rows), + "after_dedup": sum(r["after_dedup"] for r in rows), + "dropped_as_dupe": sum(r["dropped_as_dupe"] for r in rows)} + rows.append(total) + + provenance_df = pd.DataFrame( + rows, columns=["source_dataset", "total_available", "sampled", "after_dedup", "dropped_as_dupe"]) + + return deduped_df, provenance_df diff --git a/src/data/provenance/diskcheck.py b/src/data/provenance/diskcheck.py new file mode 100644 index 0000000000000000000000000000000000000000..60ee5cb558e4da2bc57c16d5ee1efe9c8f8c50b1 --- /dev/null +++ b/src/data/provenance/diskcheck.py @@ -0,0 +1,38 @@ +# src/shards/diskcheck.py +"""Multi-filesystem disk accounting. The failed corpus run's telemetry watched ONE +filesystem while the client archived on another; fs_stats observes every filesystem +that holds run data (run root, system temp) and assert_headroom +stops BEFORE an operation that cannot fit its peak + emergency reserve.""" +import os +import shutil +import tempfile + +RESERVE_BYTES = 5 * 2**30 + + +def fs_stats(paths): + out = {} + for p in paths: + p = os.path.realpath(p) + if not os.path.exists(p): + continue + dev = os.stat(p).st_dev + key = str(dev) + if key not in out: + du = shutil.disk_usage(p) + out[key] = {"paths": [], "used": du.used, "free": du.free, "total": du.total} + out[key]["paths"].append(p) + return out + + +def default_watch_paths(run_root): + paths = [run_root, tempfile.gettempdir()] + return paths + + +def assert_headroom(path, needed_bytes, reserve_bytes=RESERVE_BYTES): + free = shutil.disk_usage(os.path.realpath(path)).free + if free < needed_bytes + reserve_bytes: + raise RuntimeError( + f"disk headroom: need {needed_bytes + reserve_bytes} bytes on {path}, " + f"free {free} — stopping before the operation (state preserved)") diff --git a/src/data/provenance/families.py b/src/data/provenance/families.py new file mode 100644 index 0000000000000000000000000000000000000000..4ec625ff6c569b9c2f541497c4a992bce4602e71 --- /dev/null +++ b/src/data/provenance/families.py @@ -0,0 +1,32 @@ +# src/curator/families.py +import re + +FAMILIES = [ + "sd1x", "sd2x", "sdxl", "sd3", "flux", "dalle", "midjourney", + "kandinsky", "gan", "other_diffusion", "commercial_other", "unknown", +] + +# Ordered rules: first match wins. (name_lower) -> family. +# Order matters: sdxl/sd3 before the generic sd1x/sd2x check; specific before generic. +_RULES = [ + (r"stable-diffusion-xl|sdxl|sd-xl", "sdxl"), + (r"stable-diffusion-3|sd-?3|sd3", "sd3"), + (r"stable-diffusion-2|sd-?2|sd2", "sd2x"), + (r"stable-diffusion-v?1|sd-?1|sd1|stable-diffusion-1", "sd1x"), + (r"flux", "flux"), + (r"dall-?e|dalle", "dalle"), + (r"midjourney|mj[_-]?v?\d", "midjourney"), + (r"kandinsky", "kandinsky"), + (r"stylegan|prousgan|progan|biggan|cyclegan|stargan|gaugan|gansformer|projectedgan|\bgan\b", "gan"), + (r"\badm\b|glide|vqdm|wukong|\bldm\b|ddpm|ddim|denoising.?diffusion", "other_diffusion"), + (r"firefly|ideogram|\bimagen\b", "commercial_other"), # imagen = Google's commercial model +] + +def classify_family(model_name: str) -> str: + if not model_name: + return "unknown" + n = model_name.lower() + for pat, fam in _RULES: + if re.search(pat, n): + return fam + return "unknown" diff --git a/src/data/provenance/feature_cache.py b/src/data/provenance/feature_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..30ed195a4c9ac4598ee2826c9abd3ba15c2c5e10 --- /dev/null +++ b/src/data/provenance/feature_cache.py @@ -0,0 +1,36 @@ +# src/curator/feature_cache.py +import hashlib, os + +def cache_path(img_id, backbone_id, resolution, cache_dir) -> str: + key = f"{img_id}|{backbone_id}|{resolution}" + h = hashlib.sha1(key.encode()).hexdigest()[:16] + return os.path.join(cache_dir, f"{h}.pt") + +def shard_indices(df, shard_idx, n_shards): + return list(range(shard_idx, len(df), n_shards)) + +def extract_shard(manifest_df, shard_idx, n_shards, cfg, cache_dir): + """GPU-host real path: load backbone, forward each shard image, save features. + Local: exercised only via the pure helpers above (this body needs a backbone).""" + import torch + from model.backbone import build_backbone + from data.preprocess import preprocess_image + from PIL import Image + os.makedirs(cache_dir, exist_ok=True) + bb = build_backbone(cfg) + if torch.cuda.is_available(): + bb = bb.cuda() + idxs = shard_indices(manifest_df, shard_idx, n_shards) + for i in idxs: + row = manifest_df.iloc[i] + cp = cache_path(row["img_id"] if "img_id" in row else row["path"], + cfg.backbone, cfg.resolution, cache_dir) + if os.path.exists(cp): + continue + with Image.open(row["path"]) as im: + x = preprocess_image(im.convert("RGB"), cfg.resolution).unsqueeze(0) + if torch.cuda.is_available(): + x = x.cuda() + with torch.no_grad(): + feat = bb(x).squeeze(0).cpu().half() + torch.save(feat, cp) diff --git a/src/data/provenance/finalize.py b/src/data/provenance/finalize.py new file mode 100644 index 0000000000000000000000000000000000000000..610852a5f03af0d674186f62950cc83ffc39d412 --- /dev/null +++ b/src/data/provenance/finalize.py @@ -0,0 +1,38 @@ +# src/curator/finalize.py +"""Cross-artifact finalization (manifest-only; no image I/O, no shard rewrites). +The official holdout always wins: finalization removes candidate TRAINING rows that +collide with holdout rows (exact identity first, then pHash<=threshold), never the +reverse. Splits are assigned only AFTER subtraction, and the finalized manifest pins +the immutable candidate dataset version — 'latest' resolution is refused.""" +from data.provenance.dedup_audit import cross_leak_mask +from data.provenance.splitter import assign_mixed_split + + +def _ident(df): + col = "img_id" if "img_id" in df.columns else "rel_path" + return df[col].astype(str) + + +def finalize_training_manifest(cand_df, holdout_df, candidate_ref, threshold=6, + val_frac=0.15, seed=0): + if "@" not in candidate_ref: + raise ValueError( + f"candidate_ref must pin a version (slug@version), got {candidate_ref!r}") + cand = cand_df.copy() + n0 = len(cand) + hold_ids = set(_ident(holdout_df)) if len(holdout_df) else set() + cand = cand[~_ident(cand).isin(hold_ids)] + n_exact = n0 - len(cand) + if len(holdout_df) and len(cand): + mask = cross_leak_mask(cand["phash"].to_numpy(), + holdout_df["phash"].to_numpy(), threshold=threshold) + cand = cand[~mask] + n_phash = n0 - n_exact - len(cand) + final = assign_mixed_split(cand.reset_index(drop=True), val_frac=val_frac, seed=seed) + for part in ("train", "val"): + labs = set(final[final["split"] == part]["label"]) + assert labs == {0, 1}, f"finalize: split {part} lacks both classes ({labs})" + final["candidate_dataset"] = candidate_ref + report = {"n_exact_dropped": int(n_exact), "n_phash_dropped": int(n_phash), + "n_final": int(len(final))} + return final, report diff --git a/src/data/provenance/journal.py b/src/data/provenance/journal.py new file mode 100644 index 0000000000000000000000000000000000000000..aba2072f3718ef6f8145297f4c3e7d476ea05c6f --- /dev/null +++ b/src/data/provenance/journal.py @@ -0,0 +1,48 @@ +# src/shards/journal.py +"""Append-only shard progress journal. States are monotonic; deletion decisions +downstream depend on the DURABLY persisted state (fsync per append), never on +in-memory state. A compact parquet snapshot can be regenerated at any time.""" +import json +import os +import time + +STATES = ("planned", "writing", "verified", "indexed", "raw_cleaned", "remote_confirmed") +_ORD = {s: i for i, s in enumerate(STATES)} + + +class Journal: + def __init__(self, run_root): + self.path = os.path.join(run_root, "journal.jsonl") + self._latest = {} + if os.path.exists(self.path): + with open(self.path, encoding="utf-8") as f: + for line in f: + if line.strip(): + rec = json.loads(line) + self._latest[rec["shard_name"]] = rec + + def state(self, shard_name): + rec = self._latest.get(shard_name) + return rec["state"] if rec else None + + def load(self): + return dict(self._latest) + + def advance(self, shard_name, state, **info): + if state not in _ORD: + raise ValueError(f"unknown state {state!r}") + cur = self.state(shard_name) + if cur is not None and _ORD[state] <= _ORD[cur]: + raise ValueError(f"non-monotonic {shard_name}: {cur} -> {state}") + rec = {"shard_name": shard_name, "state": state, "ts": time.time(), **info} + prev = self._latest.get(shard_name, {}) + rec = {**{k: v for k, v in prev.items() if k not in ("state", "ts")}, **rec} + with open(self.path, "a", encoding="utf-8") as f: + f.write(json.dumps(rec) + "\n") + f.flush() + os.fsync(f.fileno()) + self._latest[shard_name] = rec + + def snapshot_parquet(self, path): + import pandas as pd + pd.DataFrame(list(self._latest.values())).to_parquet(path, index=False) diff --git a/src/data/provenance/manifest_build.py b/src/data/provenance/manifest_build.py new file mode 100644 index 0000000000000000000000000000000000000000..410533cba8380c4b662b967c1c379ddfccac6bbf --- /dev/null +++ b/src/data/provenance/manifest_build.py @@ -0,0 +1,54 @@ +# src/curator/manifest_build.py +import os +import pandas as pd +from PIL import Image +from data.provenance.families import classify_family + +MANIFEST_COLUMNS = ["path","label","source_dataset","generator_id","family", + "real_source","native_w","native_h","jpeg_bytes_per_pixel", + "ref_type","mount_slug","rel_path","split"] + +def jpeg_bytes_per_pixel(path) -> float: + try: + size = os.path.getsize(path) + with Image.open(path) as im: + w, h = im.size + return float(size) / float(w * h) if w and h else float("nan") + except Exception: + return float("nan") + + +_EXT = {"JPEG": ".jpg", "PNG": ".png", "WEBP": ".webp", "BMP": ".bmp", "GIF": ".gif", "TIFF": ".tiff"} + +def save_source_image(raw: bytes, path_noext: str): + """Persist a source image blob VERBATIM (no re-encode) and return (path, w, h). + + Re-encoding is the wrong move for this corpus: JPEG bakes a scored corruption into + the 'clean' ground truth; PNG re-encode of already-JPEG sources ~4x's disk (blows the + typical working-disk limits) for ZERO fidelity gain — the JPEG artifacts are already in + the source pixels. Writing the source bytes unchanged adds nothing and preserves the + original size. We sniff format+dims via PIL (header only, no full decode) to pick the + extension and fill native_w/h; a corrupt/undecodable blob raises (caller skips it).""" + import io + with Image.open(io.BytesIO(raw)) as im: # header sniff; verify() would consume it + fmt, (w, h) = im.format, im.size + im.verify() # raises on a truncated/corrupt blob + ext = _EXT.get(fmt, ".img") + path = path_noext + ext + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as f: + f.write(raw) + return path, w, h + +def assemble_manifest(records: pd.DataFrame) -> pd.DataFrame: + df = records.copy() + df["family"] = df["generator_id"].fillna("none").map(classify_family) + # Resolution columns come from the sources (Tasks 3/4); default if a record lacks them. + if "ref_type" not in df.columns: + df["ref_type"] = "materialized" + if "mount_slug" not in df.columns: + df["mount_slug"] = "" + if "rel_path" not in df.columns: + df["rel_path"] = df["path"].map(lambda p: os.path.basename(str(p))) + df["split"] = "" + return df[MANIFEST_COLUMNS] diff --git a/src/data/provenance/orchestrate.py b/src/data/provenance/orchestrate.py new file mode 100644 index 0000000000000000000000000000000000000000..896efbd03b0ec2e9d91af5a3e99fdfb360b8d912 --- /dev/null +++ b/src/data/provenance/orchestrate.py @@ -0,0 +1,122 @@ +# src/shards/orchestrate.py +"""Per-shard transactions: plan -> write+verify -> journal -> delete raws. Raw files +for a shard are deleted ONLY after the journal durably records the verified shard, +so conversion never holds a second corpus-sized copy and a crash at any point leaves +either the raws or a verified shard (never neither).""" +import hashlib +import os +import shutil +import zipfile + +import pandas as pd + +from data.provenance import diskcheck +from data.provenance.writer import plan_shards, write_shard + + +def _sha256_file(path): + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def resume_pending(journal, shard_dir, raw_root): + """Startup/retry recovery: drop stale .partial files, and finish raw deletion for + shards journaled verified/indexed but not raw_cleaned. A digest mismatch on a + completed shard is a hard failure.""" + for f in os.listdir(shard_dir): + if f.endswith(".partial"): + os.remove(os.path.join(shard_dir, f)) + for name, rec in journal.load().items(): + if rec["state"] in ("verified", "indexed"): + path = os.path.join(shard_dir, name) + if _sha256_file(path) != rec["sha256"]: + raise AssertionError(f"digest mismatch on completed shard {name}") + _delete_members(raw_root, rec.get("members", [])) + if rec["state"] == "verified": + journal.advance(name, "indexed") + journal.advance(name, "raw_cleaned") + + +def _delete_members(raw_root, member_names): + freed = 0 + for rel in member_names: + p = os.path.join(raw_root, rel.replace("/", os.sep)) + if os.path.exists(p): + freed += os.path.getsize(p) + os.remove(p) + # prune empty dirs bottom-up + for dp, _dirs, _files in os.walk(raw_root, topdown=False): + if dp != raw_root and not os.listdir(dp): + os.rmdir(dp) + return freed + + +def shard_corpus(manifest_df, raw_root, shard_dir, journal, target_bytes=1 << 30, + prefix="corpus", log=print): + """Shard every manifest row (transactionally, resumable via the journal) and + return (candidates_df, shards_df). candidates_df drops `path` (raws are deleted) + and carries shard_name/member_name/byte_size/crc32 with ref_type='shard'.""" + df = manifest_df.copy() + if "byte_size" not in df.columns: + # Resume path: raws of completed shards are deleted — recover their sizes from + # the verified shards' central directories so packing stays deterministic. + sizes = {} + if os.path.isdir(shard_dir): + for f in os.listdir(shard_dir): + if f.endswith(".shard"): + with zipfile.ZipFile(os.path.join(shard_dir, f)) as zf: + for zi in zf.infolist(): + sizes[zi.filename] = zi.file_size + def _size(row): + if os.path.exists(row["path"]): + return os.path.getsize(row["path"]) + if row["rel_path"] in sizes: + return sizes[row["rel_path"]] + raise FileNotFoundError( + f"no raw file and no shard member for {row['rel_path']}") + df["byte_size"] = df.apply(_size, axis=1) + plan = plan_shards(df, target_bytes=target_bytes, prefix=prefix) + shard_records = [] + for name, group in plan.groupby("shard_name", sort=True): + state = journal.state(name) + if state == "raw_cleaned": + # resumed past this shard: recover per-member CRCs + ledger row from the + # verified archive itself (raws are gone; the shard is the source of truth) + with zipfile.ZipFile(os.path.join(shard_dir, name)) as zf: + crc = {zi.filename: zi.CRC for zi in zf.infolist()} + plan.loc[group.index, "crc32"] = group["member_name"].map(crc) + rec0 = journal.load()[name] + shard_records.append({"shard_name": name, "n_members": len(crc), + "payload_bytes": int(group["byte_size"].sum()), + "zip_bytes": os.path.getsize(os.path.join(shard_dir, name)), + "sha256": rec0["sha256"]}) + continue + members = list(zip(group["member_name"], group["path"])) + diskcheck.assert_headroom(shard_dir, needed_bytes=int(group["byte_size"].sum())) + if state is None: + journal.advance(name, "planned", n_members=len(members)) + journal.advance(name, "writing") + rec = write_shard(os.path.join(shard_dir, name), members) + journal.advance(name, "verified", sha256=rec["sha256"], + members=[m["member_name"] for m in rec["members"]]) + journal.advance(name, "indexed") + freed = _delete_members(raw_root, [m for m, _ in members]) + free_after = shutil.disk_usage(shard_dir).free + journal.advance(name, "raw_cleaned", bytes_reclaimed=freed, + free_after=free_after) + log(f"[shard] {name}: {rec['n_members']} members, " + f"{rec['zip_bytes']} zip bytes, freed {freed}, free {free_after}") + crc = {m["member_name"]: m["crc32"] for m in rec["members"]} + plan.loc[group.index, "crc32"] = group["member_name"].map(crc) + shard_records.append({"shard_name": name, "n_members": rec["n_members"], + "payload_bytes": rec["payload_bytes"], + "zip_bytes": rec["zip_bytes"], "sha256": rec["sha256"]}) + cand = plan.drop(columns=["path"]).copy() + cand["ref_type"] = "shard" + cand["crc32"] = cand["crc32"].astype("int64") + return cand.reset_index(drop=True), pd.DataFrame( + shard_records, columns=["shard_name", "n_members", "payload_bytes", + "zip_bytes", "sha256"]) diff --git a/src/data/provenance/push.py b/src/data/provenance/push.py new file mode 100644 index 0000000000000000000000000000000000000000..683124c5f41d38467b06ce60221255c4b22bedb1 --- /dev/null +++ b/src/data/provenance/push.py @@ -0,0 +1,11 @@ +# src/data/provenance/push.py +"""Manifest write-out. The corpus is built and consumed locally; there is no +remote publishing step.""" +import os + + +def write_manifest(df, out_dir) -> str: + os.makedirs(out_dir, exist_ok=True) + path = os.path.join(out_dir, "manifest.parquet") + df.to_parquet(path, index=False) + return path diff --git a/src/data/provenance/rangezip.py b/src/data/provenance/rangezip.py new file mode 100644 index 0000000000000000000000000000000000000000..a8c6ca1e3f758487c0e9653cef407888b0769732 --- /dev/null +++ b/src/data/provenance/rangezip.py @@ -0,0 +1,266 @@ +# src/curator/rangezip.py +"""Random-access reader for a REMOTE zip over HTTP byte ranges, using ONLY explicit +ranges (bytes=a-b) — never suffix ranges (bytes=-N), which ModelScope's file endpoint +rejects with 400 (remotezip uses suffix ranges, hence its failure here). + +Why this works at all: a ZIP has a Central Directory at the END mapping every member +to a byte offset, and each member is compressed INDEPENDENTLY. So to pull one image out +of a 25 GB zip we: (1) range-read the tail to find + parse the central directory, then +(2) range-read just that member's local-header+data and inflate it. Transfer is ~the +member size, not the whole archive. + +Only the subset of ZIP needed for WildFake is implemented: STORED (0) + DEFLATE (8), +32-bit offsets, and ZIP64 end-of-central-directory when present. No encryption. +""" +import struct +import zlib +import urllib.request + +_EOCD_SIG = b"PK\x05\x06" +_EOCD64_LOC_SIG = b"PK\x06\x07" +_EOCD64_SIG = b"PK\x06\x06" +_CEN_SIG = b"PK\x01\x02" + + +class _ShortBuffer(Exception): + """Raised by _decode_at when the fetched buffer doesn't hold the full member; carries + the exact byte span now known so the caller re-fetches precisely (no truncation).""" + def __init__(self, needed): + self.needed = needed + super().__init__(f"buffer short, need {needed} bytes") + + +class RangeZip: + def __init__(self, url, headers=None, session=None): + self.url = url + self.headers = headers or {} + # CONNECTION REUSE: one persistent keep-alive session for ALL range reads to this + # zip's host, so the TLS handshake to ModelScope (China, ~2s) is paid ONCE instead + # of per request. This is the dominant cost — reusing the SSL/TCP connection is what + # makes many range reads cheap. Falls back to urllib (still per-request) if requests + # is unavailable. + self._session = session if session is not None else self._make_session() + self.size = self._content_length() + self._entries = self._read_central_directory() # name -> (offset, comp_size, method) + + @staticmethod + def _make_session(): + try: + import requests + from requests.adapters import HTTPAdapter + s = requests.Session() + # a small pool of persistent connections, reused across range GETs (keep-alive). + ad = HTTPAdapter(pool_connections=8, pool_maxsize=32, max_retries=3) + s.mount("https://", ad); s.mount("http://", ad) + return s + except Exception: + return None # urllib fallback (no pooling) + + # --- HTTP helpers (explicit ranges only) -------------------------------- + def _content_length(self): + if self._session is not None: + r = self._session.head(self.url, headers=self.headers, allow_redirects=True, timeout=60) + r.raise_for_status() + return int(r.headers["Content-Length"]) + req = urllib.request.Request(self.url, method="HEAD", headers=self.headers) + with urllib.request.urlopen(req) as r: + return int(r.headers["Content-Length"]) + + def _get_range(self, lo, hi): + """Fetch bytes [lo, hi] inclusive via an explicit Range header, over the persistent + keep-alive session (reused SSL/TCP connection) when available.""" + h = dict(self.headers); h["Range"] = f"bytes={lo}-{hi}" + if self._session is not None: + r = self._session.get(self.url, headers=h, timeout=300) + r.raise_for_status() + return r.content + req = urllib.request.Request(self.url, headers=h) + with urllib.request.urlopen(req) as r: + return r.read() + + # Tail window read to bootstrap the directory. Sized so the CENTRAL DIRECTORY of most + # zips fits inside it (WildFake: reals 3-17MB, typical fakes 5-16MB), letting us slice + # the directory straight out of this one read — no second GET. Only the giant + # 200k+-member zips (CD 20-52MB) need a follow-up, and then only for the missing prefix. + _TAIL = 16 * 1024 * 1024 + + # --- central directory -------------------------------------------------- + def _read_central_directory(self): + tail_len = min(self._TAIL, self.size) + tail_start = self.size - tail_len + tail = self._get_range(tail_start, self.size - 1) # ONE read, big window + i = tail.rfind(_EOCD_SIG) + if i < 0: + raise ValueError("EOCD not found (not a zip or comment too large)") + eocd = tail[i:i + 22] + cd_size = struct.unpack(" find the ZIP64 EOCD via its locator. + if cd_off == 0xFFFFFFFF or cd_size == 0xFFFFFFFF or total == 0xFFFF: + j = tail.rfind(_EOCD64_LOC_SIG, 0, i) + if j < 0: + raise ValueError("ZIP64 locator not found") + eocd64_off = struct.unpack("= tail_start: + head = tail[eocd64_off - tail_start: eocd64_off - tail_start + 56] + else: + head = self._get_range(eocd64_off, eocd64_off + 56 - 1) + if head[:4] != _EOCD64_SIG: + raise ValueError("bad ZIP64 EOCD signature") + cd_size = struct.unpack("= tail_start: + cd = tail[cd_off - tail_start: cd_off - tail_start + cd_size] + else: + missing = self._get_range(cd_off, tail_start - 1) # only the part below the tail + cd = missing + tail[:max(0, (cd_off + cd_size) - tail_start)] + return self._parse_central(cd) + + def _parse_central(self, cd): + entries, p = {}, 0 + while p + 4 <= len(cd) and cd[p:p + 4] == _CEN_SIG: + method = struct.unpack(" len(buf): + raise _ShortBuffer(30) # not even the fixed header present + n = struct.unpack(" len(buf): + raise _ShortBuffer(30 + n + m + comp_size) # exact span now known + raw = buf[start:end] + if method == 0: + return raw + if method == 8: + return zlib.decompress(raw, -15) + raise ValueError(f"unsupported compression method {method}") + + def read(self, name): + """DECOMPRESSED bytes of one member via a range fetch of its local header + data. + Retries once with the exact span if the first (slack-estimated) window was short.""" + lho, comp_size, method = self._entries[name] + span = self._LH_SLACK + comp_size + window = self._get_range(lho, min(lho + span, self.size) - 1) + try: + return self._decode_at(window, lho, lho, comp_size, method) + except _ShortBuffer as e: + window = self._get_range(lho, min(lho + e.needed, self.size) - 1) + return self._decode_at(window, lho, lho, comp_size, method) + + def read_many(self, names, max_gap=1 << 20, workers=16): + """Fetch MANY members by COALESCING nearby members into contiguous segments, then + fetching those segments IN PARALLEL over the shared keep-alive session. + + The parallelism is at the SEGMENT level (not the zip level): a single giant WildFake + zip has hundreds of segments, and with only ~12 zips, per-zip concurrency alone + barely distributes — one big zip would serialize. Running its segment GETs across + `workers` threads (all reusing one pooled SSL connection) is what actually saturates + the link. Returns {name: decompressed_bytes}. + """ + from concurrent.futures import ThreadPoolExecutor + # each member spans [lho, lho + LH_SLACK + comp_size); sort by start offset + spans = [] + for nm in names: + lho, comp_size, method = self._entries[nm] + end = lho + self._LH_SLACK + comp_size + spans.append((lho, end, nm, comp_size, method)) + spans.sort() + + # 1) build the list of coalesced segments (each = one contiguous range + its members) + segments = [] + i = 0 + while i < len(spans): + seg_lo, seg_hi = spans[i][0], spans[i][1] + members = [spans[i]] + j = i + 1 + while j < len(spans) and spans[j][0] <= seg_hi + max_gap: + seg_hi = max(seg_hi, spans[j][1]) + members.append(spans[j]) + j += 1 + segments.append((seg_lo, min(seg_hi, self.size), members)) + i = j + + # 2) fetch + decode each segment; segments run in PARALLEL over the pooled session + out, lock = {}, __import__("threading").Lock() + def _do(seg): + seg_lo, seg_hi, members = seg + buf = self._get_range(seg_lo, seg_hi - 1) # ONE GET per segment + local = {} + for lho, _end, nm, comp_size, method in members: + try: + local[nm] = self._decode_at(buf, seg_lo, lho, comp_size, method) + except _ShortBuffer: + local[nm] = self.read(nm) + with lock: + out.update(local) + if workers > 1 and len(segments) > 1: + with ThreadPoolExecutor(max_workers=workers) as ex: + list(ex.map(_do, segments)) + else: + for seg in segments: + _do(seg) + return out + + +def _zip64_patch(extra, comp_size, lho): + p = 0 + while p + 4 <= len(extra): + hid, sz = struct.unpack("= 3: + comp_size, lho = vals[1], vals[2] + elif len(vals) == 2: + comp_size, lho = vals[0], vals[1] + elif len(vals) == 1: + lho = vals[0] + break + p += 4 + sz + return comp_size, lho diff --git a/src/data/provenance/resolver.py b/src/data/provenance/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..431bd2e4c46dea9f001743d4898eb44c15e60583 --- /dev/null +++ b/src/data/provenance/resolver.py @@ -0,0 +1,53 @@ +# src/curator/resolver.py +import os +from PIL import Image + +# Optional per-slug marker sub-path used to DISCOVER a mounted dataset's root when +# Mount layouts are non-uniform across hosts. The corpus is ALL-MATERIALIZED from +# official sources (commercial-safe), so no default mount markers are configured; +# the mechanism is retained for optional reference-mounting of any future +# commercial-safe corpus directory (pass `markers={slug: marker}` to SourceResolver). +MOUNT_MARKERS: dict[str, str] = {} + +class SourceResolver: + def __init__(self, input_root="./data/raw", fakes_slug="track5-corpus-fakes", markers=None): + self.input_root = input_root + self.fakes_slug = fakes_slug + self.markers = markers if markers is not None else MOUNT_MARKERS + self._root_cache = {} + + def discover_root(self, slug) -> str: + if slug in self._root_cache: + return self._root_cache[slug] + marker = self.markers[slug] + # find a dir whose path ends with the marker; its PARENT-of-marker is the root + marker_norm = marker.replace("\\", "/") + for dp, dirs, _ in os.walk(self.input_root): + rel = os.path.relpath(dp, self.input_root) + rel_norm = rel.replace("\\", "/") + # Anchor to a path-component boundary so marker "dataset" does NOT match + # ".../coco-2017-dataset" (no "/dataset" segment) while ".../lhq-1024/dataset" does. + if rel_norm == marker_norm or rel_norm.endswith("/" + marker_norm): + # root = dp with the marker suffix stripped, so root/rel_path (which + # includes the marker prefix) resolves + depth = len(marker.replace("\\", "/").split("/")) + root = dp + for _ in range(depth): + root = os.path.dirname(root) + self._root_cache[slug] = root + return root + raise FileNotFoundError(f"resolver: could not discover mount root for {slug} (marker {marker}) under {self.input_root}") + + def _path(self, row) -> str: + if row["ref_type"] == "mount": + return os.path.join(self.discover_root(row["mount_slug"]), row["rel_path"]) + if row["ref_type"] == "materialized": + return os.path.join(self.input_root, self.fakes_slug, row["rel_path"]) + raise ValueError(f"unknown ref_type {row['ref_type']}") + + def resolve(self, row) -> Image.Image: + p = self._path(row) + if not os.path.exists(p): + raise FileNotFoundError(f"resolver: {p} (ref_type={row['ref_type']}, slug={row.get('mount_slug')})") + with Image.open(p) as im: + return im.convert("RGB").copy() diff --git a/src/data/provenance/run_curator.py b/src/data/provenance/run_curator.py new file mode 100644 index 0000000000000000000000000000000000000000..bec84885b61cc737b56811d12bb12e02a41bc4d3 --- /dev/null +++ b/src/data/provenance/run_curator.py @@ -0,0 +1,110 @@ +# src/curator/run_curator.py +import argparse, os, tempfile +import pandas as pd +from data.provenance import sources, subsample, manifest_build, audit, push, dedup_audit +from data.provenance.splitter import assign_mixed_split + + +def _write_stage_manifest(frame, out_dir, stage): + path = os.path.join(out_dir, f"manifest_{stage}.parquet") + frame.to_parquet(path, index=False) + counts = (frame.groupby(["label", "source_dataset", "family"], dropna=False) + .size().rename("count").reset_index()) + counts.insert(0, "stage", stage) + return counts + +def _write_build_metadata(tables_dir, **config): + import json, subprocess, time, uuid + try: + commit = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, + text=True).stdout.strip() + except Exception: # noqa: BLE001 + commit = "unknown" + meta = {"code_commit": commit or "unknown", "build_id": uuid.uuid4().hex, + "built_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "schema_version": 1, "license_policy": "commercial-only-2026-08-30", + "config": config} + with open(os.path.join(tables_dir, "build_metadata.json"), "w") as f: + json.dump(meta, f, indent=1) + + +def run_curator(per_family_cap=30000, seed=0, out_dir=None, input_root="./data/raw", + do_push=False, allow_shortcut=False, corpus_sources=None, max_examine=None, + progress_cb=None, sharded=False, run_root=None): + out_dir = out_dir or tempfile.mkdtemp(prefix="corpus_") + # Corpus is ALL-MATERIALIZED from official commercial-safe sources (JPEG q95) into + # a versioned corpus directory. Sources emit BOTH real (label 0) and fake (label 1) + # rows — e.g. SID_Set/WildFake/OmniFake carry reals; CommForensics is fake-only. + fakes_out_dir = os.path.join(out_dir, "track5-corpus-fakes") + if corpus_sources is None: + corpus_sources = sources.default_corpus_sources(per_family_cap) + if sharded: + # Sharded tail (spec: sharded-corpus-pipeline): candidate corpus only — NO split + # (the finalizer assigns splits after official-holdout subtraction), deterministic + # ZIP shards + audit tables in a flat publish dir, raws deleted per verified shard. + from data.provenance import cleanup as _cleanup, diskcheck, journal as _journal + from data.provenance.orchestrate import shard_corpus + root = _cleanup.make_run_root(run_root or tempfile.gettempdir()) + tables_dir = os.path.join(root, "publish") + os.makedirs(tables_dir, exist_ok=True) + print("[disk]", diskcheck.fs_stats(diskcheck.default_watch_paths(root)), flush=True) + else: + tables_dir = out_dir + raw = sources.pull_all(corpus_sources, fakes_out_dir, per_family_cap, seed, + max_examine=max_examine, progress_cb=progress_cb) + # --- pure logic --- + # assemble_manifest FIRST so `family` (and split/ref columns) exist before + # fake-family capping, which must never touch explicitly licensed real quotas. + mani = manifest_build.assemble_manifest(raw) + stage_counts = [_write_stage_manifest(mani, tables_dir, "raw")] + mani = subsample.cap_fake_families(mani, per_family_cap, seed) + stage_counts.append(_write_stage_manifest(mani, tables_dir, "capped")) + # Dedup AFTER cap (bounds the O(n^2) pHash scan to the capped set) but BEFORE + # splitting. OmniFake overlaps + # WildFake/GenImage; near-identical images collapse to one, priority-ordered. + # provenance = the {sampled, after_dedup, dropped_as_dupe} attribution deliverable. + mani, provenance = dedup_audit.dedup_manifest(mani) + stage_counts.append(_write_stage_manifest(mani, tables_dir, "deduped")) + pd.concat(stage_counts, ignore_index=True).to_parquet( + os.path.join(tables_dir, "stage_counts.parquet"), index=False) + prov_path = os.path.join(tables_dir, "provenance.parquet") + provenance.to_parquet(prov_path, index=False) + if sharded: + # GC checkpoint: cap/dedup rejects are unreferenced now — reclaim before sharding + from data.provenance.cleanup import reclaim_unreferenced + n_gc, b_gc = reclaim_unreferenced(fakes_out_dir, set(mani["rel_path"])) + print(f"[gc] post-dedup reclaim: {n_gc} files, {b_gc} bytes", flush=True) + j = _journal.Journal(root) + mani_ns = mani.drop(columns=["split"], errors="ignore") + cand, shards_df = shard_corpus(mani_ns, fakes_out_dir, tables_dir, j) + cand.to_parquet(os.path.join(tables_dir, "manifest_candidates.parquet"), index=False) + shards_df.to_parquet(os.path.join(tables_dir, "shards.parquet"), index=False) + _write_build_metadata(tables_dir, per_family_cap=per_family_cap, seed=seed) + print("[disk]", diskcheck.fs_stats(diskcheck.default_watch_paths(root)), flush=True) + return tables_dir + mani = assign_mixed_split(mani, val_frac=0.15, seed=seed) + # --- audit gate --- + rep = audit.audit_corpus(mani) + print(rep["report"]); print("balance:", rep["class_balance"]); print("families:", rep["per_family_counts"]) + if rep["flag"] and not allow_shortcut: + raise SystemExit("DDA shortcut flagged — mitigate and re-run, or pass --allow-shortcut.") + path = push.write_manifest(mani, out_dir) + print("manifest:", path) + print("provenance:", prov_path); print(provenance.to_string(index=False)) + return out_dir + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--per-family-cap", type=int, default=30000) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out-dir", default=None) + ap.add_argument("--push", action="store_true") + ap.add_argument("--allow-shortcut", action="store_true") + ap.add_argument("--max-examine", type=int, default=None, + help="cap examples streamed per source (bandwidth guard)") + a = ap.parse_args() + run_curator(per_family_cap=a.per_family_cap, seed=a.seed, out_dir=a.out_dir, + do_push=a.push, allow_shortcut=a.allow_shortcut, max_examine=a.max_examine) + +if __name__ == "__main__": + main() diff --git a/src/data/provenance/selection.py b/src/data/provenance/selection.py new file mode 100644 index 0000000000000000000000000000000000000000..0a9bb452db0ac6144813b8c3c9c5c6617d899a97 --- /dev/null +++ b/src/data/provenance/selection.py @@ -0,0 +1,52 @@ +# src/curator/selection.py +"""The DECISION between the two passes: given a scan inventory (what we HAVE), compute +a selection plan (what we'll PULL). Pure logic on the scan DataFrame — no I/O — so the +sampling policy is reviewable before any image is touched. + +inventory columns: [source_dataset, group_key, label, count] (from DatasetSource.scan) +plan: {(source_dataset, group_key, label): quota} (consumed by DatasetSource.pull) +""" + + +def plan_selection(inventory, per_group=None, balance=False): + """Cap each (source, group, label) at min(available, per_group); optionally balance + classes by trimming the majority class to the minority total. + + per_group=None -> take everything available (quota = count). balance=True -> after + capping, scale the larger class's quotas down proportionally so n_fake == n_real + (as evenly as possible), preserving per-group proportions.""" + plan = {} + for r in inventory.itertuples(): + avail = int(r.count) + quota = avail if per_group is None else min(avail, per_group) + if quota > 0: + plan[(r.source_dataset, r.group_key, int(r.label))] = quota + + if not balance: + return plan + + n = {0: 0, 1: 0} + for (_s, _g, lbl), q in plan.items(): + n[lbl] += q + if n[0] == 0 or n[1] == 0: + return plan # can't balance with a missing class; leave as-is + keep, drop = (0, 1) if n[0] < n[1] else (1, 0) + target = n[keep] # trim the majority class down to the minority total + # Distribute `target` across the majority groups proportionally, largest-remainder. + maj = {k: v for k, v in plan.items() if k[2] == drop} + total = sum(maj.values()) + scaled, rema = {}, [] + for k, v in maj.items(): + exact = v * target / total + base = int(exact) + scaled[k] = base + rema.append((exact - base, k)) + short = target - sum(scaled.values()) + for _frac, k in sorted(rema, reverse=True)[:short]: + scaled[k] += 1 + for k, v in scaled.items(): + if v > 0: + plan[k] = v + else: + plan.pop(k, None) # a group scaled to 0 drops out + return plan diff --git a/src/data/provenance/sources.py b/src/data/provenance/sources.py new file mode 100644 index 0000000000000000000000000000000000000000..8aa65cbd393344ce0c84a8ba91fb0ccb7d79378a --- /dev/null +++ b/src/data/provenance/sources.py @@ -0,0 +1,502 @@ +"""Per-dataset source abstraction. Each dataset = a DatasetSource subclass; the +base class owns the shared record-building + per-family cap loop (DRY). Heavy +imports (datasets, huggingface_hub) are deferred inside _iter_examples so importing +this module stays cheap and pytest-collection-safe (no downloads at import).""" +import abc, os +import pandas as pd +from data.provenance.families import classify_family +from data.provenance.manifest_build import jpeg_bytes_per_pixel, save_source_image + +RECORD_COLUMNS = ["path","label","source_dataset","generator_id", + "real_source","native_w","native_h","jpeg_bytes_per_pixel", + "ref_type","mount_slug","rel_path"] + +# --- Defensive field readers (shared) --------------------------------------- +# The commercial-safe sources have heterogeneous / under-documented schemas +# (ModelScope cards are bare, OmniFake's viewer is broken). These readers coerce +# a raw example into (label, generator, PIL image) across the likely field names, +# so subclasses' _to_fields stay small and testable with hand-built dicts. + +def _as_pil(obj): + """Coerce a PIL image / bytes / {'bytes'|'path'} dict / path-string -> PIL.Image.""" + import io + from PIL import Image + if isinstance(obj, Image.Image): + return obj + if isinstance(obj, (bytes, bytearray)): + return Image.open(io.BytesIO(bytes(obj))) + if isinstance(obj, dict): + if obj.get("bytes"): + return Image.open(io.BytesIO(obj["bytes"])) + if obj.get("path"): + return Image.open(obj["path"]) + if isinstance(obj, str): + return Image.open(obj) + raise TypeError(f"_as_pil: cannot coerce {type(obj)!r} to a PIL image") + +def _read_label(ex, default=None): + """Best-effort 0(real)/1(fake) from an example. Returns None (sentinel) when NO + label field is resolvable — callers MUST skip such rows, never default them to + fake (defaulting a real image into the fake class poisons the real/fake boundary).""" + for k in ("label", "labels", "target", "is_fake", "fake"): + if ex.get(k) is not None: + v = ex[k] + if isinstance(v, str): + s = v.strip().lower() + if s in ("real", "0", "authentic", "genuine"): return 0 + if s in ("fake", "1", "synthetic", "ai", "generated"): return 1 + try: + return 1 if int(v) != 0 else 0 + except (TypeError, ValueError): + pass + for k in ("category", "class", "folder", "subset", "type"): + v = ex.get(k) + if isinstance(v, str): + return 0 if v.strip().lower() in ("real", "0_real", "authentic") else 1 + return default + +def _read_generator(ex): + """Best-effort generator/model name string (feeds classify_family).""" + for k in ("generator", "model", "model_name", "gen_model", "method", + "source_model", "generator_name", "label_name", "category", + "class", "folder"): + v = ex.get(k) + if isinstance(v, str) and v: + return v + return "" + +def _read_image(ex): + """Return the raw image object from an example (caller passes to _as_pil).""" + for k in ("image", "img", "image_data", "image_bytes", "picture"): + if ex.get(k) is not None: + return ex[k] + raise KeyError("_read_image: no image field in example") + +class DatasetSource(abc.ABC): + name: str = "source" + + @abc.abstractmethod + def _iter_examples(self): + """Yield raw examples (network I/O lives here, in subclasses).""" + @abc.abstractmethod + def _to_fields(self, example) -> dict: + """Map one example -> {label, generator_id, real_source, image: PIL.Image}.""" + + def _group_key(self, f) -> str: + """The split/selection group for a materialized-field dict: family for fakes + (label 1), real_source for reals (label 0). This is the unit scan counts and + plan_selection quotas over.""" + return classify_family(f["generator_id"]) if f["label"] == 1 else f["real_source"] + + def scan(self, max_examine=None) -> pd.DataFrame: + """PASS 1 — inventory only, NO images saved. Returns a DataFrame with columns + [source_dataset, group_key, label, count]: how many examples this source has per + (group, label). Default implementation walks _iter_examples()/_to_fields() but + decodes NO pixels beyond what _to_fields needs; subclasses with cheap metadata + (e.g. WildFake's label CSVs) should OVERRIDE scan() to avoid touching images. + max_examine bounds examined examples for streaming sources.""" + counts = {} + for n_seen, ex in enumerate(self._iter_examples()): + if max_examine is not None and n_seen >= max_examine: + break + f = self._to_fields(ex) + if f is None: + continue + key = (self._group_key(f), int(f["label"])) + counts[key] = counts.get(key, 0) + 1 + rows = [{"source_dataset": self.name, "group_key": g, "label": l, "count": c} + for (g, l), c in counts.items()] + return pd.DataFrame(rows, columns=["source_dataset", "group_key", "label", "count"]) + + def pull(self, out_dir, seed=0, per_family_cap=None, max_examine=None, + selection=None, progress_cb=None, progress_every=100) -> pd.DataFrame: + # PASS 2 — materialize. Two quota modes: + # selection (two-pass): {(source, group_key, label): quota} from plan_selection. + # per_family_cap (legacy one-pass): flat cap per family. + # max_examine bounds examples EXAMINED (streamed) — the bandwidth guard for a + # streaming source that would otherwise transfer bytes of every example. + # progress_cb({done, ...}) fires every progress_every materialized rows so a long + # network stream (e.g. CF PublicEval) is monitorable, same contract WildFake uses. + os.makedirs(out_dir, exist_ok=True) + counts, rows, n_skipped = {}, [], 0 + import time as _t + _t0 = _t.time() + for n_seen, ex in enumerate(self._iter_examples()): + if max_examine is not None and n_seen >= max_examine: + break + f = self._to_fields(ex) + if f is None: # unknown-label row: SKIP, do not + n_skipped += 1 # materialize and do not count vs cap + continue # (defaulting to fake would poison reals) + # Counter keyed by GROUP_KEY (family for fakes, real_source for reals) so + # per-group quotas don't collide — all reals share fam=="unknown", so a + # fam-keyed counter would merge distinct real sources under one quota. + gkey = self._group_key(f) + if selection is not None: + quota = selection.get((self.name, gkey, int(f["label"])), 0) + if counts.get(gkey, 0) >= quota: + continue + elif (f["label"] == 1 and per_family_cap is not None + and counts.get(gkey, 0) >= per_family_cap): + continue + elif (f["label"] == 0 and getattr(self, "real_cap", None) is not None + and counts.get(gkey, 0) >= self.real_cap): + # per-real-source cap (label 0): bounds a materialized real source (e.g. CF + # PublicEval reals) without a two-pass selection. Keyed by real_source group, + # so distinct real sources each get their own quota. None => uncapped (SID + # controls its own real count via its n_real quota, so it leaves this unset). + continue + # Namespace rel_path by source name so it round-trips through the pushed fakes + # dataset root: file lands at //fakes/x. and the resolver + # opens input_root///fakes/x.. + # PREFER VERBATIM source bytes (f["raw"], set by _to_fields when it has the + # original blob) — no re-encode, so no baked-in corruption and no PNG bloat. Only + # fall back to a lossless-PNG re-encode when a source hands us a PIL with no bytes + # (e.g. SID); JPEG re-encode is never used (it would corrupt the clean ground truth). + base = f"{self.name}/fakes/{self.name}_{gkey}_{counts.get(gkey,0)}" + p_noext = os.path.join(out_dir, base) + raw = f.get("raw") + if raw is not None: + p, w, h = save_source_image(raw, p_noext) + else: + from PIL import Image + p = p_noext + ".png" + os.makedirs(os.path.dirname(p), exist_ok=True) + f["image"].convert("RGB").save(p, format="PNG", optimize=False) + with Image.open(p) as im: + w, h = im.size + rel = os.path.relpath(p, out_dir).replace(os.sep, "/") + rows.append({"path": p, "label": f["label"], "source_dataset": self.name, + "generator_id": f["generator_id"], "real_source": f["real_source"], + "native_w": w, "native_h": h, "jpeg_bytes_per_pixel": jpeg_bytes_per_pixel(p), + "ref_type": "materialized", "mount_slug": "", "rel_path": rel}) + counts[gkey] = counts.get(gkey, 0) + 1 + if progress_cb is not None and len(rows) % progress_every == 0: + dt = max(_t.time() - _t0, 1e-9) + progress_cb({"done": len(rows), "imgs_per_s": len(rows) / dt}) + print(f"[{self.name}] skipped {n_skipped} unlabeled examples") + return pd.DataFrame(rows, columns=RECORD_COLUMNS) + +def _iter_parquet_rows(fileobj, columns, batch_size=256): + """Yield rows (dicts of `columns`) from a parquet fileobj ONE BATCH AT A TIME. + + Peak RAM is a single `batch_size`-row batch, NEVER a whole row-group. This is the + load-bearing memory bound: CommForensics PublicEval shards are packed as a SINGLE + row-group of up to 6.5GB (8640 rows of raw image blobs), so HF `datasets` + streaming (which materializes a whole row-group's Arrow table before yielding row 1) + climbs to ~30GB and OOM-kills a modest host. pyarrow.iter_batches decodes only the + requested columns and only one batch of rows at a time -> bounded, provably.""" + import pyarrow.parquet as pq + pf = pq.ParquetFile(fileobj) + for batch in pf.iter_batches(batch_size=batch_size, columns=columns): + d = batch.to_pydict() + for i in range(len(batch)): + yield {c: d[c][i] for c in columns} + + +class CommunityForensicsSource(DatasetSource): + name = "commforensics" + # LICENSE-MANDATORY (docs/specs/DATASET_LICENSES.md): materialize ONLY the + # redistribution-safe PublicEval split (schema: image_data / model_name / nsfw_flag / + # real_source / label "1:Fake 0:Real"). The Systematic / Manual / Commercial splits are + # FORBIDDEN (fakes paired with reals of unknown redistributability). We read PublicEval's + # parquet shards DIRECTLY via pyarrow (batch-bounded RAM, see _iter_parquet_rows) rather + # than HF `datasets` streaming, which OOMs on the single-6.5GB-row-group shard layout. + # The separate -Small repo is DROPPED: 760GB of 4GB single-row-group shards, same hazard + # for marginal gain — PublicEval alone is 21 generators (Dalle2/3, FLUX, MidjourneyV5/6, + # Firefly, Imagen3, Ideogram, kandinsky, stable_cascade, GALIP, DFGAN, ...) + reals. + HF_REPO = "OwensLab/CommunityForensics" + HF_SPLIT_DIR = "data/publicEval" # PublicEval shards only + N_SHARDS = 6 + # HELD-OUT GENERATORS: CF PublicEval contains Dalle3 + Imagen3, the exact families the + # OOS holdout quarantines (WildFake DALLE3/GigaGAN/Imagen). Training on them would make the + # "unseen generator" holdout no longer unseen -> excluded here to keep the eval honest. + # (Dalle2 is kept — a genuinely different, non-held-out model.) + # + # CANONICAL match: exact case-sensitive strings were brittle — "Imagen" vs + # "Imagen3", "DALLE3" vs "Dalle3" vs "dall-e 3" would slip through and leak the holdout. We + # normalize (lowercase, strip non-alphanumerics) and match the held-out FAMILY prefixes, so + # every spelling/version of a held-out family is caught. GigaGAN added even though CF likely + # lacks it — a no-op if absent, a leak-stopper if present. + _HELD_OUT_CANON = ("dalle3", "imagen", "gigagan") # family prefixes, canonicalized + _COLUMNS = ["image_data", "model_name", "nsfw_flag", "real_source", "label"] + + def __init__(self, families_wanted=None, include_reals=True, real_cap=None, + include_fakes=True): + self.families_wanted = families_wanted + self.include_reals = bool(include_reals) + self.include_fakes = bool(include_fakes) # False = reals-only (additive CF-reals set) + # real_cap: per-real-source quota (label 0), enforced in the base pull(). Bounds a + # materialized CF PublicEval real pull without a two-pass selection. None = uncapped. + self.real_cap = int(real_cap) if real_cap is not None else None + + @staticmethod + def _canon(name): + """Lowercase + strip non-alphanumerics so 'DALL-E 3'/'Dalle3'/'dalle_3' all -> 'dalle3'.""" + import re + return re.sub(r"[^a-z0-9]", "", str(name or "").lower()) + + def _is_held_out(self, model_name): + """True if model_name canonicalizes to (or starts with) any held-out family. 'Imagen3' + and 'Imagen' both match family 'imagen'; 'Dalle3' matches 'dalle3' but 'Dalle2' does NOT + (canon 'dalle2' has no held-out prefix).""" + c = self._canon(model_name) + return any(c.startswith(h) for h in self._HELD_OUT_CANON) + + def _shards(self): + """Yield (path, cleanup) per PublicEval parquet shard, smallest-first (a capped + pull fills its quota from the fewest shards). Each shard is DOWNLOADED TO LOCAL + DISK first (hf_hub_download: sequential streamed GET with built-in retry/resume), + then read locally — NOT read over an fsspec https handle. Parquet's columnar + layout makes iter_batches seek around the file; over HTTP that is hundreds of + random range reads that time out (FSTimeoutError observed). One sequential + download to disk + local batched read is fast, resilient, and still RAM-bounded. + cleanup() deletes the shard so peak disk is ONE shard (~3-6GB), not all 27GB. + Overridden in tests.""" + from huggingface_hub import hf_hub_download + from concurrent.futures import ThreadPoolExecutor + # smallest-first by known size ordering (shard 0 = 3.3GB ... shard 2 = 6.5GB) + order = (0, 4, 5, 1, 3, 2)[:self.N_SHARDS] + + def _dl(i): + return hf_hub_download(self.HF_REPO, + f"{self.HF_SPLIT_DIR}/HFCF_PublicEval_{i}.parquet", + repo_type="dataset") + + # PREFETCH: download shard N+1 in the background while shard N is being read/decoded, + # overlapping the multi-minute download with processing. Peak disk = 2 shards (~12GB, + # safe) since cleanup deletes shard N before N+2 starts. One prefetch worker (serial + # downloads still saturate the HF pipe; more would just contend). + with ThreadPoolExecutor(max_workers=1) as pool: + fut = pool.submit(_dl, order[0]) + for k, i in enumerate(order): + path = fut.result() # current shard (already downloading) + if k + 1 < len(order): + fut = pool.submit(_dl, order[k + 1]) # kick off the NEXT while we process + def _cleanup(p=path): + try: os.remove(p) + except OSError: pass + yield path, _cleanup + + def _iter_examples(self): + for path, cleanup in self._shards(): + try: + with open(path, "rb") as f: + for ex in _iter_parquet_rows(f, self._COLUMNS): + if ex.get("nsfw_flag"): + continue # NSFW: never materialize + if not self.include_reals and _read_label(ex) == 0: + continue # aggregate real license chain unverified + if not self.include_fakes and _read_label(ex) == 1: + continue # reals-only (additive CF-reals dataset) + if self._is_held_out(ex.get("model_name")): + continue # keep the OOS holdout genuinely unseen + yield ex + finally: + cleanup() # bounded disk: one shard at a time + def _to_fields(self, ex): + # Record per-image model_name (OpenRAIL-M pass-through). Respect the label field + # so any real rows in a safe subset land as label 0. + lbl = _read_label(ex) + if lbl is None: # unresolvable label -> skip (never default fake) + return None + img = ex.get("image_data") if ex.get("image_data") is not None else ex.get("image") + # image_data is the ORIGINAL encoded blob (bytes or {"bytes":...}); pass it through as + # `raw` so pull() writes it VERBATIM (no re-encode, no baked corruption, no PNG bloat). + raw = img if isinstance(img, (bytes, bytearray)) else (img.get("bytes") if isinstance(img, dict) else None) + return {"label": lbl, + "generator_id": "none" if lbl == 0 else (ex.get("model_name", "") or ""), + "real_source": (ex.get("real_source", "none") or "none") if lbl == 0 else "none", + "image": _as_pil(img), "raw": raw} + +class SIDSetSource(DatasetSource): + name = "sid_set" + def __init__(self, n_real, n_fake): + self.quotas = (int(n_real), int(n_fake)) + + def _iter_examples(self): + from datasets import Image as HFImage, load_dataset + dataset = load_dataset("saberzl/SID_Set", split="train", streaming=True) + dataset = dataset.cast_column("image", HFImage(decode=False)) + seen = [0, 0] + for ex in dataset: + lbl = int(ex["label"]) + if lbl not in (0, 1) or seen[lbl] >= self.quotas[lbl]: + continue + seen[lbl] += 1 + yield ex + if tuple(seen) == self.quotas: + break + + def _to_fields(self, ex): + lbl = int(ex["label"]) + image = ex["image"] + raw = image.get("bytes") if isinstance(image, dict) else None + return {"label": lbl, "generator_id": "none" if lbl == 0 else "sid_synth", + "real_source": ex.get("real_source", "sid"), "image": _as_pil(image), + "raw": raw} + +# WildFakeSource lives in data.curator.wildfake — it fetches ONLY selected members via +# RangeZip (range requests over the ModelScope zips), never whole 25-54GB archives, and +# handles the mapped per-generator taxonomy + §5.4 held-out exclusion. Re-exported here +# so default_corpus_sources / callers keep one import site. +from data.provenance.wildfake import WildFakeSource # noqa: E402,F401 + +class OmniFakeSource(DatasetSource): + """OmniFake (MIT) via the OFFICIAL HF datasets loader: MoeNew/OmniFake (streaming). + 45 generators + real. It is imagefolder-style: the class/label encodes real-vs-fake + and the folder name IS the generator (e.g. FLUX_Dev, DALLE_3, StyleGAN_3, real).""" + name = "omnifake" + hf_id = "MoeNew/OmniFake" + def __init__(self, split="train"): self.split = split + def _iter_examples(self): + from datasets import load_dataset + ds = load_dataset(self.hf_id, split=self.split, streaming=True) + # imagefolder datasets expose an int `label` + a `label`->foldername mapping via + # features; surface the folder name so classify_family sees the generator string. + names = None + try: + feat = ds.features.get("label") + names = getattr(feat, "names", None) + except Exception: + names = None + for ex in ds: + if names is not None and isinstance(ex.get("label"), int): + ex = dict(ex); ex["label_name"] = names[ex["label"]] + yield ex + def _to_fields(self, ex): + # The folder name is the ground truth for BOTH real/fake and generator. + name = ex.get("label_name") or _read_generator(ex) + is_real = isinstance(name, str) and name.strip().lower() in ("real", "0_real", "authentic") + return {"label": 0 if is_real else 1, + "generator_id": "none" if is_real else (name or ""), + "real_source": "omnifake" if is_real else "none", + "image": _as_pil(_read_image(ex))} + +# NOTE: DeepFace_part_v4 was RETIRED (probe finding): it is 15x HiFiMask_*.zip = face +# ANTI-SPOOFING / presentation-attack data (real faces behind masks), NOT AIGC +# fully-synthetic-vs-genuine. WildFake's Images/Real/ffhq.zip already covers real faces +# Apache-clean, so DeepFace added no fitting signal. Do not re-add without a real->fake +# face-generation source. + +class CifakeSource(DatasetSource): + """CIFAKE (MIT), aux 32x32. Reads a local CIFAKE directory + (birdy654/cifake-real-and-ai-generated-synthetic-images). Label from the + FAKE/REAL folder name. No network in _iter_examples — walks the mounted dir.""" + name = "cifake" + def __init__(self, root, split="train"): + self.root = root; self.split = split + def _iter_examples(self): + import glob + base = os.path.join(self.root, self.split) + # CIFAKE layout: //{REAL,FAKE}/*.jpg + for cls in ("FAKE", "REAL"): + for fp in glob.glob(os.path.join(base, cls, "*")): + yield {"path": fp, "cls": cls} + def _to_fields(self, ex): + is_real = ex["cls"].strip().upper() == "REAL" + return {"label": 0 if is_real else 1, + "generator_id": "none" if is_real else "cifake_synth", + "real_source": "cifake" if is_real else "none", + "image": _as_pil(ex["path"])} + +# NOTE: real images now come from the SAME materialized commercial-safe sources that +# provide fakes (SID_Set / WildFake / OmniFake carry real rows at label 0; WildFake's +# Images/Real/{coco,ffhq,celebahq,afhq,church}.zip are the primary diverse-real supply). +# There is no separate real-pool recipe source — each source's _to_fields sets label + +# real_source. + +# Generators held OUT of training (the OOS holdout axes). Their fakes never enter the +# training corpus; the holdout build (WildFakeSource(holdout_only=True)) selects them + +# COCO-val2017 reals. DALLE3 is the spec §5.4 set; GigaGAN (GAN) + Imagen (research +# diffusion) add two more OOS axes (user decision, 2026-08-29). +WILDFAKE_HELD_OUT = frozenset({"DALLE3", "GigaGAN", "Imagen"}) + +def default_corpus_sources(per_family_cap): + """Commercial-safe training defaults: aggregate sources contribute fakes only; + SID_Set supplies independently capped real and fully-synthetic rows.""" + return [ + WildFakeSource(per_gen_cap=per_family_cap, held_out=WILDFAKE_HELD_OUT, + exclude_val2017=True, include_reals=False), + CommunityForensicsSource(include_reals=False), + SIDSetSource(n_real=per_family_cap, n_fake=per_family_cap), + ] + +def wildfake_holdout_source(per_gen_cap=3000, spec_exact=False): + """The §5.4 OOS holdout: held-out generators' fakes (DALLE3/GigaGAN/Imagen) + + COCO-val2017 reals. Materialized separately, NEVER trained on. + + per_gen_cap (default 3000) bounds each generator — without it this pulls ALL held-out + images (DALLE3 8843 + GigaGAN 27610 + Imagen 47435 = ~84k), slow and unnecessary for a + stable AUC. + + spec_exact=True materializes the MANDATED §5.4 set: DALLE3 = all 8843, COCO-val2017 = all + 4998 (via per-group overrides = None → take-all), while GigaGAN/Imagen stay bounded at + per_gen_cap (they are our bonus OOS axes, not part of §5.4). Use this for the graded number.""" + overrides = None + if spec_exact: + # None => take-all for the §5.4-mandated groups; GigaGAN/Imagen fall back to per_gen_cap. + overrides = {"DALLE3": None, "coco": None} + return WildFakeSource(per_gen_cap=per_gen_cap, held_out=WILDFAKE_HELD_OUT, + exclude_val2017=True, holdout_only=True, per_gen_overrides=overrides) + +def official_holdout_source(): + """The OFFICIAL holdout artifact (spec: sharded-corpus-pipeline): ALL WildFake + COCO-val2017 reals + ALL DALLE3 fakes. No quota or cap code path for these two + groups (None == take-all); every other held-out group is excluded (0).""" + return WildFakeSource(per_gen_cap=None, held_out=WILDFAKE_HELD_OUT, + exclude_val2017=True, holdout_only=True, + per_gen_overrides={"DALLE3": None, "coco": None, + "GigaGAN": 0, "Imagen": 0}) + + +def cf_reals_source(real_cap=None): + """ADDITIVE dataset source: Community Forensics PublicEval REALS only (fake rows + excluded), bounded per real_source by real_cap. These are reals from the SAME + acquisition pipeline as the CF fakes already in the corpus — the fix for the source + confound (val real-side AUC 0.618: the detector partly reads SID-pipeline = real). + PublicEval is the redistribution-safe split (user-approved 2026-08-31); held-out + generators are irrelevant to reals but the filter is harmless. Published as its own + dataset (track5-cf-reals) and mounted ALONGSIDE the existing corpus — no rebuild, + no overwrite; the training box's RoutingResolver discovers it by shard prefix.""" + return CommunityForensicsSource(include_reals=True, include_fakes=False, + real_cap=real_cap, families_wanted=None) + + +def sid_fakes_source(n_fake=2000): + """ADDITIVE dataset source: SID_Set FAKES only (n_real=0), tagged generator_id + 'sid_synth'. The other half of the source-confound fix: putting fakes into the + real-dominant SID pipeline breaks 'SID-pipeline = real' (symmetric to cf_reals_source + breaking 'CF-pipeline = fake'). This is also the SID-fake scale fix (the live corpus + has only 199 SID fakes — the family-cap bug; the per-generator cap fix lands only on a + rebuild, but this additive dataset delivers the full 2000 WITHOUT one). Published as + track5-sid-fakes, mounted alongside the corpus. SIDSetSource's own quota yields exactly + n_fake fakes (no cap_fake_families in the standalone build path).""" + return SIDSetSource(n_real=0, n_fake=int(n_fake)) + + +def oos_validation_source(n_per_gen=3000): + """Separate OOS validation artifact: fixed samples of GigaGAN + Imagen. These are + validation data, not official holdout, and never enter training.""" + return WildFakeSource(per_gen_cap=None, held_out=WILDFAKE_HELD_OUT, + exclude_val2017=True, holdout_only=True, + per_gen_overrides={"GigaGAN": int(n_per_gen), + "Imagen": int(n_per_gen), + "DALLE3": 0, "coco": 0}) + + +def pull_all(sources, out_dir, per_family_cap=None, seed=0, max_examine=None, + selection=None, progress_cb=None) -> pd.DataFrame: + # Pass the SAME dataset root to every source; pull() namespaces by source name via + # rel_path (/fakes/...), so files land at out_dir//fakes/... and round-trip. + # progress_cb is forwarded only to sources whose pull() accepts it (WildFake); others + # ignore it. Kept per-source so long network stages stream live progress to wandb. + import inspect + parts = [] + for s in sources: + kw = dict(seed=seed, per_family_cap=per_family_cap, + max_examine=max_examine, selection=selection) + if progress_cb is not None and "progress_cb" in inspect.signature(s.pull).parameters: + kw["progress_cb"] = progress_cb + parts.append(s.pull(out_dir, **kw)) + return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame(columns=RECORD_COLUMNS) diff --git a/src/data/provenance/splitter.py b/src/data/provenance/splitter.py new file mode 100644 index 0000000000000000000000000000000000000000..ccf59bbd41ef90bd2e835b820fdc41656aa1f3d3 --- /dev/null +++ b/src/data/provenance/splitter.py @@ -0,0 +1,118 @@ +# src/manifest/splitter.py +import numpy as np +import pandas as pd + +def assign_split(df, mode, val_frac, group_col, holdout_values, seed): + df = df.copy() + if mode == "holdout": + df["split"] = np.where(df["generator_id"].isin(holdout_values), "val", "train") + return df + if mode == "grouped": + rng = np.random.default_rng(seed) + # each group's stratum label = most common label within the group + grp_label = df.groupby(group_col)["label"].agg(lambda s: s.mode().iloc[0]) + val_groups = set() + for lbl, groups_in_stratum in grp_label.groupby(grp_label): + gvals = np.asarray(groups_in_stratum.index) # ndarray -> safe to shuffle + rng.shuffle(gvals) + n = len(gvals) + n_val = max(1, int(round(n * val_frac))) + if n >= 2: + n_val = min(n_val, n - 1) # keep >=1 in train per label + val_groups.update(gvals[:n_val]) + df["split"] = np.where(df[group_col].isin(val_groups), "val", "train") + return df + raise ValueError(f"unknown split mode {mode}") + + +def assign_mixed_split(df, val_frac=0.15, seed=0): + """Group-hold out fake families while randomly splitting a single-source real pool. + + SID reals all have ``real_source='sid'``. Treating that one source as a grouped unit + would put every real in one split. Fakes still need family-level isolation, while + genuine SID rows can be deterministically sampled because there is no second real + source whose acquisition signature could leak across the boundary. + """ + if not 0 < val_frac < 1: + raise ValueError("val_frac must be between 0 and 1") + + out = df.copy() + out["split"] = "train" + rng = np.random.default_rng(seed) + + real_indices = out.index[out["label"] == 0].to_numpy(copy=True) + rng.shuffle(real_indices) + if len(real_indices) >= 2: + n_real_val = min(max(1, int(round(len(real_indices) * val_frac))), + len(real_indices) - 1) + out.loc[real_indices[:n_real_val], "split"] = "val" + + fake = out[out["label"] == 1].copy() + if not fake.empty: + fake["_family"] = fake["family"].astype(str).str.strip().str.lower() + if fake["_family"].nunique() >= 2: + fake = assign_split(fake, mode="grouped", val_frac=val_frac, + group_col="_family", holdout_values=None, seed=seed) + out.loc[fake.index, "split"] = fake["split"] + elif len(fake) >= 2: + fake_indices = fake.index.to_numpy(copy=True) + rng.shuffle(fake_indices) + n_fake_val = min(max(1, int(round(len(fake_indices) * val_frac))), + len(fake_indices) - 1) + out.loc[fake_indices[:n_fake_val], "split"] = "val" + + for label in (0, 1): + label_rows = out[out["label"] == label] + if len(label_rows) >= 2 and set(label_rows["split"]) != {"train", "val"}: + raise RuntimeError(f"mixed split failed to retain label {label} in both splits") + return out + + +def assign_hybrid_split(df, grouped_frac=0.1, random_frac=0.1, seed=0): + """Per-epoch val = TWO slices measuring different things (user decision): + - val_grouped: whole generators/real-sources held out of train (cross-generator + generalization signal). group = generator_id for fakes, real_source for reals. + - val_random: a stratified-by-label random slice of the REMAINING rows + (in-distribution recall signal). + The rest is train. Returns df with split in {train, val_grouped, val_random}. + Both val slices carry both classes (stratified); grouped holds whole groups out.""" + df = df.copy().reset_index(drop=True) + rng = np.random.default_rng(seed) + # CASE-CANONICAL group key: the same generator under different casing (e.g. WildFake's + # "DALLE2" vs CommForensics' "Dalle2") must be ONE group, or the whole-generator holdout + # leaks — training on one casing while "holding out" the other is not a genuine unseen + # generator. Lowercase the grouping key so such variants collapse to one held-out group. + group = df.apply(lambda r: (r["generator_id"] if r["label"] == 1 + else r["real_source"]).strip().lower(), axis=1) + df["_grp"] = group + + # 1) val_grouped: hold out whole groups, ~grouped_frac of rows, stratified by the + # group's modal label so BOTH classes are represented in the grouped-val slice. + grp_label = df.groupby("_grp")["label"].agg(lambda s: s.mode().iloc[0]) + val_groups = set() + for lbl, groups_in_stratum in grp_label.groupby(grp_label): + gvals = np.asarray(groups_in_stratum.index) + rng.shuffle(gvals) + n = len(gvals) + n_val = max(1, int(round(n * grouped_frac))) + if n >= 2: + n_val = min(n_val, n - 1) # keep >=1 group in train per label + val_groups.update(gvals[:n_val]) + split = np.where(df["_grp"].isin(val_groups), "val_grouped", "train") + + # 2) val_random: from the remaining TRAIN rows, take ~random_frac (of the whole df) + # at random, stratified by label so both classes appear. + n_random_target = int(round(len(df) * random_frac)) + train_idx = np.where(split == "train")[0] + if n_random_target > 0 and len(train_idx) > 0: + labels = df["label"].to_numpy() + picked = [] + for lbl in (0, 1): + pool = train_idx[labels[train_idx] == lbl] + k = min(len(pool), max(1, n_random_target // 2)) + if len(pool): + picked.extend(rng.choice(pool, size=k, replace=False)) + split[np.array(picked, dtype=int)] = "val_random" + + df["split"] = split + return df.drop(columns=["_grp"]) diff --git a/src/data/provenance/subsample.py b/src/data/provenance/subsample.py new file mode 100644 index 0000000000000000000000000000000000000000..8f97e4ef18a9d69bbf00c4ad98af41252ef40c38 --- /dev/null +++ b/src/data/provenance/subsample.py @@ -0,0 +1,34 @@ +import pandas as pd + + +def cap_fake_families(df, cap, seed): + """Cap synthetic rows per GENERATOR while retaining every real row. + + The cap exists to stop any single prolific generator from dominating the fake + class. It is keyed on ``generator_id`` — one quota per distinct generator — NOT on + the coarse ``family`` label. + + Keying on ``family`` silently conflated every generator that classifies to + ``family="unknown"`` (SID's ``sid_synth`` plus the ~20k WildFake/CommForensics rows + whose model name matches no family rule) into ONE bucket, so a cap of N kept only N + rows *across all of them*. That is exactly how the 2,000 requested SID fakes + collapsed to 199 (2000 * 2000 / 22000 ≈ 182). Per-generator keying gives each + generator its own quota, preserving generator diversity (the unseen-generator axis + the hidden test grades) instead of pooling distinct models. + + Real supply (label 0) is controlled by each licensed source's explicit quota and is + never capped here — reals all carry ``generator_id="none"``, so capping them would + collapse every genuine image into one bucket and recreate the scarcity this curator + is intended to avoid. + """ + if cap is None: + return df.copy().reset_index(drop=True) + if cap <= 0: + raise ValueError("cap must be positive or None") + + parts = [df[df["label"] == 0]] + for _gen, group in df[df["label"] == 1].groupby("generator_id", sort=False): + if len(group) > cap: + group = group.sample(n=cap, random_state=seed) + parts.append(group) + return pd.concat(parts).sort_index().reset_index(drop=True) diff --git a/src/data/provenance/val_holdout.py b/src/data/provenance/val_holdout.py new file mode 100644 index 0000000000000000000000000000000000000000..b1448b8298eb0595e0d0642ce303264d9ff5a2fd --- /dev/null +++ b/src/data/provenance/val_holdout.py @@ -0,0 +1,192 @@ +"""Track5 §5.4 Validation/Demo Holdout — a QUARANTINED benchmark. + +Spec (techjam spec, §5.4): a fixed subset of WildFake for participants to +*demonstrate* performance — "Do not use the following data during training." +Non-AIGC COCO val2017 = 4998 reals, AIGC DALL-E Advanced = 8843 fakes. + +Design: +- The holdout is MATERIALIZED into its own dataset (`wildfake_valset`), separate + from the training corpus, and every row is stamped `split="val_holdout"` — a + distinct quarantine tag, NOT train/val. +- The load-bearing deliverable is `assert_no_holdout_in_training`: a hard guard + that raises if any holdout row (by split OR by source_dataset) reaches a + training manifest, so the demo data can never leak into training. + +Same manifest schema as the training corpus (MANIFEST_COLUMNS) so the resolver +and eval harness operate on it unchanged. Import-safe: the ModelScope SDK import +is deferred inside WildFakeValSource._iter_examples. +""" +import pandas as pd +from data.provenance.families import classify_family +from data.provenance.manifest_build import MANIFEST_COLUMNS +from data.provenance.sources import DatasetSource, RECORD_COLUMNS + +# --- Quarantine constants (the whole point: this data never trains) --------- +VAL_HOLDOUT_SPLIT = "val_holdout" +WILDFAKE_VALSET_SOURCE = "wildfake_valset" +DALLE_ADVANCED_GENERATOR = "dalle_advanced" # classify_family(...) -> "dalle" +COCO_VAL2017_REAL_SOURCE = "coco_val2017" + +# WildFake ModelScope dataset + the two partitions this holdout draws from. +# RUNTIME-VERIFY: the exact internal partition/subset identifiers below are the +# spec's human labels, NOT confirmed against WildFake's on-disk parquet schema +# (needs network). Confirm before the real pull; see WildFakeValSource. +WILDFAKE_MODELSCOPE_ID = "hy2628982280/WildFake" +COCO_VAL2017_PARTITION = "COCO_val2017" # RUNTIME-VERIFY exact name +DALLE_ADVANCED_PARTITION = "DALLE_Advanced" # RUNTIME-VERIFY exact name + + +def build_val_holdout_manifest(records: pd.DataFrame) -> pd.DataFrame: + """Pure logic: raw holdout records -> a MANIFEST_COLUMNS manifest. + + Normalizes identity columns purely from `label` (0=COCO real, 1=DALL-E fake) + so the manifest is correct regardless of what the raw records carried, and + stamps every row with the quarantine split. + """ + df = records.copy() + is_fake = df["label"].astype(int) == 1 + df["source_dataset"] = WILDFAKE_VALSET_SOURCE + df["generator_id"] = is_fake.map( + lambda f: DALLE_ADVANCED_GENERATOR if f else "none") + df["real_source"] = is_fake.map( + lambda f: "none" if f else COCO_VAL2017_REAL_SOURCE) + df["family"] = df["generator_id"].map(classify_family) + df["split"] = VAL_HOLDOUT_SPLIT + # Backfill any resolution columns a raw record might lack, mirroring + # assemble_manifest's defensiveness. + if "ref_type" not in df.columns: + df["ref_type"] = "materialized" + if "mount_slug" not in df.columns: + df["mount_slug"] = "" + if "rel_path" not in df.columns: + df["rel_path"] = df["path"] + return df[MANIFEST_COLUMNS] + + +def _holdout_mask(df: pd.DataFrame) -> pd.Series: + """Row-wise True where a row is holdout — by split OR by source_dataset. + Both signals are checked so a row that lost one tag is still caught.""" + mask = pd.Series(False, index=df.index) + if "split" in df.columns: + mask = mask | (df["split"] == VAL_HOLDOUT_SPLIT) + if "source_dataset" in df.columns: + mask = mask | (df["source_dataset"] == WILDFAKE_VALSET_SOURCE) + return mask + + +def assert_no_holdout_in_training(train_manifest_df: pd.DataFrame) -> None: + """HARD GUARD — raise ValueError if any quarantined holdout row is present. + + Call this on any manifest destined for training. Detects the holdout by + `split == "val_holdout"` OR `source_dataset == "wildfake_valset"`.""" + mask = _holdout_mask(train_manifest_df) + n = int(mask.sum()) + if n: + raise ValueError( + f"Train-exclusion guard tripped: {n} quarantined val_holdout row(s) " + f"found in a training manifest (split=='{VAL_HOLDOUT_SPLIT}' or " + f"source_dataset=='{WILDFAKE_VALSET_SOURCE}'). §5.4 forbids training " + f"on the demo holdout.") + + +def filter_out_holdout(df: pd.DataFrame) -> pd.DataFrame: + """Defensive: drop any holdout rows, returning a guard-clean training df.""" + return df[~_holdout_mask(df)].reset_index(drop=True) + + +def drop_holdout_image_leaks(train_df: pd.DataFrame, holdout_df: pd.DataFrame, + threshold: int = 6) -> pd.DataFrame: + """IMAGE-level leakage guard (complements the TAG-level assert_no_holdout_in_training). + + TWO passes, cheap-and-exact first: + 1. EXACT identity: drop any train row whose stable id (img_id, else rel_path) matches a + holdout row's. This catches the case both sets draw the SAME source image (e.g. both pull + COCO reals from WildFake) — no file I/O, 100% reliable. This is the leak the pHash pass + under-caught in practice (95 shared COCO reals survived a threshold-6 pHash sweep because + phash requires reading every file, which is slow/fragile at 12k holdout images). + 2. RE-ENCODE identity: drop any remaining train row whose DCT pHash is within `threshold` + (Hamming) of ANY holdout image — catches a re-encoded copy that arrives with clean tags + and a different id. Reuses dedup_audit's pHash. + """ + if len(train_df) == 0 or len(holdout_df) == 0: + return train_df.reset_index(drop=True) + + def _idcol(df): + for c in ("img_id", "rel_path"): + if c in df.columns and df[c].notna().any(): + return df[c].astype(str) + return df["path"].astype(str) + + # --- pass 1: exact id match (no file reads) --- + hold_ids = set(_idcol(holdout_df)) + tr_ids = _idcol(train_df) + keep_exact = ~tr_ids.isin(hold_ids) + df = train_df[keep_exact].reset_index(drop=True) + if len(df) == 0: + return df + + # --- pass 2: pHash re-encode match on what remains --- + from data.provenance.dedup_audit import phash, hamming + holdout_hashes = [phash(p) for p in holdout_df["path"]] + def _leaks(p): + h = phash(p) + return any(hamming(h, hh) <= threshold for hh in holdout_hashes) + keep = ~df["path"].map(_leaks) + return df[keep].reset_index(drop=True) + + +class WildFakeValSource(DatasetSource): + """Materializes the §5.4 holdout from WildFake (ModelScope parquet) into its + own `wildfake_valset` dataset: COCO val2017 reals (label 0) + DALL-E Advanced + fakes (label 1). Heavy ModelScope import is DEFERRED inside _iter_examples + for import-safety (this module imports with no network / no SDK present).""" + name = WILDFAKE_VALSET_SOURCE + + def __init__(self, coco_partition=COCO_VAL2017_PARTITION, + dalle_partition=DALLE_ADVANCED_PARTITION, + modelscope_id=WILDFAKE_MODELSCOPE_ID): + self.coco_partition = coco_partition + self.dalle_partition = dalle_partition + self.modelscope_id = modelscope_id + + def _iter_examples(self): + # DEFERRED heavy import — keeps `import data.provenance.val_holdout` cheap + # and network-free (pytest-collection-safe). + from modelscope.msdatasets import MsDataset + # RUNTIME-VERIFY: WildFake's exact partition/subset names are not + # confirmed offline. We tag each example with which partition it came + # from so _to_fields can assign the label deterministically. If the + # subset_name identifiers below are wrong, this is the single place to + # fix — do NOT silently guess elsewhere. + for partition, tag in ((self.coco_partition, "real"), + (self.dalle_partition, "fake")): + ds = MsDataset.load(self.modelscope_id, subset_name=partition, + split="train") + for ex in ds: + ex = dict(ex) + ex["_holdout_tag"] = tag + yield ex + + def _to_fields(self, ex) -> dict: + is_fake = ex.get("_holdout_tag") == "fake" + return { + "label": 1 if is_fake else 0, + "generator_id": DALLE_ADVANCED_GENERATOR if is_fake else "none", + "real_source": "none" if is_fake else COCO_VAL2017_REAL_SOURCE, + "image": ex["image"], + } + + +def evaluate_on_holdout(predict_fn, holdout_df, cfg, logger=None) -> dict: + """Thin wrapper: run the existing eval harness on the quarantined holdout. + No harness logic is duplicated — this only exists to name the intent.""" + from eval.harness import evaluate + return evaluate(predict_fn, holdout_df, cfg, logger=logger) + + +__all__ = [ + "VAL_HOLDOUT_SPLIT", "WILDFAKE_VALSET_SOURCE", "DALLE_ADVANCED_GENERATOR", + "COCO_VAL2017_REAL_SOURCE", "build_val_holdout_manifest", + "assert_no_holdout_in_training", "filter_out_holdout", "drop_holdout_image_leaks", + "WildFakeValSource", "evaluate_on_holdout", "RECORD_COLUMNS", +] diff --git a/src/data/provenance/wildfake.py b/src/data/provenance/wildfake.py new file mode 100644 index 0000000000000000000000000000000000000000..3079aa0a9dfe4cc62c899fd2f4eacc89d86d5424 --- /dev/null +++ b/src/data/provenance/wildfake.py @@ -0,0 +1,393 @@ +# src/curator/wildfake.py +"""WildFake member classification, derived from the mapped taxonomy +(docs/specs/wildfake_membermap.log). Given a zip's source-relative path and an image +member's in-zip path, return (label, generator, timing, real_source). + +Two member-path shapes exist in WildFake: + STRUCTURED: ///... (DALLE, GAN_based, Other_based, + personalizedSD, SDwithAdaptor) -> generator = segment after the timing dir. + FLAT: .png or /imgs/.png (ADM, VQDM, Imagen, Midjourney/SD + part zips) -> no per-generator subfolder, so generator = the ZIP's stem and + timing comes from the zip PATH (…/Typical|Advanced/… when present). +Reals live under Images/Real/.zip -> label 0, real_source = zip stem. +""" +import os + +_TIMINGS = ("Typical", "Advanced") + + +def _zip_stem(zip_rel): + """'Images/Diffusion_based/DALLE.zip' -> 'DALLE'; multi-part + 'Images/Diffusion_based/Midjourney/Advanced/part_1.zip' -> 'Midjourney'.""" + parts = zip_rel.split("/") + # For part-zips the generator is the dir two-or-more levels up (…/Midjourney/Advanced/part_1.zip + # -> Midjourney); for single zips it's the filename stem. + stem = os.path.splitext(parts[-1])[0] + if stem.startswith("part_") or stem.startswith("part"): + # walk up past the timing dir to the generator dir + for i in range(len(parts) - 2, 0, -1): + if parts[i] not in _TIMINGS: + return parts[i] + return stem + + +def _zip_timing(zip_rel): + """Timing from the zip PATH for flat part-zips (…/Advanced/part_1.zip -> 'Advanced').""" + for seg in zip_rel.split("/"): + if seg in _TIMINGS: + return seg + return None + + +def classify_member(zip_rel, member): + """-> {label, generator, timing, real_source}. Pure string logic (no I/O).""" + is_real = "/Real/" in zip_rel or zip_rel.startswith("Images/Real") + if is_real: + return {"label": 0, "generator": "none", "timing": None, + "real_source": _zip_stem(zip_rel)} + segs = member.split("/") + # STRUCTURED: find a timing segment; generator is the NEXT segment. + for i, s in enumerate(segs[:-1]): + if s in _TIMINGS: + gen = segs[i + 1] if i + 1 < len(segs) else _zip_stem(zip_rel) + return {"label": 1, "generator": gen, "timing": s, "real_source": "none"} + # FLAT: generator = zip stem; timing from the zip path (part-zips) else None. + return {"label": 1, "generator": _zip_stem(zip_rel), + "timing": _zip_timing(zip_rel), "real_source": "none"} + + +# --- §5.4 holdout member predicates (for quarantine + dalle3 exclusion) ----------- +def is_val2017_member(member): + """COCO val2017 reals (the §5.4 demo reals, 4998).""" + m = member.replace("\\", "/") + return "/val2017/" in m or m.startswith("coco/coco2017/val2017") + + +def is_dalle3_member(member): + """DALLE3 = §5.4 DALL-E-Advanced fakes (8843), the held-out OOS generator.""" + return "DALLE3" in member or "/Advanced/DALLE3" in member + + +def plan_wildfake_members(zip_to_members, per_gen_cap, held_out=frozenset(), + exclude_val2017=True, seed=0, holdout_only=False, stratify=False, + oversample=0, per_gen_overrides=None, include_reals=True): + """PASS-1.5 selection for WildFake: from {zip_rel: [member,...]} choose which members + to fetch, PER GENERATOR, honoring held-out exclusion. Pure/offline (no network). + + held_out: generator names kept OUT of training (their members are dropped unless + holdout_only, which selects ONLY them + the val2017 reals = the OOS set). + exclude_val2017: drop COCO val2017 reals from a training pull (they're §5.4 holdout). + stratify: sample the per-generator cap PROPORTIONALLY across the Typical/Advanced + timing strata (WildFake distinguishes older vs newer generations), so a + generator's quota isn't accidentally drawn all from one era. False = flat + uniform sample per generator (original behavior). + Returns a list of (zip_rel, member, fields), capped at per_gen_cap per generator. + """ + import random + rng = random.Random(seed) + # bucket members by generator (fakes)/real_source (reals) AND timing stratum + buckets = {} + for zip_rel, members in zip_to_members.items(): + for m in members: + c = classify_member(zip_rel, m) + if c["label"] == 0 and not include_reals: + continue + gen = c["generator"] if c["label"] == 1 else c["real_source"] + is_ho = (gen in held_out) or is_dalle3_member(m) + is_v2017 = is_val2017_member(m) + if holdout_only: + if not (is_ho or is_v2017): + continue + else: + if is_ho: + continue + if is_v2017 and exclude_val2017: + continue + if c["label"] == 0 and not is_v2017 and "val2017" in m: + continue + key = (zip_rel, gen, c["label"]) + strat = c.get("timing") or "none" + buckets.setdefault(key, {}).setdefault(strat, []).append((m, c)) + + def _emit(zip_rel, gen, label, chosen, plan): + for m, c in chosen: + plan.append((zip_rel, m, {"label": c["label"], + "generator_id": "none" if c["label"] == 0 else gen, + "real_source": gen if c["label"] == 0 else "none"})) + + # OVERSAMPLE: fetch a few EXTRA candidates per generator beyond per_gen_cap so the + # save loop can replace undecodable images and still land EXACTLY at the cap (count + # stability). oversample=0 preserves original behavior. + # per_gen_overrides: {gen_or_real_source: cap} overrides per_gen_cap for named groups — + # used for the spec-exact §5.4 holdout (DALLE3 all 8843, coco all 4998, GigaGAN/Imagen + # bounded). A None override means "take all" for that group. + overrides = per_gen_overrides or {} + plan = [] + for (zip_rel, gen, label), strata in buckets.items(): + all_items = [it for lst in strata.values() for it in lst] + cap = overrides[gen] if gen in overrides else per_gen_cap + if cap == 0: + continue # explicit 0 = group EXCLUDED (oversample must not resurrect it) + eff_cap = None if cap is None else cap + max(0, oversample) + if eff_cap is None or len(all_items) <= eff_cap: + _emit(zip_rel, gen, label, all_items, plan) + continue + if not stratify or len(strata) <= 1: + _emit(zip_rel, gen, label, rng.sample(all_items, eff_cap), plan) + continue + # STRATIFIED: split the cap across timing strata proportionally to their size, + # largest-remainder for the leftover, so both eras are represented. + total = len(all_items) + alloc, rema = {}, [] + for s, lst in strata.items(): + exact = eff_cap * len(lst) / total + alloc[s] = min(len(lst), int(exact)) + rema.append((exact - int(exact), s)) + short = eff_cap - sum(alloc.values()) + for _frac, s in sorted(rema, reverse=True): + if short <= 0: + break + if alloc[s] < len(strata[s]): + alloc[s] += 1; short -= 1 + for s, k in alloc.items(): + _emit(zip_rel, gen, label, rng.sample(strata[s], k), plan) + return plan + + +class WildFakeSource: + """WildFake source over the ModelScope raw zips. BULK-downloads each generator zip + whole (one sustained connection — far faster than per-member range requests, which + pay a full TLS handshake to ModelScope/China per image), extracts only the SAMPLED + members (per-generator cap + held-out exclusion), then DELETES the zip. `download_workers` + whole zips run concurrently. scan()/pull() build on plan_wildfake_members; zip_opener/ + zip_list are injectable for unit tests.""" + name = "wildfake" + modelscope_id = "hy2628982280/WildFake" + + def __init__(self, per_gen_cap=2000, held_out=frozenset(), exclude_val2017=True, + seed=0, holdout_only=False, zip_opener=None, zip_list=None, + download_workers=2, mode="segmented", stratify=True, + seg_workers=16, max_gap=256 * 1024, per_gen_overrides=None, + include_reals=True): + self.per_gen_cap = per_gen_cap + self.per_gen_overrides = per_gen_overrides or {} + self.held_out = frozenset(held_out) + self.exclude_val2017 = exclude_val2017 + self.seed = seed + self.holdout_only = holdout_only + self.include_reals = bool(include_reals) + self._zip_opener = zip_opener + self._zip_list = zip_list + # download_workers: number of ZIPS processed concurrently. + self.download_workers = max(int(download_workers), 1) + # seg_workers/max_gap: WITHIN a zip, read_many fetches coalesced segments across + # seg_workers threads. Bandwidth probe measured the ModelScope CDN throttles + # ~8MB/s PER CONNECTION, scaling to ~97MB/s at 16 conns -> seg_workers=16. max_gap + # 256KB (not 1MB) keeps segments smaller/more-numerous so there's enough of them to + # actually use 16 workers (over-coalescing into 3 huge segments would cap it at 3). + self.seg_workers = max(int(seg_workers), 1) + self.max_gap = int(max_gap) + # mode: "segmented" (default, FAST) = RangeZip + coalesced parallel range reads. + # "bulk" = download whole zip then extract. stratify: proportional Typical/Advanced. + self.mode = mode + self.stratify = stratify + + # --- ModelScope wiring (deferred; overridden in tests) ------------------ + def _list_zips(self): + if self._zip_list is not None: + return self._zip_list + from modelscope.hub.api import HubApi + api = HubApi() + all_zips = [getattr(f, "path", None) or str(f) + for f in api.list_repo_files(self.modelscope_id, repo_type="dataset") + if (str(getattr(f, "path", f)).startswith("Images/") + and str(getattr(f, "path", f)).lower().endswith(".zip"))] + return self._dedupe_parts(all_zips) + + @staticmethod + def _dedupe_parts(zips): + """Multi-part generator zips (…/Midjourney/Advanced/part_1.zip … part_7.zip) are + REDUNDANT splits of ONE generator — part_1 alone holds 30k+ images, far more than + any per-gen cap needs. Opening all 14 SD/Midjourney 52GB parts just to parse their + central directories is THE build bottleneck (dir-parse is 30s/zip). Keep only the + FIRST part per (generator, timing) group; non-part zips pass through untouched.""" + import os + seen_groups, kept = set(), [] + for z in zips: + base = os.path.basename(z) + if base.startswith("part_") or base.startswith("part"): + group = os.path.dirname(z) # …/Midjourney/Advanced (gen+timing) + if group in seen_groups: + continue # already have a part for this generator + seen_groups.add(group) + kept.append(z) + return kept + + def _open_zip(self, zip_rel): + """Return a zip handle exposing .image_members(), .read(name), optionally + .read_many(names), .close(). Test opener injected verbatim. Otherwise: + segmented -> RangeZip (reads central dir only; read_many coalesces sampled + members into a few contiguous range GETs — no over-download, no + per-image handshake). THE fast path. + bulk -> download the whole zip, extract from disk, delete on close.""" + if self._zip_opener is not None: + return self._zip_opener(zip_rel) + from modelscope.hub.api import HubApi + api = HubApi() + if self.mode == "segmented": + from data.provenance.rangezip import RangeZip + url = api.get_dataset_file_url(zip_rel, "WildFake", "hy2628982280", "master") + rz = RangeZip(url) + seg_workers, max_gap = self.seg_workers, self.max_gap + + class _Seg: + def image_members(self): return rz.image_members() + def read(self, name): return rz.read(name) + def read_many(self, names): + # MEASURED (bandwidth probe): ModelScope CDN throttles ~8MB/s PER + # CONNECTION but scales near-linearly to ~16 conns (97MB/s). So we want + # many concurrent range GETs. Small max_gap => more, smaller segments + # => higher effective concurrency (up to seg_workers). + return rz.read_many(names, max_gap=max_gap, workers=seg_workers) + def close(self): pass # nothing downloaded to delete + return _Seg() + # bulk fallback + import os, zipfile + from modelscope import dataset_file_download + local = dataset_file_download(dataset_id=self.modelscope_id, file_path=zip_rel) + z = zipfile.ZipFile(local) + + class _Handle: + def image_members(self): + return [n for n in z.namelist() + if n.lower().endswith((".jpg", ".jpeg", ".png", ".webp", ".bmp"))] + def read(self, name): return z.read(name) + def close(self): + z.close() + try: os.remove(local) + except OSError: pass + return _Handle() + + def _members_of(self, handle): + return handle.image_members() + + def scan(self, max_examine=None): + """Inventory per generator. Downloads each zip to read its member list, then + releases it (bounded disk). For a pure-metadata scan without downloads, use the + label CSVs (kernel-side); this scan is the source-agnostic fallback.""" + import pandas as pd + counts = {} + for zr in self._list_zips(): + h = self._open_zip(zr) + try: + for m in h.image_members(): + c = classify_member(zr, m) + if c["label"] == 0 and not self.include_reals: + continue + gkey = c["generator"] if c["label"] == 1 else c["real_source"] + counts[(gkey, c["label"])] = counts.get((gkey, c["label"]), 0) + 1 + finally: + h.close() + rows = [{"source_dataset": self.name, "group_key": g, "label": l, "count": c} + for (g, l), c in counts.items()] + return pd.DataFrame(rows, columns=["source_dataset", "group_key", "label", "count"]) + + def pull(self, out_dir, seed=0, per_family_cap=None, max_examine=None, selection=None, + progress_cb=None, progress_every=200): + """BULK per-zip materialization: for each zip, download it whole (one fast + connection), plan+extract only the members we sampled (held-out excluded), write + JPEGs, then DELETE the zip before the next — peak disk = one zip. INTERLEAVED: + the next zip downloads (in a background thread) while the current one extracts, + so download and extract overlap and fetching starts within seconds (no upfront + all-directories barrier). + + progress_cb: fn(dict) called after each zip with {done,total,imgs_per_s,elapsed_s, + zips_done,n_zips,eta_s} so a long build streams live to wandb.""" + import os, time, pandas as pd + from concurrent.futures import ThreadPoolExecutor, as_completed + from PIL import Image + from data.provenance.sources import RECORD_COLUMNS, _as_pil + from data.provenance.manifest_build import jpeg_bytes_per_pixel, save_source_image + os.makedirs(out_dir, exist_ok=True) + + zips = self._list_zips() + n_zips = len(zips) + t0 = time.perf_counter() + + def _process_zip(zr): + """Download one whole zip, extract its planned members, delete it. Fully + independent -> safe to run many in parallel. Each (generator lives in exactly + ONE zip after part-dedup) so per-zip filename counters never collide.""" + handle = self._open_zip(zr) + zip_rows, local_counts = [], {} + try: + plan = plan_wildfake_members({zr: handle.image_members()}, + per_gen_cap=self.per_gen_cap, + held_out=self.held_out, + exclude_val2017=self.exclude_val2017, + seed=self.seed, holdout_only=self.holdout_only, + stratify=self.stratify, + per_gen_overrides=self.per_gen_overrides, + include_reals=self.include_reals, + # +5% reserves (min 8) to replace undecodable blobs + oversample=max(8, (self.per_gen_cap or 0) // 20)) + # Batch-fetch all sampled members in a FEW coalesced range GETs when the + # handle supports it (segmented mode); else read one at a time (bulk/test). + want = [m for _zr, m, _f in plan] + if hasattr(handle, "read_many") and want: + blobs = handle.read_many(want) + else: + blobs = {m: handle.read(m) for m in want} + for _zr, m, f in plan: + gkey = f["generator_id"] if f["label"] == 1 else f["real_source"] + # COUNT-STABLE: the plan is OVERSAMPLED (cap + buffer). Stop saving a + # generator once it reaches the true cap — extra candidates are reserves + # that only get used to replace undecodable images below. + if self.per_gen_cap is not None and local_counts.get(gkey, 0) >= self.per_gen_cap: + continue + zt = _zr.replace("/", "_").replace(".zip", "") + # VERBATIM source bytes (no re-encode): JPEG re-encode bakes a scored + # corruption into the 'clean' ground truth; PNG re-encode ~4x's disk (blows + # working-disk limits) for zero fidelity gain. save_source_image writes the + # source blob unchanged with the sniffed extension. + base = f"{self.name}/fakes/{self.name}_{gkey}_{zt}_{local_counts.get(gkey,0)}" + p_noext = os.path.join(out_dir, base) + try: + # A corrupt/undecodable blob must NOT abort the whole run — skip it; a + # reserve candidate takes its slot so the per-generator count stays exact. + p, w, h = save_source_image(blobs[m], p_noext) + rel = os.path.relpath(p, out_dir).replace(os.sep, "/") + except Exception as e: + print(f"[wildfake] skip undecodable {_zr}:{m} ({type(e).__name__}: {e})", flush=True) + import glob as _g + for _f in _g.glob(p_noext + ".*"): # remove any partial write + try: os.remove(_f) + except OSError: pass + continue + zip_rows.append({"path": p, "label": f["label"], "source_dataset": self.name, + "generator_id": f["generator_id"], "real_source": f["real_source"], + "native_w": w, "native_h": h, + "jpeg_bytes_per_pixel": jpeg_bytes_per_pixel(p), + "ref_type": "materialized", "mount_slug": "", "rel_path": rel}) + local_counts[gkey] = local_counts.get(gkey, 0) + 1 + finally: + handle.close() + return zip_rows + + # N whole zips downloaded+extracted CONCURRENTLY (each a sustained connection; + # parallel streams saturate total bandwidth). Peak disk ~= download_workers zips. + rows, zips_done = [], 0 + with ThreadPoolExecutor(max_workers=self.download_workers) as ex: + futs = {ex.submit(_process_zip, zr): zr for zr in zips} + for fut in as_completed(futs): + rows.extend(fut.result()) + zips_done += 1 + if progress_cb: + dt = time.perf_counter() - t0 + progress_cb({"phase": "fetch", "done": len(rows), "zips_done": zips_done, + "n_zips": n_zips, "imgs_per_s": len(rows) / dt if dt > 0 else 0.0, + "elapsed_s": dt, + "eta_s": (n_zips - zips_done) * (dt / zips_done) if zips_done else None}) + print(f"[wildfake] zip {zips_done}/{n_zips} {futs[fut]} -> {len(rows)} total " + f"imgs ({time.perf_counter()-t0:.0f}s)", flush=True) + return pd.DataFrame(rows, columns=RECORD_COLUMNS) diff --git a/src/data/provenance/writer.py b/src/data/provenance/writer.py new file mode 100644 index 0000000000000000000000000000000000000000..8445a17201e3a4adcab5d401ff9db946dc3b1134 --- /dev/null +++ b/src/data/provenance/writer.py @@ -0,0 +1,98 @@ +# src/shards/writer.py +"""Deterministic ZIP64 shard planning + writing (spec: sharded-corpus-pipeline). + +Members are ZIP_STORED (images are already compressed), ordered by a stable sort, +with fixed timestamps/attrs so identical inputs reproduce identical SHA-256 digests. +write_shard is transactional up to the atomic rename: it writes `.partial`, +re-opens and verifies (names, order, count, sizes, per-member CRC by reading every +member), computes the digest, then renames. On any failure only the .partial is +removed; raw inputs are never touched here (deletion is the orchestrator's job, +and only AFTER the journal records the verified shard).""" +import hashlib +import os +import zipfile + +_FIXED_DT = (1980, 1, 1, 0, 0, 0) + + +def validate_member_name(name: str) -> str: + if not name or "\\" in name or name.startswith("/"): + raise ValueError(f"unsafe member name: {name!r}") + parts = name.split("/") + if any(p in ("", ".", "..") for p in parts): + raise ValueError(f"unsafe member name: {name!r}") + return name + + +def plan_shards(df, target_bytes=1 << 30, prefix="corpus"): + """Assign every manifest row to a deterministic shard. Returns a sorted copy with + shard_name, member_name, byte_size columns. Close-before-exceed packing; a single + oversized member occupies its own shard.""" + out = df.copy() + out["member_name"] = [validate_member_name(r) for r in out["rel_path"]] + if out["member_name"].duplicated().any(): + dupes = out.loc[out["member_name"].duplicated(), "member_name"].head(3).tolist() + raise ValueError(f"duplicate member names: {dupes}") + ident = out["img_id"] if "img_id" in out.columns else out["rel_path"] + out["_grp"] = out["family"].where(out["label"] == 1, out["real_source"]) + out["_ident"] = ident.astype(str) + out = out.sort_values(["source_dataset", "label", "_grp", "_ident", "rel_path"], + kind="mergesort").reset_index(drop=True) + # honor a pre-supplied byte_size column (resume path: raws of completed shards are + # deleted, their sizes recovered from the shard central directories by the caller) + if "byte_size" not in out.columns: + out["byte_size"] = [os.path.getsize(p) for p in out["path"]] + names, cur_bytes, cur_idx = [], 0, 0 + for size in out["byte_size"]: + if cur_bytes > 0 and cur_bytes + size > target_bytes: + cur_idx += 1 + cur_bytes = 0 + names.append(f"{prefix}-{cur_idx:05d}.shard") # .shard NOT .zip: keeps generic file-sync tools from auto-extracting the archives + cur_bytes += size + out["shard_name"] = names + return out.drop(columns=["_grp", "_ident"]) + + +def write_shard(zip_path: str, members) -> dict: + """Write members [(member_name, src_path)...] to zip_path via .partial + verify + + digest + atomic rename. Returns shard record with per-member sizes and CRCs.""" + partial = zip_path + ".partial" + try: + with open(partial, "wb") as fh: + with zipfile.ZipFile(fh, "w", compression=zipfile.ZIP_STORED, + allowZip64=True) as zf: + for name, src in members: + validate_member_name(name) + zi = zipfile.ZipInfo(name, date_time=_FIXED_DT) + zi.compress_type = zipfile.ZIP_STORED + zi.create_system = 3 + zi.external_attr = 0o644 << 16 + with open(src, "rb") as f: + zf.writestr(zi, f.read()) + fh.flush() + os.fsync(fh.fileno()) + # --- verify by re-reading the closed archive, never trusting write state --- + rec_members, payload = [], 0 + with zipfile.ZipFile(partial) as zf: + if zf.namelist() != [n for n, _ in members]: + raise AssertionError(f"shard member order mismatch in {zip_path}") + for name, src in members: + info = zf.getinfo(name) + data = zf.read(name) # zipfile CRC-checks on read + if len(data) != os.path.getsize(src): + raise AssertionError(f"size mismatch for {name}") + rec_members.append({"member_name": name, "byte_size": len(data), + "crc32": info.CRC}) + payload += len(data) + h = hashlib.sha256() + with open(partial, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + os.replace(partial, zip_path) # atomic on same filesystem + return {"sha256": h.hexdigest(), "zip_bytes": os.path.getsize(zip_path), + "payload_bytes": payload, "n_members": len(rec_members), + "members": rec_members} + except BaseException: + if os.path.exists(partial): + os.remove(partial) + raise diff --git a/src/data/shard_reader.py b/src/data/shard_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..3d8ab51cbc623bc9fc34cdc562fc1dcba31c3d59 --- /dev/null +++ b/src/data/shard_reader.py @@ -0,0 +1,49 @@ +# src/data/shard_reader.py +"""Read training images directly from mounted ZIP shards — no extraction. + +Handles are opened lazily per worker process and kept in a small LRU (default 8). +The pid check recreates handles after DataLoader fork/spawn: a zipfile handle +shared across processes corrupts reads. Rows resolve via shard_name + member_name +from the candidate manifest.""" +import io +import os +import zipfile +from collections import OrderedDict + +from PIL import Image + + +class ShardResolver: + def __init__(self, shard_dir, max_handles=8): + self.shard_dir = shard_dir + self.max_handles = int(max_handles) + self._handles = OrderedDict() + self._pid = os.getpid() + + def _get_zip(self, shard_name): + if os.getpid() != self._pid: # forked/spawned worker: stale handles + self._handles = OrderedDict() + self._pid = os.getpid() + zf = self._handles.get(shard_name) + if zf is None: + zf = zipfile.ZipFile(os.path.join(self.shard_dir, shard_name)) + self._handles[shard_name] = zf + while len(self._handles) > self.max_handles: + _, old = self._handles.popitem(last=False) + old.close() + else: + self._handles.move_to_end(shard_name) + return zf + + def resolve(self, row): + zf = self._get_zip(row["shard_name"]) + data = zf.read(row["member_name"]) + img = Image.open(io.BytesIO(data)) + img.load() + return img.convert("RGB") + + def __getstate__(self): # picklable into DataLoader workers + return {"shard_dir": self.shard_dir, "max_handles": self.max_handles} + + def __setstate__(self, state): + self.__init__(state["shard_dir"], state["max_handles"]) diff --git a/src/eval/__init__.py b/src/eval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/eval/error_gallery.py b/src/eval/error_gallery.py new file mode 100644 index 0000000000000000000000000000000000000000..24382dfd416e7ac02a854c3452eaf3b426766bec --- /dev/null +++ b/src/eval/error_gallery.py @@ -0,0 +1,69 @@ +# src/eval/error_gallery.py +"""Offline error-gallery construction from the predictions parquet. + +The predictions table (one row per image x condition, with provenance) is the ONLY +input the model produces — galleries are built here, offline, AFTER a run, and can be +rebuilt at any bands without re-running training. Images are pulled on demand from the +mounted source via SourceResolver (rel_path is the stable address), so nothing is +saved during the run. + +Banding is by PROBABILITY, not a 0.5 threshold (ROC AUC is rank-only, so 0.5 is +meaningless). Two bands the user asked for, stratified per dataset: + - confident_wrong : LABEL-RELATIVE — a fake (label 1) with prob < conf, or a real + (label 0) with prob > 1-conf. The embarrassing, high-confidence misses. + - borderline : prob in [border_lo, border_hi]. The model is genuinely unsure. +""" +import pandas as pd + + +def select_gallery_rows(pred_df, n_per=30, conf=0.15, border=(0.4, 0.6), + condition="clean", seed=0): + """Return a DataFrame of selected rows with an added `band` column. + + n_per caps rows per (band x source_dataset). Selection is on `condition` (default + "clean"). confident_wrong sorted by how wrong (most confident-wrong first); + borderline sampled. Deterministic given seed. + """ + df = pred_df[pred_df["condition"] == condition].copy() + lo, hi = border + is_fake = df["label"] == 1 + # LABEL-RELATIVE confident-wrong: fake&low OR real&high. + cw = df[(is_fake & (df["prob"] < conf)) | (~is_fake & (df["prob"] > 1 - conf))].copy() + cw["band"] = "confident_wrong" + # "wrongness" = distance of prob from the correct end (1 for fakes, 0 for reals). + cw["_wrong"] = (cw["label"] - cw["prob"]).abs() + bl = df[(df["prob"] >= lo) & (df["prob"] <= hi)].copy() + bl["band"] = "borderline" + bl["_wrong"] = (bl["prob"] - 0.5).abs() # closeness to 0.5 (most-unsure first) + + out = [] + for band, g, ascending in (("confident_wrong", cw, False), ("borderline", bl, True)): + for _src, grp in g.groupby("source_dataset"): + # confident_wrong: most-wrong first (desc). borderline: most-unsure first (asc dist). + grp = grp.sort_values("_wrong", ascending=ascending).head(n_per) + out.append(grp) + if not out: + return pd.DataFrame(columns=list(pred_df.columns) + ["band"]) + res = pd.concat(out, ignore_index=True).drop(columns=["_wrong"]) + return res + + +def render_gallery(selected_df, resolver, out_dir, thumb=256): + """Pull each selected image via `resolver` (offline, from the mounted source) and + write a 256px thumbnail into out_dir///. Filename encodes prob+label so + the miss is legible at a glance. Returns {(band,source): count}. Never crashes on a + single bad image (logs + skips) — the selection table is the source of truth.""" + import os + counts = {} + for _, r in selected_df.iterrows(): + try: + img = resolver.resolve(r) # rel_path/source_dataset -> PIL + img = img.convert("RGB"); img.thumbnail((thumb, thumb)) + d = os.path.join(out_dir, r["band"], str(r["source_dataset"])) + os.makedirs(d, exist_ok=True) + tag = f"L{int(r['label'])}_p{float(r['prob']):.2f}_{os.path.basename(str(r['img_id']))}.jpg" + img.save(os.path.join(d, tag), quality=88) + k = (r["band"], r["source_dataset"]); counts[k] = counts.get(k, 0) + 1 + except Exception as e: # noqa: BLE001 — one bad image must not kill the gallery + print(f"[gallery] skip {r.get('img_id')}: {e}") + return counts diff --git a/src/eval/explain.py b/src/eval/explain.py new file mode 100644 index 0000000000000000000000000000000000000000..0e892ad834b8755c132a535ac1a5007008557917 --- /dev/null +++ b/src/eval/explain.py @@ -0,0 +1,99 @@ +# src/eval/explain.py +"""Explainability for the detector — the business-facing "why was this flagged" layer. + +The student's AttentionPoolPatchHead already emits a per-patch attention map (the MIL +pooling weights). We turn it into (a) a spatial heatmap overlay a human moderator can read, +(b) a compact per-image rationale (verdict + confidence + where the model looked), and +(c) an aggregate reliance summary (how real/fake separation holds up per corruption). No +extra forward pass — the attention is a free byproduct of the pooling head. + +Why this matters for the business case: a content-moderation / provenance system cannot act +on a black-box score. A localized "this region looks synthetic" plus a robustness profile is +what makes the flag actionable and auditable. +""" +import numpy as np + + +def attention_heatmap(attn, img_size): + """Upsample a per-patch attention vector (length = n_side², a square grid) to a + (H, W) heatmap normalized to [0,1]. Nearest-neighbour upsample keeps patch structure.""" + n = len(attn) + side = int(round(np.sqrt(n))) + if side * side != n: + raise ValueError(f"attention length {n} is not a square patch grid") + grid = np.asarray(attn, dtype=float).reshape(side, side) + rng = grid.max() - grid.min() + grid = (grid - grid.min()) / rng if rng > 0 else np.zeros_like(grid) + W, H = img_size + ys = (np.arange(H) * side // H).clip(0, side - 1) + xs = (np.arange(W) * side // W).clip(0, side - 1) + return grid[ys][:, xs] + + +def explain_prediction(prob, attn, img_size, threshold=0.5): + """Per-image explanation: verdict, calibrated-ish confidence, attention concentration + (focus_fraction = share of attention mass in the top 10% of patches — high means the + decision hinges on a localized region), and a short human rationale.""" + attn = np.asarray(attn, dtype=float) + k = max(1, int(round(0.1 * len(attn)))) + focus_fraction = float(np.sort(attn)[::-1][:k].sum() / (attn.sum() + 1e-9)) + verdict = "ai-generated" if prob >= threshold else "authentic" + confidence = float(abs(prob - 0.5) * 2) # 0 at the boundary, 1 at the extremes + where = ("concentrated in a localized region" if focus_fraction > 0.4 + else "spread broadly across the image") + rationale = (f"Flagged as {verdict} (p={prob:.2f}); the detector's evidence is " + f"{where} (focus={focus_fraction:.0%}).") + return {"verdict": verdict, "confidence": confidence, + "focus_fraction": focus_fraction, "heatmap": attention_heatmap(attn, img_size), + "rationale": rationale} + + +def aggregate_reliance(pred_df): + """Per-condition real/fake separation = mean prob(fake rows) − mean prob(real rows). + A large separation on clean that DEGRADES under a corruption tells the business exactly + which real-world transformation weakens the detector — the actionable robustness story.""" + summ = {} + for cond, g in pred_df.groupby("condition"): + fake = g[g["label"] == 1]["prob"] + real = g[g["label"] == 0]["prob"] + sep = float(fake.mean() - real.mean()) if len(fake) and len(real) else float("nan") + summ[cond] = {"separation": sep, "n": int(len(g)), + "fake_mean": float(fake.mean()) if len(fake) else float("nan"), + "real_mean": float(real.mean()) if len(real) else float("nan")} + return summ + + +def deletion_faithfulness(predict_fn, img, attn, n_regions=8, seed=0): + """Deletion-test faithfulness of an attention explanation. + + An attention map is only a real explanation if the regions it highlights are the ones + the model actually USES. We measure this by DELETING (graying out) the top-attention + patches and seeing how much the prediction drops, versus deleting random patches. Score + = (prob_full − prob_delete_top) − (prob_full − prob_delete_random) = the extra prediction + change attributable to removing the ATTENDED region. Higher = more faithful attention. + Returns a scalar; positive means attention aligns with the model's causal evidence.""" + import numpy as np + from PIL import Image + rng = np.random.default_rng(seed) + a = np.asarray(img.convert("RGB")).copy() + H, W = a.shape[:2] + n = len(attn) + side = int(round(np.sqrt(n))) + attn = np.asarray(attn, dtype=float) + + def _gray_patches(idxs): + out = a.copy() + for idx in idxs: + r, c = idx // side, idx % side + y0, y1 = r * H // side, (r + 1) * H // side + x0, x1 = c * W // side, (c + 1) * W // side + out[y0:y1, x0:x1] = 128 # neutral gray = "delete" + return Image.fromarray(out) + + k = max(1, min(n_regions, n)) + top_idx = np.argsort(attn)[::-1][:k] + rand_idx = rng.choice(n, size=k, replace=False) + p_full = float(predict_fn([img])[0]) + p_top = float(predict_fn([_gray_patches(top_idx)])[0]) + p_rand = float(predict_fn([_gray_patches(rand_idx)])[0]) + return (p_full - p_top) - (p_full - p_rand) # = p_rand - p_top diff --git a/src/eval/harness.py b/src/eval/harness.py new file mode 100644 index 0000000000000000000000000000000000000000..ac72e2eeeed9744dec36581fc947c11f0012141a --- /dev/null +++ b/src/eval/harness.py @@ -0,0 +1,210 @@ +# src/eval/harness.py +import numpy as np +from PIL import Image +from data.aug.augment import CORRUPTIONS, EVAL_GRID, apply_corruption +from eval.metrics import roc_auc, aggregate_scores +from eval.predictions import build_predictions_df +from eval.stream import plan_eval_batches +from analysis.format_probe import format_of as _format_of + +def evaluate(predict_fn, val_df, cfg, logger=None, corruption_fn=None, resolver=None, + watchdog=None) -> dict: + """Streaming bounded-memory evaluation over clean + the 6-corruption grid. + + corruption_fn(img, name, severity)->PIL lets the caller SWAP the corruption + implementation (default: the canonical single-impl `apply_corruption`). Passing an + implementation-holdout function (see evaluate_impl_holdout) measures AUC under + operators NOT seen during training — the corruption-implementation-invariance result.""" + corrupt = corruption_fn or apply_corruption + labels = val_df["label"].to_numpy() + # sharded rows (ref_type='shard') carry no usable path: resolve via the resolver + shard_rows = None + if resolver is not None and "shard_name" in val_df.columns: + shard_rows = val_df.reset_index(drop=True).to_dict("records") + paths = val_df["path"].tolist() if "path" in val_df.columns else [""] * len(val_df) + n = len(paths) + # img_id column if present, else fall back to path (aligned to `paths` order). + if "img_id" in val_df: + img_ids = val_df["img_id"].tolist() + else: + img_ids = list(paths) + # Provenance for offline error-gallery stratification + reconstruction. rel_path is + # the STABLE source address (source_dataset + rel_path -> SourceResolver re-fetches + # the image); image_path is an ephemeral /tmp path. Absent columns default to "". + def _col(name): + return val_df[name].tolist() if name in val_df else [""] * n + source_datasets = _col("source_dataset") + generator_ids = _col("generator_id") + rel_paths = _col("rel_path") + # native_w + format carry the codec/size axes so the format-shortcut probe (the decisive + # holdout test: JPEG-only, size-matched AUC) is derivable OFFLINE from the predictions table. + # native_w defaults to NaN; format is derived from the path extension when the column is absent. + native_ws = val_df["native_w"].tolist() if "native_w" in val_df else [float("nan")] * n + + # The 19 eval conditions, index 0 = clean. condition_idx -> (name, severity, key). + conditions = [("clean", None, "clean")] + for name in CORRUPTIONS: + for sev in EVAL_GRID[name]: + conditions.append((name, sev, f"{name}@{sev}")) + n_cond = len(conditions) + + # BOUNDED-MEMORY STREAMING (fixes host-RAM OOM): the old harness decoded EVERY val + # image into a list held across all conditions (+a full copy per condition) — tens + # of GB on a large val set. Instead we stream (image x condition) work through + # plan_eval_batches under a hard cap of `pool_max` decoded images resident at once; + # each image is decoded on first use and evicted the moment its last condition is + # scored. Memory is O(pool_max), independent of val-set size. + pool_max = max(getattr(cfg, "eval_pool_max", 64), 1) + batch = max(getattr(cfg, "eval_batch_size", 32), 1) + prob_mat = np.full((n_cond, n), np.nan, dtype=float) # [condition][image] + cache = {} # item_idx -> decoded PIL.Image (at most pool_max live) + done = np.zeros(n, dtype=int) # conditions scored per item -> evict at n_cond + + # FAIR EVAL: normalize real+fake through an identical resize+recompress so absolute + # resolution/format can't be the discriminator (the holdout's reals were 200px JPEG, fakes + # large PNG -> resolution-only AUC 1.0). Opt-in; any AUC after this is real detection signal. + fair = getattr(cfg, "fair_eval", False) + fair_res = getattr(cfg, "resolution", 518) + # PARALLEL CPU PREP (cfg.eval_workers>0): decode + corruption are PIL C calls that + # release the GIL, so a THREAD pool parallelizes them for real — the serial loop left + # 29 of 30 cores idle through every eval phase (GPU util 0%, the v12 4h stall's true + # anatomy). Decode stays serialized per item (cache coherence); corruption fans out. + eval_workers = int(getattr(cfg, "eval_workers", 0)) + pool = None + if eval_workers > 0: + from concurrent.futures import ThreadPoolExecutor + pool = ThreadPoolExecutor(max_workers=eval_workers) + + def _decode(item): + if item not in cache: + if resolver is not None and shard_rows is not None: + im = resolver.resolve(shard_rows[item]).convert("RGB") + else: + with Image.open(paths[item]) as im2: + im = im2.convert("RGB") + if fair: + from data.preprocess import fair_eval_normalize + im = fair_eval_normalize(im, fair_res) + cache[item] = im + return cache[item] + + def _make(item_cond): + item, cond = item_cond + name, sev, _ = conditions[cond] + base_im = cache[item] + return base_im if name == "clean" else corrupt(base_im, name, sev) + + # MID-EVAL PROGRESS: eval phases used to be silent for up to hours (GPU idle, no + # output — indistinguishable from a hang). Log every ~30s: cells done/total, rate, + # ETA — to stdout (run.log) and, when a logger is given, to wandb as eval/progress. + import time as _time + total_cells = n * n_cond + cells_done = 0 + _t0 = _time.time() + _t_last = _t0 + for bat in plan_eval_batches(n, n_cond, pool_max, batch): + for item, _cond in bat: + _decode(item) # serial: fills cache once per item + coords = list(bat) + if pool is not None: + imgs = list(pool.map(_make, coords)) + else: + imgs = [_make(ic) for ic in coords] + probs = predict_fn(imgs) + cells_done += len(coords) + if watchdog is not None: + watchdog.update(cells_done=cells_done, cells_total=total_cells) + if _time.time() - _t_last >= 30: + _t_last = _time.time() + rate = cells_done / max(_time.time() - _t0, 1e-9) + eta = (total_cells - cells_done) / max(rate, 1e-9) + print(f"[eval] {cells_done}/{total_cells} cells " + f"({100*cells_done/total_cells:.0f}%) | {rate:.0f} cells/s | " + f"eta {eta/60:.1f} min", flush=True) + if logger is not None: + try: + logger.log({"eval/progress_frac": cells_done / total_cells, + "eval/cells_per_s": rate}) + except Exception: # noqa: BLE001 - progress logging never breaks eval + pass + for (item, cond), p in zip(coords, probs): + prob_mat[cond, item] = float(p) + done[item] += 1 + if done[item] >= n_cond: + cache.pop(item, None) # last condition scored -> free the image + if pool is not None: + pool.shutdown(wait=False) + + # Assemble per-condition AUCs + the long-format predictions table from prob_mat. + pred_records = [] + per = {} + for cond, (name, sev, key) in enumerate(conditions): + col = prob_mat[cond] + per[key] = roc_auc(labels, col) + for i in range(n): + pred_records.append({ + "img_id": img_ids[i], "image_path": paths[i], + "source_dataset": source_datasets[i], "generator_id": generator_ids[i], + "rel_path": rel_paths[i], "native_w": native_ws[i], + "format": _format_of(paths[i]), "label": int(labels[i]), + "condition": name, "severity": sev, "prob": float(col[i]), + }) + + agg = aggregate_scores(per) + predictions = build_predictions_df(pred_records) + result = {"per_condition": per, **agg, "predictions": predictions} + + # Attention-map diagnostics are intentionally skipped: predict_fn only returns + # sigmoid probs (attn is discarded), so measuring attention peakedness would + # require a second forward pass / harness refactor. Not worth it — see report. + if logger is not None: + metrics = {f"eval/{k}": v for k, v in per.items()} # e.g. eval/jpeg@30 + metrics.update({ + "eval/auc_clean": agg["auc_clean"], + "eval/auc_robust_mean": agg["auc_robust_mean"], + "eval/auc_robust_worst": agg["auc_robust_worst"], + "eval/final": agg["final"], + }) + logger.log(metrics) + + return result + + +def _impl_corruption_fn(impl_choice, seed=0): + """Build a corruption_fn(img,name,sev) that applies a SPECIFIC implementation per + family. impl_choice maps family -> list of allowed impl names; one is chosen + deterministically per (family, severity) so the eval is reproducible.""" + import random + from data.aug.augment_impls import apply_impl, list_impls + def fn(img, name, sev): + allowed = impl_choice.get(name) or list_impls(name) + # deterministic pick within the allowed set, varied by severity + r = random.Random(f"{name}:{sev}:{seed}") + impl = allowed[r.randrange(len(allowed))] + return apply_impl(img, name, impl, sev) + return fn + + +def evaluate_impl_holdout(predict_fn, val_df, cfg, train_impls, test_impls, + logger=None, seed=0, resolver=None) -> dict: + """Corruption-IMPLEMENTATION-invariance measurement: evaluate the SAME model under + (a) the training implementations and (b) the held-out implementations, and report the + AUC gap. A small gap = the detector generalizes across operators (did NOT overfit the + grader's exact JPEG/resize/blur); a large gap = operator overfit. + + Returns {train_impl: , test_impl: , gap: {...}}. + train_impls / test_impls are {family: [impl_name,...]} from split_impls_train_test.""" + seen = evaluate(predict_fn, val_df, cfg, logger=None, + corruption_fn=_impl_corruption_fn(train_impls, seed), resolver=resolver) + held = evaluate(predict_fn, val_df, cfg, logger=None, + corruption_fn=_impl_corruption_fn(test_impls, seed), resolver=resolver) + gap = { + "auc_robust_mean_seen": seen["auc_robust_mean"], + "auc_robust_mean_held": held["auc_robust_mean"], + "robust_mean_gap": seen["auc_robust_mean"] - held["auc_robust_mean"], + "final_seen": seen["final"], "final_held": held["final"], + "final_gap": seen["final"] - held["final"], + } + if logger is not None: + logger.log({f"implholdout/{k}": v for k, v in gap.items()}) + return {"train_impl": seen, "test_impl": held, "gap": gap} diff --git a/src/eval/metrics.py b/src/eval/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..bddbb93b73ec52721ff4a1c3fe2af0c922ce989a --- /dev/null +++ b/src/eval/metrics.py @@ -0,0 +1,58 @@ +# src/eval/metrics.py +import numpy as np +from sklearn.metrics import roc_auc_score + +def roc_auc(labels: np.ndarray, probs: np.ndarray) -> float: + labels = np.asarray(labels); probs = np.asarray(probs) + if len(np.unique(labels)) < 2: + return float("nan") + return float(roc_auc_score(labels, probs)) + +def stratified_auc(df, by="generator_id", prob_col="prob", label_col="label"): + """Per-group AUC: each FAKE group (e.g. one generator) scored vs the SHARED real pool. + Reveals which generators the model separates well — logged live to wandb per epoch. + + df must carry label/prob/`by`. Reals (label 0) form one pool reused for every group. + Groups with <2 classes (no reals, or the group has no fakes) are skipped/NaN.""" + import numpy as np + reals = df[df[label_col] == 0] + if len(reals) == 0: + return {} + real_labels = reals[label_col].to_numpy() + real_probs = reals[prob_col].to_numpy() + out = {} + for gval, grp in df[df[label_col] == 1].groupby(by): + labels = np.concatenate([real_labels, grp[label_col].to_numpy()]) + probs = np.concatenate([real_probs, grp[prob_col].to_numpy()]) + out[str(gval)] = roc_auc(labels, probs) + return out + + +def aggregate_scores(per_condition: dict[str, float]) -> dict: + """Aggregate per-condition AUCs into the competition score. + + The organizer's metric is mean over the SIX corruption FAMILIES (jpeg/blur/resize/noise/ + jitter/crop), each family weighted EQUALLY. Our eval grid probes several severities per + family (jpeg@30/50/70/90, ...), so a naive mean over all severity cells over-weights + families with more severities (jpeg had 4 votes, jitter/crop 1). We therefore average the + severities WITHIN each family first (family AUC), then take the mean over families — so the + number matches the score that will actually be computed. `robust_worst` is the worst FAMILY. + Condition keys are 'family@severity' (or a bare 'family' with no '@').""" + clean = per_condition["clean"] + families: dict[str, list] = {} + for k, v in per_condition.items(): + if k == "clean": + continue + fam = k.split("@", 1)[0] # 'jpeg@30' -> 'jpeg'; 'blur' -> 'blur' + families.setdefault(fam, []).append(v) + family_auc = {fam: float(np.nanmean(vs)) for fam, vs in families.items()} + fam_vals = list(family_auc.values()) + robust_mean = float(np.mean(fam_vals)) if fam_vals else float("nan") + robust_worst = float(np.min(fam_vals)) if fam_vals else float("nan") + return { + "auc_clean": clean, + "auc_robust_mean": robust_mean, + "auc_robust_worst": robust_worst, + "per_family_auc": family_auc, + "final": 0.5 * clean + 0.5 * robust_mean, + } diff --git a/src/eval/predictions.py b/src/eval/predictions.py new file mode 100644 index 0000000000000000000000000000000000000000..a5811f08019c14186e394c3c3f70fe2ff4320b1d --- /dev/null +++ b/src/eval/predictions.py @@ -0,0 +1,51 @@ +# src/eval/predictions.py +"""Per-image prediction table for offline error analysis. + +The eval harness already computes per-image sigmoid probabilities under every +(clean + corruption) condition and, historically, discarded them after folding +into roc_auc(). This module captures those probs into a tidy long-format table +so precision/recall and TP/FP/TN/FN sampling can be derived OFFLINE at ANY +threshold — nothing here recomputes a forward pass. + +Schema (one row per image x condition), columns in this EXACT order: + img_id, image_path, source_dataset, generator_id, rel_path, + label, condition, severity, prob + - source_dataset/generator_id: provenance, so galleries stratify by dataset/generator + - rel_path : the STABLE source-relative address (source_dataset + rel_path re-addresses + the image via SourceResolver offline; image_path is an ephemeral /tmp path) + - condition: "clean" or a corruption name (jpeg/blur/resize/noise/jitter/crop) + - severity : None for clean, else the eval-grid severity value + - prob : sigmoid prob in [0,1] for that image under that condition +""" +import pandas as pd + +COLUMNS = ["img_id", "image_path", "source_dataset", "generator_id", "rel_path", + "native_w", "format", "label", "condition", "severity", "prob"] + + +def build_predictions_df(records) -> pd.DataFrame: + """Build the predictions DataFrame from an iterable of row dicts. + + Each record must carry the COLUMNS keys. Column order is enforced so the + downstream (label, prob)-at-any-threshold contract is stable. + """ + df = pd.DataFrame(list(records), columns=COLUMNS) + return df + + +def save_predictions(df: pd.DataFrame, out_dir, prefix="") -> dict[str, str]: + """Write {prefix}predictions.parquet AND {prefix}predictions.csv into out_dir. + + Returns {"parquet": , "csv": }. Both formats are written: parquet + for typed/compact offline analysis, csv for eyeballing / spreadsheet use. + `prefix` (e.g. 'base_' / 'fuse4scalar_') keeps per-arm artifacts SEPARATE — the + v10 sequential A/B overwrote the base arm's files, destroying the paired oracle. + """ + import os + + os.makedirs(out_dir, exist_ok=True) + parquet_path = os.path.join(out_dir, f"{prefix}predictions.parquet") + csv_path = os.path.join(out_dir, f"{prefix}predictions.csv") + df.to_parquet(parquet_path, index=False) + df.to_csv(csv_path, index=False) + return {"parquet": parquet_path, "csv": csv_path} diff --git a/src/eval/stream.py b/src/eval/stream.py new file mode 100644 index 0000000000000000000000000000000000000000..586941cd65bf1a3ea664278305b259397dc1c18f --- /dev/null +++ b/src/eval/stream.py @@ -0,0 +1,53 @@ +# src/eval/stream.py +"""Bounded-memory eval scheduler. + +The old eval harness decoded EVERY val image into a Python list and held it across +all 19 corruption conditions (plus a full copy per condition) — tens of GB of host +RAM on a large val set, which OOM-killed the kernel. This schedules eval work as a +stream of (item, condition) pairs under a hard cap on how many items are "live" +(decoded and resident) at once, so peak memory is O(pool_max) regardless of val size. + +Policy (user-specified, made deterministic): keep a small pool of admitted items; +each item tracks which conditions it still owes via a countdown. We drain OLDEST +items first (age-priority) so nothing lingers and the pool never bloats — a hard +no-stall guarantee (every step does min(batch, pending) work), where a random +age-weighted policy only makes stalls unlikely. AUC is order-independent, so +deterministic draining costs nothing and buys reproducibility. +""" +from collections import OrderedDict + + +def plan_eval_batches(n_items, n_conditions, pool_max, batch): + """Yield lists of (item_idx, condition_idx) pairs, each pair emitted exactly once. + + Invariants (see tests/test_eval_stream.py): + - coverage: every (item, condition) in the n_items x n_conditions grid appears once + - memory bound: at most `pool_max` items are live (admitted, not fully drained) + at any time — this is the decoded-images-in-RAM ceiling + - progress: each yielded batch has 1..`batch` pairs (never empty, never stalls) + + item_idx admission is in ascending order; within the live pool we emit the + OLDEST-admitted items' pending conditions first (age-priority), so an item is + fully consumed and evicted before newer ones dominate. + """ + if n_items <= 0 or n_conditions <= 0: + return + pool = OrderedDict() # item_idx -> next condition_idx to emit (insertion-ordered = age) + next_item = 0 + while next_item < n_items or pool: + # Admit new items (oldest-first FIFO) until the pool is full or items run out. + while len(pool) < pool_max and next_item < n_items: + pool[next_item] = 0 + next_item += 1 + # Build one batch by draining oldest items' pending conditions first. + bat = [] + for item in list(pool.keys()): + while pool[item] < n_conditions and len(bat) < batch: + bat.append((item, pool[item])) + pool[item] += 1 + if pool[item] >= n_conditions: + del pool[item] # fully drained -> evict (frees its decoded image) + if len(bat) >= batch: + break + if bat: + yield bat diff --git a/src/inference/__init__.py b/src/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/inference/predict.py b/src/inference/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..aed13dc7c08a568a5069f3abee57922721ada74b --- /dev/null +++ b/src/inference/predict.py @@ -0,0 +1,143 @@ +# src/inference/predict.py +"""Competition inference core. + +Loads a checkpoint bundle, reconstructs the detector, and scores every image in a +directory with a clean (no-augmentation) forward pass. Emits JSON entries of the +form {"image_path": , "pred": }. + +The heavy lifting lives here (not in the root predict.py wrapper) so it is +CPU-testable with the stub backbone. +""" +import json +import os + +import torch +from PIL import Image + +from data.preprocess import preprocess_image +from eval.explain import explain_prediction +from train.checkpoint import load_checkpoint, load_ensemble_checkpoint + +# Common image extensions, matched case-insensitively. +_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} + + +def _list_images(image_dir): + """Top-level image files in image_dir, sorted for deterministic order.""" + out = [] + for name in os.listdir(image_dir): + full = os.path.join(image_dir, name) + if not os.path.isfile(full): + continue + if os.path.splitext(name)[1].lower() in _IMAGE_EXTS: + out.append(full) + return sorted(out) + + +def _resolve_device(device): + if device is None or device == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + return device + + +@torch.no_grad() +def run_inference(image_dir, checkpoint_path, out_path="predictions.json", + batch_size=16, device="auto", explain=False, + checkpoint2_path=None, progress_cb=None) -> list[dict]: + """Score every image in image_dir; write and return the JSON entries. + + Returns a list of {"image_path", "pred"} dicts in sorted-path order. Bad/ + unreadable images are skipped with a warning (one bad file never aborts the + run). An empty directory writes [] and warns. + + explain=True enriches each entry with a human-readable moderation rationale + (verdict + confidence + focus_fraction) derived from the pooling head's + attention — the "why was this flagged" layer. It is FREE: the model already + returns the attention map, so no second forward pass is made. The dense + heatmap array is intentionally NOT written to JSON (too large); it stays + available via the eval/explain API for gallery rendering. + + progress_cb(entry, attn) is invoked once per scored image (attn is the + per-patch attention vector, or None when explain=False) — the hook the live + dashboard (--serve) uses to stream results without touching the JSON output. + """ + device = _resolve_device(device) + # Ensure the output's parent dir exists so --out subdir/preds.json doesn't crash. + os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True) + # ENSEMBLE (the ship config): weights.pt bundles both heads (blocks 8 and 10 of + # the same frozen backbone); final logit = mean over members. A single-model + # bundle loads as a one-member ensemble through the same path. + members = load_ensemble_checkpoint(checkpoint_path, map_location=device) + if checkpoint2_path: # optional extra member from a separate file + members += load_ensemble_checkpoint(checkpoint2_path, map_location=device) + for m, _ in members: + m.to(device) + m.eval() + model, cfg = members[0] # first member drives attn/XAI + preprocessing cfg + + paths = _list_images(image_dir) + if not paths: + print(f"[predict] warning: no images found in {image_dir}; writing empty output.") + with open(out_path, "w") as f: + json.dump([], f) + return [] + + results = [] + batch_tensors, batch_paths = [], [] + + def _flush(): + if not batch_tensors: + return + x = torch.stack(batch_tensors).to(device) + out = model(x) + logit = out["image_logit"] + for m2, _ in members[1:]: + logit = logit + m2(x)["image_logit"] + probs = torch.sigmoid(logit / len(members)).cpu().numpy() + attns = out["attn"].cpu().numpy() if (explain or progress_cb) else None + for i, (p, prob) in enumerate(zip(batch_paths, probs)): + entry = {"image_path": p, "pred": float(prob)} + if explain: + ex = explain_prediction(float(prob), attns[i], (cfg.resolution, cfg.resolution)) + entry.update({"verdict": ex["verdict"], "confidence": ex["confidence"], + "focus_fraction": ex["focus_fraction"], "rationale": ex["rationale"]}) + results.append(entry) + if progress_cb: + progress_cb(entry, attns[i] if attns is not None else None) + batch_tensors.clear() + batch_paths.clear() + + for p in paths: + try: + with Image.open(p) as im: + tensor = preprocess_image(im, cfg.resolution) + except Exception as e: # noqa: BLE001 — one bad image must not abort the run + print(f"[predict] warning: skipping unreadable image {p}: {e}") + continue + batch_tensors.append(tensor) + batch_paths.append(p) + if len(batch_tensors) >= batch_size: + _flush() + _flush() + + with open(out_path, "w") as f: + json.dump(results, f) + return results + + +def build_predict_fn_from_checkpoint(checkpoint_path, device="auto"): + """Load a checkpoint and return predict_fn(list[PIL]) -> np.ndarray of P(fake). Reusable + by offline analyses (patch diversity, galleries) that score arbitrary PIL crops with the + trained model. Fixed at cfg.resolution (deterministic).""" + import numpy as np + from data.preprocess import preprocess_image + dev = _resolve_device(device) + model, cfg = load_checkpoint(checkpoint_path, map_location=dev) + model.to(dev).eval() + + @torch.no_grad() + def predict_fn(imgs): + x = torch.stack([preprocess_image(im, cfg.resolution) for im in imgs]).to(dev) + return torch.sigmoid(model(x)["image_logit"]).cpu().numpy() + + return predict_fn diff --git a/src/inference/server.py b/src/inference/server.py new file mode 100644 index 0000000000000000000000000000000000000000..5dcc4c3d73c89aadd775f5d290fea4ac47809c9f --- /dev/null +++ b/src/inference/server.py @@ -0,0 +1,315 @@ +# src/inference/server.py +"""Live inference dashboard (predict.py --serve). + +A small stdlib-only web UI over a persistent inference worker: + + - QUEUE section: everything waiting to be scored. Seeded from --image-dir at + startup; grows at runtime — drag-and-drop, file-picker upload, or paste an + image straight onto the page (POST /upload). + - COMPLETED section: each scored image with the XAI layer — verdict, + probability, rationale, and the attention heatmap rendered as a red alpha + overlay on the grayscale image. + +The worker runs fully asynchronously to the page and rewrites the JSON output +(atomically) after every image, so the artifact predict.py promises exists at +all times, exactly as in headless mode. +""" +import io +import json +import os +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import numpy as np +import torch +from PIL import Image + +from data.preprocess import preprocess_image +from eval.explain import explain_prediction +from inference.predict import _list_images, _resolve_device +from train.checkpoint import load_ensemble_checkpoint + +_PAGE = """ +AIGC detector — live inference + +

AIGC detector — live inference

+
+

drop images here, click to pick, or just paste (Ctrl-V) anywhere

+ +

Completed

+

Queue

(empty)
+""" + + +class InferenceWorker: + """Persistent scoring loop over a grow-at-runtime queue. One model load; + each completion appends to results and atomically rewrites the JSON.""" + + def __init__(self, image_dir, checkpoint_path, checkpoint2_path=None, + out_path="predictions.json", device="auto", explain=True, + batch_size=16): + self.image_dir = image_dir + self.upload_dir = os.path.join(image_dir, "uploads") + self.out_path = out_path + self.explain = explain + self.lock = threading.Lock() + self.queue = list(_list_images(image_dir)) # pending paths, FIFO + self.done = [] # completed entries + self.attns = [] # per-entry attention (np or None) + self._overlay_cache = {} + self._wake = threading.Event() + self.device = _resolve_device(device) + self.members = load_ensemble_checkpoint(checkpoint_path, map_location=self.device) + if checkpoint2_path: + self.members += load_ensemble_checkpoint(checkpoint2_path, + map_location=self.device) + for m, _ in self.members: + m.to(self.device) + m.eval() + self.cfg = self.members[0][1] + + def enqueue_bytes(self, data, filename): + """Persist an uploaded/pasted image and add it to the queue.""" + os.makedirs(self.upload_dir, exist_ok=True) + safe = os.path.basename(filename) or f"upload-{int(time.time()*1000)}.png" + path = os.path.join(self.upload_dir, safe) + if os.path.exists(path): # de-collide repeated names + stem, ext = os.path.splitext(safe) + path = os.path.join(self.upload_dir, f"{stem}-{int(time.time()*1000)}{ext}") + with open(path, "wb") as f: + f.write(data) + with self.lock: + self.queue.append(path) + self._wake.set() + return path + + def _score(self, path): + with Image.open(path) as im: + entry, evid = score_image(self.members, im, self.device, self.explain) + entry = {"image_path": path, **entry} + return entry, evid + + def _write_json(self): + with self.lock: + payload = [{k: v for k, v in e.items()} for e in self.done] + tmp = self.out_path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f) + os.replace(tmp, self.out_path) # atomic on POSIX and Windows + + def run_forever(self): + os.makedirs(os.path.dirname(self.out_path) or ".", exist_ok=True) + self._write_json() # artifact exists from t=0 + while True: + with self.lock: + path = self.queue.pop(0) if self.queue else None + if path is None: + self._wake.wait(timeout=1.0) + self._wake.clear() + continue + try: + entry, attn = self._score(path) + except Exception as e: # noqa: BLE001 — one bad image never kills the worker + print(f"[serve] warning: skipping unreadable image {path}: {e}") + continue + with self.lock: + self.done.append(entry) + self.attns.append(attn) + self._write_json() + + # --- view helpers ------------------------------------------------------- + def snapshot(self): + with self.lock: + return { + "queue": [os.path.basename(p) for p in self.queue], + "done": [{**{k: v for k, v in e.items() if k != "image_path"}, + "name": os.path.basename(e["image_path"])} + for e in self.done], + } + + def overlay_png(self, idx, max_side=480): + with self.lock: + if idx in self._overlay_cache: + return self._overlay_cache[idx] + if idx >= len(self.done): + return None + entry, evid = self.done[idx], self.attns[idx] + with Image.open(entry["image_path"]) as im: + out = render_overlay(im, evid, max_side) + buf = io.BytesIO() + out.save(buf, format="PNG") + png = buf.getvalue() + with self.lock: + self._overlay_cache[idx] = png + return png + + +@torch.no_grad() +def score_image(members, img, device, explain=True): + """Score one PIL image with the ensemble. Returns (entry, evidence): entry is the + predict.py JSON record (minus image_path) and evidence the ensemble-mean SIGNED + per-patch logits (positive = fake evidence) for the overlay.""" + cfg = members[0][1] + x = preprocess_image(img, cfg.resolution).unsqueeze(0).to(device) + out = members[0][0](x) + logit = out["image_logit"] + plog = out["patch_logits"] + for m2, _ in members[1:]: + o2 = m2(x) + logit = logit + o2["image_logit"] + plog = plog + o2["patch_logits"] + prob = float(torch.sigmoid(logit / len(members))[0]) + attn = out["attn"][0].cpu().numpy() # rationale (focus_fraction) + evid = (plog[0] / len(members)).cpu().numpy() + entry = {"pred": prob} + if explain: + ex = explain_prediction(prob, attn, (cfg.resolution, cfg.resolution)) + entry.update({"verdict": ex["verdict"], "confidence": ex["confidence"], + "focus_fraction": ex["focus_fraction"], + "rationale": ex["rationale"]}) + return entry, evid + + +def render_overlay(img, evid, max_side=480): + """Grayscale image with the SIGNED evidence map alpha-overlaid: red where the + ensemble sees fake evidence, blue where it sees real, alpha proportional to + |evidence| so neutral regions stay grayscale — same rendering as the docs' + error tiles. Returns a PIL image.""" + img = img.convert("RGB") + img.thumbnail((max_side, max_side)) + W, H = img.size + base = np.asarray(img.convert("L"), dtype=float)[..., None].repeat(3, -1) + if evid is not None: + g = np.asarray(evid, dtype=float) + side = int(round(np.sqrt(len(g)))) + grid = g.reshape(side, side) + m = max(1e-6, float(np.abs(grid).max())) + ys = (np.arange(H) * side // H).clip(0, side - 1) + xs = (np.arange(W) * side // W).clip(0, side - 1) + up = grid[ys][:, xs] / m # [-1,1] signed + r = np.clip(up, 0, 1)[..., None]; b = np.clip(-up, 0, 1)[..., None] + color = np.concatenate([255 * r, np.zeros_like(r), 255 * b], axis=-1) + alpha = 0.65 * np.abs(up)[..., None] + base = base * (1 - alpha) + color * alpha + return Image.fromarray(base.clip(0, 255).astype("uint8")) + + +def _make_handler(worker): + class Handler(BaseHTTPRequestHandler): + def log_message(self, *a): # keep the console clean for inference logs + pass + + def _send(self, code, ctype, body): + self.send_response(code) + self.send_header("Content-Type", ctype) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path in ("/", "/index.html"): + self._send(200, "text/html; charset=utf-8", _PAGE.encode()) + elif self.path == "/state": + self._send(200, "application/json", + json.dumps(worker.snapshot()).encode()) + elif self.path.startswith("/overlay/"): + try: + png = worker.overlay_png(int(self.path.rsplit("/", 1)[1])) + except ValueError: + png = None + if png is None: + self._send(404, "text/plain", b"not ready") + else: + self._send(200, "image/png", png) + else: + self._send(404, "text/plain", b"not found") + + def do_POST(self): + if self.path != "/upload": + return self._send(404, "text/plain", b"not found") + n = int(self.headers.get("Content-Length", 0)) + if not 0 < n <= 64 * 1024 * 1024: + return self._send(413, "text/plain", b"bad size") + from urllib.parse import unquote + name = unquote(self.headers.get("X-Filename", "upload.png")) + path = worker.enqueue_bytes(self.rfile.read(n), name) + self._send(200, "application/json", + json.dumps({"queued": os.path.basename(path)}).encode()) + return Handler + + +def serve_inference(image_dir, port=7860, batch_size=16, **inference_kwargs): + """Start the dashboard + persistent worker. Blocks until Ctrl-C. Scored + entries are in the JSON output (rewritten after every image).""" + inference_kwargs.setdefault("explain", True) # cards always carry the XAI layer + worker = InferenceWorker(image_dir, batch_size=batch_size, **inference_kwargs) + httpd = ThreadingHTTPServer(("0.0.0.0", port), _make_handler(worker)) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + print(f"[serve] dashboard at http://localhost:{port} — {len(worker.queue)} queued, " + f"drop/paste images on the page to add more (Ctrl-C to exit)") + try: + worker.run_forever() + except KeyboardInterrupt: + pass + finally: + httpd.shutdown() + return worker.done diff --git a/src/model/__init__.py b/src/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/model/_dinov2.py b/src/model/_dinov2.py new file mode 100644 index 0000000000000000000000000000000000000000..f31db7fbc95acf2d8e3bfaeb0df55098a40f84e9 --- /dev/null +++ b/src/model/_dinov2.py @@ -0,0 +1,82 @@ +# src/model/_dinov2.py +import torch +import torch.nn as nn + +# Supported hub names; reg = "Vision Transformers Need Registers" variants +# (cleaner patch tokens + attention maps — matters for our per-patch MIL head +# and the attention attribution map). All Apache-2.0. +_VALID = {"dinov2_vitb14", "dinov2_vitb14_reg", "dinov2_vitl14", "dinov2_vitl14_reg"} + + +class DinoV2Backbone(nn.Module): + """Real DINOv2 (downloads weights on first use). Emits the patch-token grid. + Geometry (patch_size, embed_dim, register count) is READ FROM THE MODEL, not + hardcoded — so switching B<->L or reg<->no-reg is a pure cfg.backbone change.""" + def __init__(self, cfg): + super().__init__() + if cfg.backbone not in _VALID: + raise ValueError(f"unknown dinov2 variant {cfg.backbone}; valid: {_VALID}") + self.model = torch.hub.load("facebookresearch/dinov2", cfg.backbone) + self.model.requires_grad_(False) + self.model.eval() + # RINE-style intermediate-layer fusion (opt-in). n_intermediate_layers>1 fuses the + # last-N transformer blocks' patch tokens into the head's [B,N,D] via a small + # TRAINABLE fusion (fusion.py); backbone stays frozen. Targets cross-generator signal. + self.n_inter = int(getattr(cfg, "n_intermediate_layers", 1)) + self._fusion = None + if self.n_inter > 1: + from model.fusion import IntermediateFusion + self._fusion = IntermediateFusion(self.n_inter, int(self.model.embed_dim), + project=bool(getattr(cfg, "fusion_project", True))) + # OPT-IN single intermediate-layer extraction (cfg.extract_layer, 1-indexed). Reads the + # patch tokens AFTER block L instead of the final block. Two effects: (1) tests whether + # spatially-localized signal survives at an earlier layer (the final layer is globally + # homogenized — every patch-logit near-identical, per the XAI audit), and (2) CUTS + # COMPUTE — the block stack is TRUNCATED to L blocks at init, so blocks L+1..12 never run + # (a ~ (12-L)/12 backbone-FLOP saving). 0/None = full depth. Mutually exclusive with the + # RINE fusion path (n_intermediate_layers>1). + self.extract_layer = int(getattr(cfg, "extract_layer", 0) or 0) + if self.extract_layer > 0: + if self.n_inter > 1: + raise ValueError("extract_layer and n_intermediate_layers>1 are mutually exclusive") + L = self.extract_layer + if not (1 <= L <= len(self.model.blocks)): + raise ValueError(f"extract_layer {L} out of range 1..{len(self.model.blocks)}") + # physically drop blocks L+1..N so the forward stops early (real FLOP saving) + self.model.blocks = nn.ModuleList(list(self.model.blocks)[:L]) + # Derive geometry from the loaded model's own attributes. + self.patch_size = int(self.model.patch_size) # 14 + self.embed_dim = int(self.model.embed_dim) # 768 (B) / 1024 (L) + self.num_registers = int(getattr(self.model, "num_register_tokens", 0)) # 0 or 4 + if cfg.resolution % self.patch_size != 0: + raise ValueError(f"resolution {cfg.resolution} not divisible by patch_size {self.patch_size}") + grid = cfg.resolution // self.patch_size # 518//14 = 37 + self.n_patch = grid * grid # 1369 + + # MUST be no_grad (NOT inference_mode): the frozen backbone's output is + # consumed by the trainable head, whose loss.backward() saves these tensors + # for backward. inference_mode marks outputs as "inference tensors" that + # cannot be saved for backward -> RuntimeError in the head's backward pass. + # no_grad avoids that while still skipping autograd bookkeeping in the frozen + # backbone. Matches StubBackbone.forward's decorator by design. + @torch.no_grad() + def _frozen_tokens(self, x: torch.Tensor) -> torch.Tensor: + """Frozen backbone tokens under no_grad. Single layer -> [B,N,D]; fusion -> the + last-N blocks stacked as [K,B,N,D] (patch tokens only, registers excluded).""" + if self.extract_layer > 0: + # blocks truncated to L at init -> the last remaining block IS layer L; norm=True + # applies the final LN so the tokens are on the same scale the head expects. + outs = self.model.get_intermediate_layers(x, n=1, return_class_token=False, norm=True) + return outs[0] # [B, N, D] at layer L + if self.n_inter <= 1: + return self.model.forward_features(x)["x_norm_patchtokens"] + # get_intermediate_layers returns a tuple of the last-N blocks' patch tokens + outs = self.model.get_intermediate_layers(x, n=self.n_inter, return_class_token=False) + return torch.stack(list(outs), dim=0) # [K, B, N, D] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + toks = self._frozen_tokens(x) # no_grad frozen features + if self._fusion is not None: + # fusion is TRAINABLE and lives OUTSIDE no_grad so gradients reach it + return self._fusion(toks) # [K,B,N,D] -> [B,N,D] + return toks # [B, N, D] diff --git a/src/model/backbone.py b/src/model/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..2734153c4532d067d352fc75b8f4d475c6fe1b64 --- /dev/null +++ b/src/model/backbone.py @@ -0,0 +1,34 @@ +# src/model/backbone.py +import torch +import torch.nn as nn + +class StubBackbone(nn.Module): + """Shape-matched stand-in for DINOv2. Fixed random projection of pooled pixels. + Deterministic, CPU-cheap; lets the whole head/loss/eval path be tested without GPU.""" + def __init__(self, embed_dim: int = 768, n_patch: int = 1369): + super().__init__() + self.embed_dim = embed_dim + self.n_patch = n_patch + g = torch.Generator().manual_seed(0) + # project a fixed grid pooling of the image into per-patch features + self.register_buffer("proj", torch.randn(3, embed_dim, generator=g)) + self.grid = int(round(n_patch ** 0.5)) + + # MUST be no_grad (NOT inference_mode) so the frozen backbone's output can be + # consumed by the trainable head's backward. inference_mode outputs are + # "inference tensors" that cannot be saved for backward. DinoV2Backbone shares + # this contract — keep both decorators in sync. + @torch.no_grad() + def forward(self, x: torch.Tensor) -> torch.Tensor: + b = x.shape[0] + pooled = torch.nn.functional.adaptive_avg_pool2d(x, (self.grid, self.grid)) # [B,3,g,g] + pooled = pooled.permute(0, 2, 3, 1).reshape(b, self.grid * self.grid, 3) # [B,N,3] + return pooled @ self.proj # [B,N,D] + +def build_backbone(cfg) -> nn.Module: + if cfg.backbone == "stub": + return StubBackbone(embed_dim=cfg.embed_dim, n_patch=cfg.n_patch) + if cfg.backbone.startswith("dinov2_"): # routes vitb14, vitb14_reg, vitl14, vitl14_reg + from model._dinov2 import DinoV2Backbone # lazy: only imported on GPU hosts + return DinoV2Backbone(cfg) + raise ValueError(f"unknown backbone {cfg.backbone}") diff --git a/src/model/detector.py b/src/model/detector.py new file mode 100644 index 0000000000000000000000000000000000000000..f96b200730a38377f548bba9f71a0e0078d9055c --- /dev/null +++ b/src/model/detector.py @@ -0,0 +1,92 @@ +# src/model/detector.py +import torch +import torch.nn as nn +from model.backbone import build_backbone +from model.head import AttentionPoolPatchHead + +class Detector(nn.Module): + def __init__(self, cfg): + super().__init__() + self.backbone = build_backbone(cfg) + self.backbone.requires_grad_(False) + # Re-enable grad on the OPT-IN intermediate-layer fusion (a small trainable module + # living inside the frozen backbone) — the backbone weights stay frozen, the fusion + # trains. No-op when fusion is off (n_intermediate_layers<=1). + _fusion = getattr(self.backbone, "_fusion", None) + if _fusion is not None: + _fusion.requires_grad_(True) + # Size the head from the BACKBONE's embed_dim (read from the model), not cfg — + # so DINOv2 L/14 (1024) works without a cfg change. StubBackbone.embed_dim == cfg.embed_dim. + _pool = getattr(cfg, "head_pool", "attn") + if _pool == "dist": + from model.head import DistributionalPoolHead + self.head = DistributionalPoolHead(embed_dim=self.backbone.embed_dim, + hidden=getattr(cfg, "head_hidden", 0)) + elif _pool == "nbr": + from model.head import NeighborhoodPoolHead + self.head = NeighborhoodPoolHead(embed_dim=self.backbone.embed_dim, + hidden=getattr(cfg, "head_hidden", 0)) + elif _pool == "enc": + from model.head import EncoderPoolHead + self.head = EncoderPoolHead(embed_dim=self.backbone.embed_dim, + hidden=getattr(cfg, "head_hidden", 0), + enc_layers=getattr(cfg, "enc_layers", 2), + enc_heads=getattr(cfg, "enc_heads", 8), + use_pos_embed=getattr(cfg, "use_pos_embed", True), + mil_top_k=getattr(cfg, "mil_top_k", 0), + mil_centered=getattr(cfg, "mil_centered", False)) + else: + self.head = AttentionPoolPatchHead(embed_dim=self.backbone.embed_dim, + hidden=getattr(cfg, "head_hidden", 0)) + self.branches = nn.ModuleDict() # Phase-2 seam: depth/normal, frequency, ... + # OPT-IN mid-band spectral branch (cfg.use_freq_branch): a small trainable head over a + # corruption-robust radial FFT statistic, added to the image logit. Off by default so + # the shipped model stays a single RGB forward pass with no extra cost. When on, the + # branch is trained alongside the MIL head (backbone still frozen). + self._freq = None + if getattr(cfg, "use_freq_branch", False): + from model.freq import FreqBranch + fb_out = 32 + self._freq = FreqBranch(n_bands=8, out_dim=fb_out) + self._freq_logit = nn.Linear(fb_out, 1) + + # AMP: run the FROZEN backbone in fp16 autocast (T4 tensor cores ~2x matmul). + # Safe split: backbone has no grads (no fp16 gradient-scale concerns); the tiny + # trainable head consumes fp32-cast tokens, so training numerics are unchanged + # in kind. Opt-in via cfg.amp; exact fp32 path when off (clean A/B). + self._amp = bool(getattr(cfg, "amp", False)) + + def forward_features(self, x: torch.Tensor) -> torch.Tensor: + if self._amp and x.is_cuda: + with torch.autocast("cuda", dtype=torch.float16): + tokens = self.backbone(x) + return tokens.float() + return self.backbone(x) + + def forward(self, x: torch.Tensor) -> dict: + tokens = self.forward_features(x) + # Run the head under bf16 autocast when AMP is on. This is load-bearing for the + # encoder head (head_pool="enc"): nn.TransformerEncoderLayer calls SDPA, which only + # dispatches to the flash / mem-efficient kernels for half-precision q/k/v — in fp32 + # it falls back to the math backend and MATERIALIZES the full B*heads*N^2 attention + # matrix (7.7 GB at B=128, N=1369), OOM-ing the 5090. bf16 (not fp16) needs no + # GradScaler and has no softmax-overflow risk; the backbone is frozen so there are no + # fp16-gradient concerns. For the linear/pool heads this is a numerically-negligible + # speedup. Exact fp32 path preserved when amp=False (clean A/B). + if self._amp and tokens.is_cuda: + with torch.autocast("cuda", dtype=torch.bfloat16): + image_logit, patch_logits, attn = self.head(tokens) + image_logit, patch_logits, attn = image_logit.float(), patch_logits.float(), attn.float() + else: + image_logit, patch_logits, attn = self.head(tokens) + if self._freq is not None: + # additive spectral contribution (mid-band, corruption-robust). Additive so the + # ablation is exact: use_freq_branch=False recovers the pure-DINOv2 logit. + image_logit = image_logit + self._freq_logit(self._freq(x)).squeeze(-1) + out = {"image_logit": image_logit, "patch_logits": patch_logits, "attn": attn} + # pooled trainable representation (enc head stashes it) — consumed by the + # feature-consistency loss; float() so bf16-autocast runs still compare in fp32. + lp = getattr(self.head, "last_pooled", None) + if lp is not None: + out["pooled_feat"] = lp.float() + return out diff --git a/src/model/freq.py b/src/model/freq.py new file mode 100644 index 0000000000000000000000000000000000000000..cfc75ea3a22452d40391e1db0f9da8ae399176da --- /dev/null +++ b/src/model/freq.py @@ -0,0 +1,75 @@ +# src/model/freq.py +"""Cheap corruption-ROBUST spectral feature (Tier-2, opt-in). + +The classic AIGC spectral cue lives in the HIGH-frequency band — which jpeg/blur/resize +destroy, so it fails exactly the robust half of the score. Instead we extract a **mid-band** +radial power-spectrum statistic: the ring of spatial frequencies between the very-low +(content) and very-high (fragile) bands. Mid-band structure survives mild corruption while +still carrying some generator signal (esp. for GANs). This is concatenated to the pooled +DINOv2 features via a small trainable projection — it does NOT replace the backbone. + +RISK/GATE: this only earns a place if it measurably separates real/fake WITHOUT collapsing +under corruption. Use freq_separability_auc() as the P1 gate before enabling it in training. +Kept opt-in (cfg.use_freq_branch) so the default single-forward-pass model is unaffected. +""" +import numpy as np +import torch +import torch.nn as nn + + +def radial_midband_features(x, n_bands=8, lo=0.15, hi=0.6): + """Per-image radial mid-band power spectrum → [B, n_bands]. + + x: [B,3,H,W] in [0,1]. Convert to grayscale, 2D FFT, azimuthally average |F|² into radial + rings, keep the mid-band fraction [lo,hi] of the radius, bin into n_bands, log1p-compress. + Deterministic, cheap (one FFT per image), differentiable-friendly (returns a tensor).""" + if x.dim() != 4: + raise ValueError(f"expected [B,3,H,W], got {tuple(x.shape)}") + gray = x.mean(dim=1) # [B,H,W] luminance-ish + B, H, W = gray.shape + fft = torch.fft.fftshift(torch.fft.fft2(gray), dim=(-2, -1)) + power = (fft.real ** 2 + fft.imag ** 2) # [B,H,W] + # radial coordinate per pixel, normalized to [0,1] at the Nyquist corner + yy, xx = torch.meshgrid(torch.arange(H, device=x.device) - H // 2, + torch.arange(W, device=x.device) - W // 2, indexing="ij") + r = torch.sqrt((yy.float() / (H / 2)) ** 2 + (xx.float() / (W / 2)) ** 2) # [H,W] + r = r.clamp(0, 1) + out = [] + edges = torch.linspace(lo, hi, n_bands + 1) + for b in range(n_bands): + m = ((r >= edges[b]) & (r < edges[b + 1])).float() # [H,W] ring mask + denom = m.sum().clamp(min=1.0) + band = (power * m).reshape(B, -1).sum(dim=1) / denom # [B] mean power in ring + out.append(torch.log1p(band)) + feats = torch.stack(out, dim=1) # [B, n_bands] + # per-sample normalize so absolute brightness doesn't dominate + return feats / (feats.norm(dim=1, keepdim=True) + 1e-6) + + +class FreqBranch(nn.Module): + """Trainable projection of the mid-band spectral feature to `out_dim`, to concat with the + pooled DINOv2 features. Tiny (n_bands·out_dim params); frozen backbone stays frozen.""" + def __init__(self, n_bands=8, out_dim=32, lo=0.15, hi=0.6): + super().__init__() + self.n_bands, self.lo, self.hi = n_bands, lo, hi + self.proj = nn.Sequential(nn.Linear(n_bands, out_dim), nn.GELU()) + + def forward(self, x): + feats = radial_midband_features(x, self.n_bands, self.lo, self.hi) + return self.proj(feats) + + +def freq_separability_auc(imgs_real, imgs_fake, n_bands=8): + """P1 GATE: does the raw mid-band feature separate real from fake at all? Fit a quick + logistic on the n_bands feature and report CV AUC. Run this BEFORE enabling the branch in + training; if AUC ≈ 0.5 the feature is noise and should be cut. imgs_* are [N,3,H,W].""" + from sklearn.linear_model import LogisticRegression + from sklearn.model_selection import cross_val_predict + from sklearn.metrics import roc_auc_score + fr = radial_midband_features(imgs_real, n_bands).cpu().numpy() + ff = radial_midband_features(imgs_fake, n_bands).cpu().numpy() + X = np.concatenate([fr, ff]); y = np.r_[np.zeros(len(fr)), np.ones(len(ff))] + if len(np.unique(y)) < 2: + return 0.5 + p = cross_val_predict(LogisticRegression(max_iter=500), X, y, cv=3, method="predict_proba")[:, 1] + return float(roc_auc_score(y, p)) diff --git a/src/model/fusion.py b/src/model/fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..e1c0ed6ff498feddd11e9378e7d82cff6f3f6755 --- /dev/null +++ b/src/model/fusion.py @@ -0,0 +1,38 @@ +# src/model/fusion.py +"""RINE-style intermediate-layer fusion. + +The final DINOv2 layer emphasizes semantics; INTERMEDIATE ViT layers retain more localized +forensic signal (RINE, arXiv 2402.19091, reports large cross-generator gains from fusing +intermediate layers with a light trainable head — exactly our frozen-backbone / few-epoch +regime). This module fuses a stack of K frozen layer outputs into the single [B,N,D] token +grid the attention-pooling MIL head already consumes, so the head interface is unchanged and +only the small fusion is trainable — no backbone unfreezing (which would help random-split +but hurt the generator-held-out split that actually matters). +""" +import torch +import torch.nn as nn + + +class IntermediateFusion(nn.Module): + """Fuse K per-layer token grids [K, B, N, D] -> [B, N, D]. + + A learnable softmax over the K layers (a weighted combination) followed by an optional + per-token linear projection. Lightweight (K + D*D params) and fully trainable while the + backbone stays frozen.""" + def __init__(self, n_layers: int, embed_dim: int, project: bool = True): + super().__init__() + self.n_layers = n_layers + self.embed_dim = embed_dim + # learnable mixing weight per layer (softmax-normalized at forward) + self.layer_logits = nn.Parameter(torch.zeros(n_layers)) + self.proj = nn.Linear(embed_dim, embed_dim) if project else None + + def forward(self, layers: torch.Tensor) -> torch.Tensor: + # layers: [K, B, N, D] + if layers.dim() != 4 or layers.shape[0] != self.n_layers: + raise ValueError(f"expected [{self.n_layers}, B, N, D], got {tuple(layers.shape)}") + w = torch.softmax(self.layer_logits, dim=0).view(-1, 1, 1, 1) # [K,1,1,1] + fused = (layers * w).sum(dim=0) # [B, N, D] + if self.proj is not None: + fused = self.proj(fused) + return fused diff --git a/src/model/head.py b/src/model/head.py new file mode 100644 index 0000000000000000000000000000000000000000..72f164d5e60258c016adbf6a7e1aa1f2dde79ea0 --- /dev/null +++ b/src/model/head.py @@ -0,0 +1,272 @@ +# src/model/head.py +import math + +import torch +import torch.nn as nn + + +def _sincos_2d_pos_embed(dim: int, grid: int, device, dtype): + """Standard fixed 2D sin/cos positional embedding for a grid x grid token map. + + Returns [grid*grid, dim] in row-major (the same order tokens are flattened by + grid.reshape). Half the channels encode the row coordinate, half the column, so a + token's absolute (row, col) is recoverable — this is the position information the + pooling heads never had. Fixed (not learned) so it adds ZERO parameters and cannot + itself become a per-generator shortcut; only the encoder that consumes it is learned. + """ + assert dim % 4 == 0, "sincos-2d needs dim divisible by 4" + d_each = dim // 2 # channels for row, channels for col + omega = torch.arange(d_each // 2, device=device, dtype=torch.float32) + omega = 1.0 / (10000 ** (omega / (d_each / 2))) # [d_each//2] + coords = torch.arange(grid, device=device, dtype=torch.float32) + # outer products -> [grid, d_each//2] + row = torch.arange(grid, device=device).repeat_interleave(grid).float() # [N] row idx + col = torch.arange(grid, device=device).repeat(grid).float() # [N] col idx + def _emb(pos): # [N] -> [N, d_each] + ang = pos[:, None] * omega[None, :] # [N, d_each//2] + return torch.cat([torch.sin(ang), torch.cos(ang)], dim=1) + pe = torch.cat([_emb(row), _emb(col)], dim=1) # [N, dim] + return pe.to(dtype) + + +class EncoderPoolHead(nn.Module): + """DistributionalPoolHead preceded by a REAL self-attention encoder over the tokens, + WITH fixed 2D positional embeddings — the "sequence, not bag" test. + + Every other head (attn/dist/nbr) is permutation-invariant: the "attention" is a softmax + over INDEPENDENT per-token scores (a scalar gate), never token-to-token mixing. The + feature-trace proved this (logit changes <=3.8e-6 under token permutation), so those heads + are blind to CONFIGURAL structure (six fingers, fused bodies) — impossibility that lives in + the ARRANGEMENT across tokens, not in any one token. This head inserts an nn.TransformerEncoder + (true self-attention) so a token's contribution can depend on the whole set AND on absolute + position (via the 2D PE), then runs the same attn + mean/std/top-k distributional pooling. + + ZERO-INIT RESIDUAL: the encoder's output projection is zeroed, so at init the encoder is an + exact identity (tokens pass through unchanged) and the head is EXACTLY DistributionalPoolHead + over the raw tokens — the same clean A/B as nbr/dist. Any gain is attributable to the learned + self-attention+position term. The PE is added only INSIDE the residual branch, so it + contributes nothing at init and cannot shift the baseline; the encoder alone is the new + variable. Requires N to be a perfect square (the DINO patch grid); falls back to no mixing if + not (keeps the head safe on the stub backbone / odd token counts). + + Same (image_logit, patch_logits, attn) interface — attn stays the XAI map. Deployability: one + extra small transformer over ~1369 tokens, still a single backbone forward.""" + def __init__(self, embed_dim: int = 768, hidden: int = 0, top_k: int = 16, + enc_layers: int = 2, enc_heads: int = 8, use_pos_embed: bool = True, + mil_top_k: int = 0, mil_centered: bool = False): + super().__init__() + self.use_pos_embed = bool(use_pos_embed) + # MIL mode (mil_top_k > 0): image_logit = mean of the top-k patch logits, PERIOD — no + # attn-pool, no learned combine. The L8 XAI audit showed the borderline misses have + # strong localized fake evidence (max+ patches +2..+5) diluted to ~0.5 by global + # pooling, while the learned `combine` never moved off the diluting stats. MIL makes + # the evidence map load-bearing: an image is fake iff k patches SAY fake, so the + # patch-logit heatmap is faithful BY CONSTRUCTION (it literally is the decision), and + # a real image must contain no fake-looking region anywhere. + self.mil_top_k = int(mil_top_k) + # Centered MIL: plain top-k over globally-informed patch logits collapses back into + # global pooling — the encoder routes the image verdict into EVERY patch (the + # "washed red/blue" maps; milk16 GigaGAN 0.9418 ≈ baseline). Centering kills that + # channel exactly: top-k(patch − mean) is invariant to a uniform tint, so the + # localized term can only grow via patches deviating from the image's own baseline. + # The global signal is kept in a SEPARATE mean term: logit = w1·mean + w2·topk_dev + # (w learned; init 1,1 = both channels live). XAI: verdict decomposes into an + # explicit global + localized split, and the centered map shows relative structure. + self.mil_centered = bool(mil_centered) + if self.mil_centered: + self.mil_combine = nn.Linear(2, 1) + with torch.no_grad(): + self.mil_combine.weight.copy_(torch.tensor([[1.0, 1.0]])) + self.mil_combine.bias.zero_() + layer = nn.TransformerEncoderLayer( + d_model=embed_dim, nhead=enc_heads, dim_feedforward=embed_dim * 2, + dropout=0.0, batch_first=True, norm_first=True) + self.encoder = nn.TransformerEncoder(layer, num_layers=enc_layers, + enable_nested_tensor=False) + self.enc_out = nn.Linear(embed_dim, embed_dim) # residual projection, zero-init below + with torch.no_grad(): # zero residual -> starts as pure dist head + self.enc_out.weight.zero_() + self.enc_out.bias.zero_() + if hidden and hidden > 0: + self.patch_head = nn.Sequential( + nn.Linear(embed_dim, hidden), nn.GELU(), nn.Linear(hidden, 1)) + else: + self.patch_head = nn.Linear(embed_dim, 1) + self.attn_score = nn.Linear(embed_dim, 1) + self.top_k = int(top_k) + self.combine = nn.Linear(4, 1) + with torch.no_grad(): # pass-through of the attn-pooled logit + self.combine.weight.copy_(torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + self.combine.bias.zero_() + + def _encode(self, tokens): + """tokens [B,N,D] -> self-attention-enriched [B,N,D] (zero-init residual). No-op if N!=G*G.""" + b, n, d = tokens.shape + g = int(round(n ** 0.5)) + if g * g != n: + return tokens # non-grid token count: skip mixing + x = tokens + if self.use_pos_embed: + pe = _sincos_2d_pos_embed(d, g, tokens.device, tokens.dtype) # [N,D] + x = x + pe[None, :, :] + enc = self.encoder(x) # true token-to-token self-attention + return tokens + self.enc_out(enc) # zero-init residual on the RAW tokens + + def forward(self, tokens: torch.Tensor): + tokens = self._encode(tokens) + # Stash the pooled TRAINABLE representation (post-encoder token mean) for the optional + # feature-consistency loss (cfg.lambda_feat): penalizing its drift across corrupted + # variants forces the decision path onto corruption-stable (semantic/geometric) + # feature directions — fine structure can't carry the verdict if it doesn't survive + # the aug pipeline. Stashed (not returned) to keep the 3-tuple interface stable. + self.last_pooled = tokens.mean(dim=1) # [B,D] + patch_logits = self.patch_head(tokens).squeeze(-1) # [B,N] + attn = torch.softmax(self.attn_score(tokens).squeeze(-1), dim=1) # [B,N] + if self.mil_top_k > 0: + k = min(self.mil_top_k, patch_logits.shape[1]) + if self.mil_centered: + mean = patch_logits.mean(dim=1) # global channel + dev = patch_logits - mean[:, None] # DC-free deviations + topk_dev = dev.topk(k, dim=1).values.mean(dim=1) # localized channel + image_logit = self.mil_combine( + torch.stack([mean, topk_dev], dim=1)).squeeze(-1) + else: + image_logit = patch_logits.topk(k, dim=1).values.mean(dim=1) + return image_logit, patch_logits, attn + pooled = (attn * patch_logits).sum(dim=1) + mean = patch_logits.mean(dim=1) + std = patch_logits.std(dim=1) + k = min(self.top_k, patch_logits.shape[1]) + topk = patch_logits.topk(k, dim=1).values.mean(dim=1) + stats = torch.stack([pooled, mean, std, topk], dim=1) + image_logit = self.combine(stats).squeeze(-1) + return image_logit, patch_logits, attn + + +class NeighborhoodPoolHead(nn.Module): + """DistributionalPoolHead + an explicit LOCAL-SPATIAL stage before scoring. + + The feature-trace audit proved the pooling heads are pure SET operators (logit changes + <=3.8e-6 under token permutation) even though DINO's final-layer tokens are spatially + coherent — they cannot see adjacency discontinuities (resampling seams, fused + boundaries), exactly where GAN/upsampling artifacts live. This head reshapes the token + grid and applies a depthwise 3x3 conv so each patch feature is enriched by its 8 + neighbours BEFORE the per-patch fakeness scorer, then runs the same attention + + mean/std/top-k distributional pooling as DistributionalPoolHead. + + The conv is a ZERO-INITIALISED RESIDUAL (tokens + conv(tokens), conv weights/bias 0), + so at init the head is EXACTLY DistributionalPoolHead over the raw tokens — a clean A/B + where any gain is attributable to the learned neighbourhood term. Requires N to be a + perfect square (the DINO patch grid, e.g. 1369 = 37x37); falls back to no spatial mixing + if it is not (keeps the head safe on the stub backbone / odd token counts). + + Same (image_logit, patch_logits, attn) interface — attn stays the XAI map, deployability + (single forward, tiny head) unchanged.""" + def __init__(self, embed_dim: int = 768, hidden: int = 0, top_k: int = 16): + super().__init__() + # depthwise 3x3 over the grid: embed_dim groups -> per-channel local mixing, cheap + # (9*embed_dim params) and it cannot blend channels, only neighbours within a channel. + self.spatial = nn.Conv2d(embed_dim, embed_dim, kernel_size=3, padding=1, + groups=embed_dim, bias=True) + with torch.no_grad(): # zero residual -> starts as pure dist head + self.spatial.weight.zero_() + self.spatial.bias.zero_() + if hidden and hidden > 0: + self.patch_head = nn.Sequential( + nn.Linear(embed_dim, hidden), nn.GELU(), nn.Linear(hidden, 1)) + else: + self.patch_head = nn.Linear(embed_dim, 1) + self.attn_score = nn.Linear(embed_dim, 1) + self.top_k = int(top_k) + self.combine = nn.Linear(4, 1) + with torch.no_grad(): # pass-through of the attn-pooled logit + self.combine.weight.copy_(torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + self.combine.bias.zero_() + + def _mix(self, tokens): + """tokens [B,N,D] -> neighbourhood-enriched [B,N,D] (residual). No-op if N != G*G.""" + b, n, d = tokens.shape + g = int(round(n ** 0.5)) + if g * g != n: + return tokens # non-grid token count: skip spatial mixing + grid = tokens.transpose(1, 2).reshape(b, d, g, g) # [B,D,G,G] + mixed = self.spatial(grid) + return tokens + mixed.reshape(b, d, n).transpose(1, 2) + + def forward(self, tokens: torch.Tensor): + tokens = self._mix(tokens) + patch_logits = self.patch_head(tokens).squeeze(-1) # [B,N] + attn = torch.softmax(self.attn_score(tokens).squeeze(-1), dim=1) # [B,N] + pooled = (attn * patch_logits).sum(dim=1) + mean = patch_logits.mean(dim=1) + std = patch_logits.std(dim=1) + k = min(self.top_k, patch_logits.shape[1]) + topk = patch_logits.topk(k, dim=1).values.mean(dim=1) + stats = torch.stack([pooled, mean, std, topk], dim=1) + image_logit = self.combine(stats).squeeze(-1) + return image_logit, patch_logits, attn + +class AttentionPoolPatchHead(nn.Module): + """Shared per-patch fakeness scorer + learned attention pooling (MIL). + Attention weights double as the free attribution map for XAI/localization. + + hidden=0 (default) keeps the original 1.5K-param LINEAR probe (Linear(D,1)). hidden>0 uses + a small MLP per-patch scorer (Linear(D,hidden)→GELU→Linear(hidden,1)) — more capacity on + the frozen features, ablatable against the linear probe. The attention branch stays linear + so the attribution map keeps its simple interpretation.""" + def __init__(self, embed_dim: int = 768, hidden: int = 0): + super().__init__() + if hidden and hidden > 0: + self.patch_head = nn.Sequential( + nn.Linear(embed_dim, hidden), nn.GELU(), nn.Linear(hidden, 1)) + else: + self.patch_head = nn.Linear(embed_dim, 1) + self.attn_score = nn.Linear(embed_dim, 1) + + def forward(self, tokens: torch.Tensor): + patch_logits = self.patch_head(tokens).squeeze(-1) # [B,N] + attn = torch.softmax(self.attn_score(tokens).squeeze(-1), dim=1) # [B,N] + image_logit = (attn * patch_logits).sum(dim=1) # [B] + return image_logit, patch_logits, attn + + +class DistributionalPoolHead(nn.Module): + """Attention pooling + DISTRIBUTIONAL summaries of the patch-logit population + concat [attn-pooled, mean, std, top-k mean] of the per-patch + fakeness logits, then a tiny linear combiner produces the image logit. + + Rationale: attention pooling alone answers "is there a suspicious REGION?"; different + generator families leave differently-SHAPED patch-evidence distributions (GAN texture + artifacts = many mid-strength patches; diffusion semantic slips = few strong ones). + The population stats expose that shape for ~4 extra params, targeting cross-generator + generalization (our weak axis) with negligible capacity/overfit risk. + + Same interface/attribution map as AttentionPoolPatchHead — attn stays the XAI map, + and the deployability story (single forward, tiny head) is unchanged.""" + def __init__(self, embed_dim: int = 768, hidden: int = 0, top_k: int = 16): + super().__init__() + if hidden and hidden > 0: + self.patch_head = nn.Sequential( + nn.Linear(embed_dim, hidden), nn.GELU(), nn.Linear(hidden, 1)) + else: + self.patch_head = nn.Linear(embed_dim, 1) + self.attn_score = nn.Linear(embed_dim, 1) + self.top_k = int(top_k) + self.combine = nn.Linear(4, 1) + # init: pass-through of the attn-pooled logit -> starts EXACTLY as the baseline + # head, so the stats terms can only be learned INTO relevance (stable A/B). + with torch.no_grad(): + self.combine.weight.copy_(torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + self.combine.bias.zero_() + + def forward(self, tokens: torch.Tensor): + patch_logits = self.patch_head(tokens).squeeze(-1) # [B,N] + attn = torch.softmax(self.attn_score(tokens).squeeze(-1), dim=1) # [B,N] + pooled = (attn * patch_logits).sum(dim=1) # [B] + mean = patch_logits.mean(dim=1) + std = patch_logits.std(dim=1) + k = min(self.top_k, patch_logits.shape[1]) + topk = patch_logits.topk(k, dim=1).values.mean(dim=1) + stats = torch.stack([pooled, mean, std, topk], dim=1) # [B,4] + image_logit = self.combine(stats).squeeze(-1) # [B] + return image_logit, patch_logits, attn diff --git a/src/train/__init__.py b/src/train/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/train/checkpoint.py b/src/train/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..5d2f97ddf8020d30f3d55558e927ee11c5d89ae7 --- /dev/null +++ b/src/train/checkpoint.py @@ -0,0 +1,106 @@ +# src/train/checkpoint.py +"""Checkpoint bundle: everything needed to reconstruct the detector for inference. + +A bundle is a single torch.save() dict with the trainable head weights, the full +cfg (as a plain dict), the backbone id, and a best-effort git sha for provenance. +The frozen backbone weights are NOT stored — they are re-fetched from the backbone +id (dinov2 download on the GPU host / deterministic stub locally), so bundles stay small. +""" +import dataclasses +import subprocess + +import torch + +from config import Config +from model.detector import Detector + + +def _git_sha() -> str: + """Best-effort short git sha; 'unknown' if git is unavailable / not a repo.""" + try: + return ( + subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + stderr=subprocess.DEVNULL, + ) + .decode() + .strip() + ) + except Exception: # noqa: BLE001 — provenance is nice-to-have, never fatal + return "unknown" + + +def save_checkpoint(model, cfg, path) -> str: + """Serialize the head weights + cfg + provenance to `path`. Returns `path`.""" + torch.save( + { + "head_state_dict": model.head.state_dict(), + "cfg": dataclasses.asdict(cfg), + "backbone_id": cfg.backbone, + "git_sha": _git_sha(), + }, + path, + ) + return path + + +def _build_member(member_dict): + """(Detector, Config) from one member dict {head_state_dict, cfg, ...}.""" + cfg_dict = member_dict.get("cfg", {}) + known = {f.name for f in dataclasses.fields(Config)} + cfg = Config(**{k: v for k, v in cfg_dict.items() if k in known}) + model = Detector(cfg) + model.head.load_state_dict(member_dict["head_state_dict"]) + model.eval() + return model, cfg + + +def load_checkpoint(path, map_location="cpu"): + """Load a bundle -> (Detector, Config), head weights restored, eval-ready. + + cfg is rebuilt from ONLY the known Config fields, so a bundle written by a + future version with extra keys (or missing a newly-added key) still loads. + Building the Detector re-instantiates the backbone from cfg.backbone. + + A multi-member ensemble bundle (see save_ensemble_checkpoint) loads as its + FIRST member here; use load_ensemble_checkpoint to get every member. + """ + bundle = torch.load(path, map_location=map_location, weights_only=False) + if "members" in bundle: + return _build_member(bundle["members"][0]) + return _build_member(bundle) + + +def save_ensemble_checkpoint(models_cfgs, path) -> str: + """Serialize N (model, cfg) pairs into ONE bundle file (the ship format). + + weights.pt holds {"members": [{head_state_dict, cfg, backbone_id}, ...]}. + All members share the frozen backbone id; only the small head weights and + per-member cfgs (extract_layer differs) are stored. + """ + torch.save( + { + "members": [ + { + "head_state_dict": m.head.state_dict(), + "cfg": dataclasses.asdict(c), + "backbone_id": c.backbone, + } + for m, c in models_cfgs + ], + "git_sha": _git_sha(), + }, + path, + ) + return path + + +def load_ensemble_checkpoint(path, map_location="cpu"): + """Load a bundle -> list of (Detector, Config), one per member. + + A single-model bundle (save_checkpoint format) returns a one-element list, + so callers can treat every checkpoint as an ensemble. + """ + bundle = torch.load(path, map_location=map_location, weights_only=False) + members = bundle.get("members", [bundle]) + return [_build_member(m) for m in members] diff --git a/src/train/compile_backbone.py b/src/train/compile_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..0da9b4cc71959b510612042b536848bf148daaab --- /dev/null +++ b/src/train/compile_backbone.py @@ -0,0 +1,21 @@ +# src/train/compile_backbone.py +import torch + + +def maybe_compile(model, cfg): + if not cfg.compile: + return model + if not torch.cuda.is_available(): + print("[compile] no CUDA — running eager") + return model + try: + # dynamic=True: random_resolution feeds 19 distinct shapes (392..644 step 14); + # static max-autotune would RECOMPILE per shape (minutes each, silently eaten + # as "training time"). Dynamic compile trades a little kernel quality for one + # compilation that serves every resolution. + dynamic = bool(getattr(cfg, "random_resolution", False)) or None + model.backbone = torch.compile(model.backbone, mode="max-autotune", dynamic=dynamic) + print(f"[compile] backbone compiled (max-autotune, dynamic={dynamic})") + except Exception as e: # noqa: BLE001 + print(f"[compile] failed ({e}); running eager") + return model diff --git a/src/train/driver.py b/src/train/driver.py new file mode 100644 index 0000000000000000000000000000000000000000..5b52923386b14473c89cc324b03a8d7112a3234f --- /dev/null +++ b/src/train/driver.py @@ -0,0 +1,205 @@ +# src/train/driver.py +"""Multi-epoch training driver: trains, evaluates on val EACH EPOCH, and streams +per-epoch + per-generator (stratified) AUC to the logger so metrics appear LIVE on +wandb during the run (not just at the end). Optionally runs a final OOS-holdout eval. + +This is the orchestration kernel-2 calls; the per-step/per-batch mechanics live in +loop.py and the eval mechanics in eval/harness.py — this only sequences them + logs. +""" +import dataclasses +import time + +import torch +from torch.utils.data import DataLoader + +from model.detector import Detector +from data.dataset import AigcDataset +from train.loop import train_one_epoch, build_predict_fn +from train.compile_backbone import maybe_compile +from eval.harness import evaluate +from eval.metrics import stratified_auc + + +def _log(logger, metrics): + if logger is not None: + logger.log(metrics) + + +def _label_balanced_cap(df, cap, seed): + """Subsample `df` to ~`cap` rows, balanced across LABELS first, then stratified + across each label's generator groups. + + The earlier `cap // (nunique_generator * 2)` formula assumed every generator appears + under both labels. It does not: fakes spread across many generator_ids while reals + ALL collapse to generator_id="none" (one group). That gave reals a single group's + budget while fakes got one-per-generator, so a cap of 1500 yielded ~750 rows with the + real class under-sampled ~6:1 — and any Youden-J threshold calibrated on it was skewed. + Splitting the budget by label first keeps the class ratio ~1:1 regardless of how many + generators each class carries.""" + n_labels = max(1, df["label"].nunique()) + per_label = max(1, cap // n_labels) + keep = [] + for _label, ldf in df.groupby("label", sort=False): + n_gen = max(1, ldf["generator_id"].nunique()) + per_gen = max(1, per_label // n_gen) + for _gen, g in ldf.groupby("generator_id", sort=False): + keep.extend(g.sample(min(len(g), per_gen), random_state=seed).index) + return df.loc[keep].reset_index(drop=True) + + +def train_model(cfg, train_df, val_df, holdout_df=None, logger=None, device=None, + resolver=None, epoch_ckpt_dir=None): + """Train for cfg.epochs, evaluating val each epoch (live logging). Returns + {model, val_report, holdout_report}. holdout_report is None unless holdout_df given. + + `device` pins the model to a specific GPU (e.g. 'cuda:0' / 'cuda:1'). A dual-GPU instance + has TWO T4s, so two independent training arms (e.g. a base vs fusion A/B) can run CONCURRENTLY, + one per GPU, nearly halving A/B wall-clock — the arms share nothing, so no gradient sync is + needed (this is the right parallelism here, not DataParallel). Defaults to 'cuda' (device 0) + when CUDA is available, else CPU.""" + # resolver (e.g. ShardResolver) lets training read images straight from mounted + # ZIP shards — rows with ref_type='shard' resolve via shard_name+member_name. + ds = AigcDataset(train_df, cfg, train=True, resolver=resolver) + # shuffle=True: corpus is materialized per-generator; unshuffled feeds homogeneous + # runs of one generator. drop_last keeps batch stats stable. Route through make_loader + # (NOT a bare DataLoader) so the random_resolution collate actually wires in — building the + # DataLoader directly here silently dropped random_resolution in every run. + from train.loop import make_loader + loader = make_loader(ds, cfg, shuffle=True, drop_last=True) + model = Detector(cfg) + model = maybe_compile(model, cfg) # compiles model.backbone IN PLACE + if device is not None: + model = model.to(device) + elif torch.cuda.is_available(): + model = model.cuda() + # DATA-PARALLEL across all visible GPUs (e.g. 2x T4). The frozen DINOv2 backbone does a + # full 86M-param ViT forward (the dominant compute) — DP splits THAT across both cards for ~2x, + # while only the tiny trainable head/fusion grads are synced (negligible). Also halves per-GPU + # batch memory (16->8/GPU), removing the res-644 OOM risk. `base` stays the unwrapped Detector + # for checkpoint/predict (DP would prefix state_dict keys with 'module.'); we TRAIN on `train_model_ref`. + base = model + train_ref = model + if getattr(cfg, "data_parallel", False) and torch.cuda.is_available() and torch.cuda.device_count() > 1: + train_ref = torch.nn.DataParallel(model) + print(f"[data_parallel] training across {torch.cuda.device_count()} GPUs", flush=True) + opt = torch.optim.Adam([p for p in base.parameters() if p.requires_grad], lr=cfg.lr) + # MID-EPOCH LR scheduling (per-step OneCycle): with few epochs on a small corpus, a + # per-step warmup+cosine-decay converges faster and more stably than a fixed LR. Opt-in + # via cfg.lr_schedule="onecycle"; stepped every batch inside train_one_epoch. + sched = None + # fractional epochs (cfg.epoch_frac<1): cap batches per epoch; with shuffle=True the + # half-pass is a random subsample. OneCycle total_steps must match the CAPPED count + # or the schedule never leaves warmup. + frac = float(getattr(cfg, "epoch_frac", 1.0)) + max_batches = None if frac >= 1.0 else max(1, int(len(loader) * frac)) + steps_per_epoch = max_batches if max_batches is not None else len(loader) + total_steps = cfg.epochs * steps_per_epoch + # OneCycle needs enough steps to have distinct warmup+anneal phases; on a tiny corpus + # (few steps) it div-by-zeros. Require >=10 total steps, else skip (fixed LR is fine there). + if getattr(cfg, "lr_schedule", "none") == "onecycle" and total_steps >= 10: + sched = torch.optim.lr_scheduler.OneCycleLR( + opt, max_lr=cfg.lr, total_steps=total_steps, pct_start=0.25) + predict_fn = build_predict_fn(base, cfg) # eval on the UNWRAPPED model (small batches) + + # PER-EPOCH VAL SUBSAMPLE (cfg.val_eval_cap>0): the full 47k-corpus val is 7.5k + # images x 19 conditions = hours of CPU-bound corruption PER EPOCH — it stalled the + # first big run for 3.5h+ inside epoch-0 eval. A stratified per-(label,generator) + # subsample gives the same AUC signal at a fraction of the wall; the FINAL model can + # still be scored on the full val offline from the checkpoint. + val_eval_df = val_df + vcap = getattr(cfg, "val_eval_cap", 0) + if vcap and len(val_df) > vcap: + val_eval_df = _label_balanced_cap(val_df, vcap, cfg.seed) + print(f"[val_eval_cap] per-epoch val subsampled {len(val_df)} -> {len(val_eval_df)}", + flush=True) + + from train.watchdog import Watchdog + wd = Watchdog(logger=logger, interval=5).start() + # CROSS-PHASE PREFETCH (user-directed): each epoch's iterator is created EARLY — + # right after the PREVIOUS epoch's training, before its val eval — so the workers + # spawn + prefetch through val's wall time and epoch N+1's first batch is ready the + # moment val finishes (kills the visible epoch-start stall). Correct by construction: + # ds.set_epoch(N+1) runs BEFORE iter() forks the workers (same ordering as today, + # just earlier), and val eval never touches the train dataset. + def _spawn_iter(epoch): + ds.set_epoch(epoch) # vary per-image aug realization (label-blind) BEFORE fork + return iter(loader) + next_iter = _spawn_iter(0) + val_report = None + for epoch in range(cfg.epochs): + wd.set_phase(f"train epoch {epoch}") + t0 = time.time() + train_one_epoch(train_ref, next_iter, opt, cfg, epoch=epoch, logger=logger, + scheduler=sched, max_batches=max_batches) + # spawn NEXT epoch's workers now; they prefetch while the val eval runs below + next_iter = _spawn_iter(epoch + 1) if epoch + 1 < cfg.epochs else None + # --- per-epoch val eval (LIVE) --- + wd.set_phase(f"val eval epoch {epoch}") + val_report = evaluate(predict_fn, val_eval_df, cfg, logger=None, resolver=resolver, + watchdog=wd) # logger=None: we log below + preds = val_report["predictions"] + clean = preds[preds["condition"] == "clean"] + epoch_metrics = { + "val/epoch": epoch, + "val/auc_clean": val_report["auc_clean"], + "val/auc_robust_mean": val_report["auc_robust_mean"], + "val/auc_robust_worst": val_report["auc_robust_worst"], + "val/final": val_report["final"], + "val/epoch_wall_s": time.time() - t0, + } + # per-generator stratified AUC (clean), keyed val/gen/ + for gen, auc in stratified_auc(clean, by="generator_id").items(): + epoch_metrics[f"val/gen/{gen}"] = auc + # split-slice AUC if the val carries the hybrid split tag AND a shared id column + if "split" in val_df.columns and "img_id" in val_df.columns and "img_id" in clean.columns: + for slice_name in ("val_grouped", "val_random"): + mask = clean["img_id"].isin( + val_df[val_df["split"] == slice_name]["img_id"]) + sub = clean[mask] + if len(sub) and sub["label"].nunique() == 2: + from eval.metrics import roc_auc + epoch_metrics[f"val/slice/{slice_name}_auc_clean"] = roc_auc( + sub["label"].to_numpy(), sub["prob"].to_numpy()) + _log(logger, epoch_metrics) + # PER-EPOCH SNAPSHOT (opt-in): checkpoint + val predictions after EVERY epoch, + # so a multi-epoch run preserves each epoch's model/metrics — the 1-epoch answer + # comes free from a 2-epoch run, and a crash mid-epoch-2 loses nothing. + if epoch_ckpt_dir is not None: + import os + from train.checkpoint import save_checkpoint + os.makedirs(epoch_ckpt_dir, exist_ok=True) + save_checkpoint(base, cfg, os.path.join(epoch_ckpt_dir, f"epoch{epoch}.pt")) + val_report["predictions"].to_parquet( + os.path.join(epoch_ckpt_dir, f"epoch{epoch}_val_predictions.parquet"), + index=False) + + # --- final OOS holdout eval (once) --- + holdout_report = None + if holdout_df is not None and len(holdout_df): + # The holdout can be large (12k imgs x 19 conditions = CPU-bound corruption on 2 + # cores -> ~1h). A stratified subsample gives the same AUC signal far cheaper. + # cfg.holdout_eval_cap=0 -> use all; >0 -> cap per (label, generator) group. + cap = getattr(cfg, "holdout_eval_cap", 0) + if cap and len(holdout_df) > cap: + # label-balanced then per-generator stratified (see _label_balanced_cap): a bare + # per-(label,generator) split under-samples the single-group real class. + holdout_df = _label_balanced_cap(holdout_df, cap, cfg.seed) + # logger passed through: the holdout eval is the longest silent phase (~9x the + # capped val) — its eval/progress stream is the difference between monitorable + # and apparently-hung. + wd.set_phase("holdout eval") + holdout_report = evaluate(predict_fn, holdout_df, cfg, logger=logger, resolver=resolver, + watchdog=wd) + hclean = holdout_report["predictions"] + hclean = hclean[hclean["condition"] == "clean"] + hmetrics = { + "holdout/auc_clean": holdout_report["auc_clean"], + "holdout/auc_robust_mean": holdout_report["auc_robust_mean"], + "holdout/final": holdout_report["final"], + } + for gen, auc in stratified_auc(hclean, by="generator_id").items(): + hmetrics[f"holdout/gen/{gen}"] = auc + _log(logger, hmetrics) + + wd.stop() + return {"model": base, "val_report": val_report, "holdout_report": holdout_report} diff --git a/src/train/logging_utils.py b/src/train/logging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..99816f3f122760c9e6ffe1b4b7a88d9e360c116c --- /dev/null +++ b/src/train/logging_utils.py @@ -0,0 +1,114 @@ +# src/train/logging_utils.py +"""Dependency-optional Weights & Biases logger + hardware diagnostics. + +Design invariants: +- Logging NEVER crashes training/eval/smoke. Every wandb / psutil / pynvml use is + guarded by try/except with a graceful no-op fallback. +- With cfg.use_wandb=False (the default) or wandb unavailable, RunLogger is a pure + no-op that only keeps an in-memory buffer of logged dicts (for tests + a final + summary print). No network is ever touched. +""" +import collections + + +class RunLogger: + """Thin wrapper over wandb.run. No-op when wandb is off/unavailable. + + Always keeps the last `buffer_size` logged metric dicts in `self.history` so + tests and a final summary print can inspect what was logged without wandb. + """ + + def __init__(self, cfg, run_config: dict | None = None, buffer_size: int = 200): + self.run = None + self.history: collections.deque = collections.deque(maxlen=buffer_size) + if getattr(cfg, "use_wandb", False): + try: + import wandb + + self.run = wandb.init( + project=getattr(cfg, "wandb_project", "track5-aigc"), + name=(getattr(cfg, "wandb_run_name", "") or None), + config=run_config or {}, + ) + except Exception as e: # noqa: BLE001 — never crash training on log init + print(f"[RunLogger] wandb init failed ({e!r}); logging disabled.") + self.run = None + + def log(self, metrics: dict, step: int | None = None) -> None: + # Record locally first so the buffer is populated even in no-op mode. + self.history.append(dict(metrics)) + if self.run is not None: + try: + self.run.log(metrics, step=step) + except Exception as e: # noqa: BLE001 — never crash training on a log call + print(f"[RunLogger] wandb.log failed ({e!r}); continuing.") + + def finish(self) -> None: + if self.run is not None: + try: + self.run.finish() + except Exception as e: # noqa: BLE001 + print(f"[RunLogger] wandb.finish failed ({e!r}); ignoring.") + finally: + self.run = None + + +def system_stats(device) -> dict: + """Hardware diagnostics, robust to CPU-only (no CUDA) environments. + + Returns a flat dict under the `sys/` namespace. On CPU-only machines only the + psutil RAM/CPU keys are present; GPU keys appear only when CUDA is available. + Never raises — a missing dependency or GPU-stat API just drops those keys. + """ + stats: dict = {} + + # RAM + CPU via psutil (optional dependency). + try: + import psutil + + vm = psutil.virtual_memory() + stats["sys/ram_used_gb"] = vm.used / 1e9 + stats["sys/ram_percent"] = vm.percent + stats["sys/cpu_percent"] = psutil.cpu_percent() + except Exception: # noqa: BLE001 — psutil missing/unavailable -> skip gracefully + pass + + # GPU memory + utilisation via torch.cuda — only when the RUN is on CUDA + # (gated on `device`, not global availability: a CPU run on a GPU box must + # not report GPU stats it isn't using). + try: + import torch + + if getattr(device, "type", str(device)) == "cuda" and torch.cuda.is_available(): + # enumerate ALL devices: DataParallel puts real memory on cuda:1 too, and a + # cuda:0-only reading under-reports peak by ~half (user-flagged blind spot). + for d in range(torch.cuda.device_count()): + sfx = "" if d == 0 else f"_gpu{d}" # gpu0 keeps legacy keys + stats[f"sys/gpu_mem_allocated_gb{sfx}"] = torch.cuda.memory_allocated(d) / 1e9 + stats[f"sys/gpu_mem_reserved_gb{sfx}"] = torch.cuda.memory_reserved(d) / 1e9 + stats[f"sys/gpu_max_mem_allocated_gb{sfx}"] = torch.cuda.max_memory_allocated(d) / 1e9 + util = _gpu_utilization() + if util is not None: + stats["sys/gpu_util_pct"] = util + except Exception: # noqa: BLE001 — never crash on missing GPU-stat APIs + pass + + return stats + + +def _gpu_utilization(): + """Best-effort GPU utilisation %, or None if no API is available.""" + try: + import torch + + return float(torch.cuda.utilization()) + except Exception: # noqa: BLE001 — torch.cuda.utilization not always available + pass + try: + import pynvml + + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + return float(pynvml.nvmlDeviceGetUtilizationRates(handle).gpu) + except Exception: # noqa: BLE001 — pynvml missing/unavailable -> give up + return None diff --git a/src/train/loop.py b/src/train/loop.py new file mode 100644 index 0000000000000000000000000000000000000000..042043671429ca01ebb5614bbb4a8f33aa5b2d65 --- /dev/null +++ b/src/train/loop.py @@ -0,0 +1,198 @@ +# src/train/loop.py +import time + +import torch +from torch.utils.data import DataLoader +from train.loss import proposal_a_loss, distillation_loss +from train.logging_utils import system_stats +from data.preprocess import preprocess_image + +def _lambda_cons(cfg, epoch): + if cfg.epochs <= 1: + return cfg.lambda_cons_max + return cfg.lambda_cons_max * (epoch / (cfg.epochs - 1)) + +def train_one_epoch(model, loader, optimizer, cfg, epoch, logger=None, scheduler=None, + max_batches=None) -> dict: + model.train() + device = next(model.parameters()).device # Detector always has trainable head params + lam_cons = _lambda_cons(cfg, epoch) + # Head params with grad — used for the grad-norm diagnostic (no clipping). + head_params = [p for p in model.parameters() if p.requires_grad] + log_every = max(getattr(cfg, "log_every", 10), 1) + agg = {"loss":0.0, "bce_clean":0.0, "bce_aug":0.0, "cons":0.0, "n":0} + # Data-starvation instrumentation: time spent BLOCKED waiting for the next batch + # (dataloader) vs. computing. High data_wait_frac => GPU starved by the streaming + # data path => raise num_workers/prefetch. Measured as the gap between the end of + # one step and receipt of the next batch. + _t_prev_end = time.perf_counter() + for step, batch in enumerate(loader): + if max_batches is not None and step >= max_batches: + break # fractional epoch (cfg.epoch_frac): stop early + t_step = time.perf_counter() + data_wait = t_step - _t_prev_end # time the loop was blocked on the loader + clean = batch["clean_img"].to(device, non_blocking=True) + aug = batch["aug_img"].to(device, non_blocking=True) + # gpu_preprocess: uint8 batches finish preprocessing (float/normalize/resize) here + if clean.dtype == torch.uint8: + from data.preprocess import gpu_finish + _tr = batch.get("target_res") + _tr = int(_tr[0]) if torch.is_tensor(_tr) else _tr + clean = gpu_finish(clean, _tr); aug = gpu_finish(aug, _tr) + y = batch["label"].float().to(device) + # label smoothing: BCE targets eps/1-eps instead of 0/1 — bounds the optimal logits + # (~±logit(eps)) so the head isn't pushed into ±12 saturation (brittle under the + # corruption families and uncalibrated). eps=0 (default) = exact hard-label path. + eps = getattr(cfg, "label_smooth", 0.0) + if eps: + y = y * (1.0 - 2.0 * eps) + eps + out_c = model(clean); out_a = model(aug) + pw = cfg.pos_weight if getattr(cfg, "pos_weight", 0.0) else None + # KD when the batch carries a cached teacher score (precomputed offline, so the + # teacher never runs in this hot loop — keeps the data path from starving the GPU). + lam_kd = getattr(cfg, "lambda_kd", 0.0) + if lam_kd and "teacher_prob" in batch: + tp = batch["teacher_prob"].float().to(device) + loss, parts = distillation_loss(out_c["image_logit"], out_a["image_logit"], y, tp, + cfg.lambda_aug, lam_cons, lam_kd, cfg.entropy_gate, + pos_weight=pw) + else: + loss, parts = proposal_a_loss(out_c["image_logit"], out_a["image_logit"], y, + cfg.lambda_aug, lam_cons, cfg.entropy_gate, + pos_weight=pw) + # FEATURE-consistency (cfg.lambda_feat>0): cosine drift of the pooled trainable + # representation across the clean/corrupted pair, BOTH classes, label-blind. Unlike the + # logit-level cons term, this constrains the REPRESENTATION: feature directions that + # don't survive corruption are suppressed as unreliable, so the verdict must ride + # corruption-stable (semantic/geometric) content. Cosine (not MSE) is scale-invariant — + # no degenerate shrink-the-norms minimum. + lam_feat = getattr(cfg, "lambda_feat", 0.0) + if lam_feat and "pooled_feat" in out_c: + import torch.nn.functional as _F + cons_feat = (1.0 - _F.cosine_similarity( + out_c["pooled_feat"], out_a["pooled_feat"], dim=-1)).mean() + loss = loss + lam_feat * cons_feat + parts = {**parts, "cons_feat": cons_feat} + optimizer.zero_grad(); loss.backward() + # Grad-norm diagnostic: total L2 norm over head params. max_norm is huge so + # this does NOT clip — it only returns the norm. Computed after backward, + # before optimizer.step(), and only when we're going to log it. + grad_norm = None + will_log = logger is not None and (step % log_every == 0) + if will_log and head_params: + # max_norm huge => measures the grad L2 norm without actually clipping. + # Guarded so a diagnostic can never crash training (airtight "logging never crashes"). + try: + grad_norm = float( + torch.nn.utils.clip_grad_norm_(head_params, max_norm=1e9) + ) + except Exception: + grad_norm = None + optimizer.step() + if scheduler is not None: + scheduler.step() # per-STEP (mid-epoch) LR update + bs = len(y) + agg["loss"] += loss.item()*bs + for k in ("bce_clean","bce_aug","cons"): + agg[k] += parts[k].item()*bs + if "kd" in parts: + agg["kd"] = agg.get("kd", 0.0) + parts["kd"].item()*bs + agg["n"] += bs + if will_log: + compute_dt = time.perf_counter() - t_step + total_dt = data_wait + compute_dt # full wall time attributable to this step + metrics = { + "train/loss": loss.item(), + "train/bce_clean": parts["bce_clean"].item(), + "train/bce_aug": parts["bce_aug"].item(), + "train/cons": parts["cons"].item(), + "train/lambda_cons": lam_cons, + "train/batch_size": bs, + "train/throughput_img_per_s": bs / total_dt if total_dt > 0 else 0.0, + # GPU-saturation signals: data_wait_frac near 0 => GPU-bound (good); + # near 1 => starved waiting on the data path (raise workers/prefetch). + "perf/data_wait_frac": data_wait / total_dt if total_dt > 0 else 0.0, + "perf/data_wait_s": data_wait, + "perf/compute_s": compute_dt, + } + if grad_norm is not None: + metrics["train/grad_norm"] = grad_norm + metrics.update(system_stats(device)) + logger.log(metrics) + _t_prev_end = time.perf_counter() # for next step's data_wait measurement + n = max(agg.pop("n"), 1) + return {k: v/n for k, v in agg.items()} + +def build_predict_fn(model, cfg): + @torch.no_grad() + def predict_fn(imgs): + model.eval() + device = next(model.parameters()).device + # Chunk the forward pass by eval_batch_size. Stacking the WHOLE val set into + # one forward OOMs on real val sets (T4: 235x518px images = 2.77 GiB in + # a single alloc on top of the compiled model's CUDA-graph pools). + bs = max(getattr(cfg, "eval_batch_size", 32), 1) + probs = [] + for i in range(0, len(imgs), bs): + chunk = imgs[i:i + bs] + x = torch.stack([preprocess_image(im, cfg.resolution) for im in chunk]).to(device) + out = model(x) + probs.append(torch.sigmoid(out["image_logit"]).cpu().numpy()) + import numpy as np + return np.concatenate(probs) if probs else np.empty((0,)) + return predict_fn + +def make_loader(dataset, cfg, shuffle=True, drop_last=False) -> DataLoader: + # Per-batch random resolution (opt-in): resize each batch to a random valid (patch-multiple) + # size so the model sees varied token-grid densities and can't anchor on a fixed 518 layout + # or absolute resolution (a general robustness aug + a resolution-shortcut defense). + collate = None + if getattr(cfg, "random_resolution", False) and shuffle: # train-time only + from data.aug.random_resolution import random_resolution_collate + lo = getattr(cfg, "res_lo", 392); hi = getattr(cfg, "res_hi", 644) + collate = random_resolution_collate(patch=14, lo=lo, hi=hi, seed=cfg.seed) + # prefetch_factor=1 (torch default 2): each prefetched batch is workers x batch x + # 2 views of full-res fp32 tensors in HOST RAM — at batch>=64/res 644 the default + # doubles a multi-GB footprint for little pipeline benefit and can trip the host + # OOM-killer (bridge run died 31s in, no traceback = SIGKILL). + _workers = getattr(cfg, "num_workers", 0) + return DataLoader( + dataset, + batch_size=cfg.batch_size, + shuffle=shuffle, + drop_last=drop_last, + prefetch_factor=1 if _workers > 0 else None, + # ANTI-HANG: a worker blocked forever (hung FUSE read on a mounted shard froze + # the 47k run mid-epoch for 3.5h, process alive, zero errors) must RAISE, not + # freeze — a loud crash is retryable, a silent hang eats the session. + timeout=600 if _workers > 0 else 0, + num_workers=_workers, + pin_memory=getattr(cfg, "pin_memory", False), + collate_fn=collate, + ) + +def find_max_batch(model, cfg, start=None) -> int: + """VRAM batch-size autotuner. On CPU (no CUDA) this is a no-op returning + cfg.batch_size. On CUDA it probes descending powers-of-two forward passes, + returning the largest that fits. Never crashes training.""" + if not torch.cuda.is_available(): + return cfg.batch_size + try: + bs = start or cfg.batch_size * 4 + res = cfg.resolution + while bs >= 1: + try: + torch.cuda.empty_cache() + x = torch.randn(bs, 3, res, res, device="cuda") + model(x) + return bs + except (torch.cuda.OutOfMemoryError, RuntimeError) as e: + if "out of memory" not in str(e).lower() and not isinstance( + e, torch.cuda.OutOfMemoryError + ): + raise + torch.cuda.empty_cache() + bs //= 2 + return 1 + except Exception: # noqa: BLE001 — never crash training + return cfg.batch_size diff --git a/src/train/loss.py b/src/train/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..486eda499b4893c8ec5a9d7f10e8548f641e5a0a --- /dev/null +++ b/src/train/loss.py @@ -0,0 +1,54 @@ +# src/train/loss.py +import torch +import torch.nn.functional as F + +def _binary_entropy(p: torch.Tensor) -> torch.Tensor: + p = p.clamp(1e-6, 1 - 1e-6) + return -(p * p.log2() + (1 - p) * (1 - p).log2()) + +def proposal_a_loss(clean_logit, aug_logit, labels, lambda_aug, lambda_cons, entropy_gate, + pos_weight=None): + # pos_weight scales the positive(fake=1) BCE term. A fake-skewed train set collapses + # the operating point to "fake"; pos_weight<1 rebalances it. None = unweighted. + pw = None if pos_weight is None else torch.as_tensor( + pos_weight, dtype=clean_logit.dtype, device=clean_logit.device) + bce_clean = F.binary_cross_entropy_with_logits(clean_logit, labels, pos_weight=pw) + bce_aug = F.binary_cross_entropy_with_logits(aug_logit, labels, pos_weight=pw) + + pc = torch.sigmoid(clean_logit) + pa = torch.sigmoid(aug_logit) + # Consistency pulls the AUG prediction toward the CLEAN one, treating clean as the + # ANCHOR: pc is detached so the corrupted view is regularized toward the clean view, + # NOT the reverse (a symmetric term lets the corrupted branch degrade + # the clean prediction). Directed KL( Bern(pc_detached) || Bern(pa) ). + pc_anchor = pc.detach() + def _kl(p, q): + p = p.clamp(1e-6, 1-1e-6); q = q.clamp(1e-6, 1-1e-6) + return p*(p/q).log() + (1-p)*((1-p)/(1-q)).log() + cons_per = _kl(pc_anchor, pa) + keep = (_binary_entropy(pa) <= entropy_gate).float() # gate ambiguous aug views + denom = keep.sum().clamp(min=1.0) + cons = (cons_per * keep).sum() / denom + + loss = bce_clean + lambda_aug * bce_aug + lambda_cons * cons + return loss, {"bce_clean": bce_clean, "bce_aug": bce_aug, "cons": cons} + + +def distillation_loss(clean_logit, aug_logit, labels, teacher_prob, + lambda_aug, lambda_cons, lambda_kd, entropy_gate, pos_weight=None): + """proposal_a_loss + a KNOWLEDGE-DISTILLATION term: the student's clean prediction is + pulled toward the teacher's SOFT score (CommForensics-384, a 4803-generator prior). We + distill the scalar teacher probability (binary cross-entropy against the soft target, + NOT feature-MSE) — robust, ROC-AUC-aligned, and the mechanism behind the amortization + claim (student inherits the teacher's cross-generator generalization). lambda_kd=0 + recovers proposal_a_loss exactly.""" + base, parts = proposal_a_loss(clean_logit, aug_logit, labels, lambda_aug, lambda_cons, + entropy_gate, pos_weight=pos_weight) + # soft-target BCE: -[t·log s + (1-t)·log(1-s)], t=teacher prob, s=student clean prob + t = torch.as_tensor(teacher_prob, dtype=clean_logit.dtype, device=clean_logit.device).clamp(1e-6, 1-1e-6) + kd = F.binary_cross_entropy_with_logits( + clean_logit, t, reduction="mean") - ( + -(t * t.log() + (1 - t) * (1 - t).log())).mean() # subtract teacher entropy -> KD≈0 at match + kd = kd.clamp(min=0.0) + parts = {**parts, "kd": kd} + return base + lambda_kd * kd, parts diff --git a/src/train/preflight.py b/src/train/preflight.py new file mode 100644 index 0000000000000000000000000000000000000000..f630f1e480e3962515b6bc340a49ac7c7f0fc27c --- /dev/null +++ b/src/train/preflight.py @@ -0,0 +1,214 @@ +# src/train/preflight.py +"""Frontloaded runtime assertions — fail fast BEFORE the expensive training loop. + +Every deep failure in this project surfaced late: the P100/sm_60 CUDA error died 455s into the +run at the DINOv2 conv2d; a swallowed-config bug silently disabled augmentations for a whole run. +The cost of a late failure on a remote GPU host is a 15-30 min round trip. This module runs all the cheap, +high-signal checks — data integrity + a REAL forward-pass smoke through the actual model+cfg — so +a bad config, empty mount, class imbalance, leak, un-decodable image, or GPU/kernel incompatibility +raises in seconds, not deep inside the training loop. + +Call `preflight_check(cfg, train_df, val_df, holdout_df, device="cuda")` right after building cfg +and before `train_model`. It raises AssertionError with a specific message on any failure, and +returns a dict of what it verified (logged for the run record). +""" +import torch +from PIL import Image + +from data.preprocess import preprocess_image + + +def tune_batch_size(model, cfg, device="cuda", candidates=(128, 96, 64, 48, 32, 16, 8), + mem_frac=0.80, _probe=None) -> dict: + """Pick the largest batch size that fits at the WORST-CASE resolution. + + random_resolution varies the train res up to cfg.res_hi (644) — memory scales + ~ (res/518)^2, so a batch tuned at 518 can OOM on a 644 batch mid-epoch. This + probes descending candidates at res_hi through the SAME DataParallel wrapper and + dual forward+backward shape the train step uses, and returns the largest batch + whose peak stays under mem_frac of the smallest GPU's capacity. Deterministic + fail-fast: OOM is decided in preflight seconds, not at epoch 1's unlucky batch. + + Returns {batch_size, peak_gb, capacity_gb, resolution, n_gpus}. On CPU (no CUDA) + returns cfg.batch_size unchanged — the probe is a GPU-memory question. + `_probe(bs)->peak_bytes|raise OOM` is injectable for tests.""" + res = int(getattr(cfg, "res_hi", cfg.resolution)) if getattr(cfg, "random_resolution", False) \ + else int(cfg.resolution) + if _probe is None and not torch.cuda.is_available(): + return {"batch_size": int(cfg.batch_size), "peak_gb": None, "capacity_gb": None, + "resolution": res, "n_gpus": 0} + n_gpu = 1 + dp = model + if _probe is None: + if (getattr(cfg, "data_parallel", False) and torch.cuda.device_count() > 1): + dp = torch.nn.DataParallel(model) + n_gpu = torch.cuda.device_count() + capacity = min(torch.cuda.get_device_properties(d).total_memory + for d in range(n_gpu)) + else: + capacity = 16 * 2**30 # test injection: nominal T4 + + def _real_probe(bs): + torch.cuda.empty_cache() + for d in range(n_gpu): + torch.cuda.reset_peak_memory_stats(d) + x = torch.randn(bs, 3, res, res, device=device) + dp.train() + out_c = dp(x) # the train step runs TWO forwards + out_a = dp(x) # (clean + aug) before one backward + loss = out_c["image_logit"].sum() + out_a["image_logit"].sum() + loss.backward() + peak = max(torch.cuda.max_memory_allocated(d) for d in range(n_gpu)) + model.zero_grad(set_to_none=True) + del x, out_c, out_a, loss + torch.cuda.empty_cache() + return peak + + # HOST-RAM cap: the GPU probe says nothing about the DataLoader side, where + # workers x prefetch batches of 2 full-res fp32 views live in host memory. The + # bridge run's batch-128 pick passed the GPU probe then host-OOM-killed 31s into + # training (no traceback = SIGKILL). Cap candidates so in-flight loader bytes + # (workers x prefetch(1) x batch x 2 views x res^2 x 3 x 4B, +1 collating) stay + # under half the AVAILABLE host RAM. + try: + import psutil + avail = psutil.virtual_memory().available + workers = max(1, int(getattr(cfg, "num_workers", 0))) + per_img = 2 * res * res * 3 * 4 # clean+aug fp32 views + max_bs_host = int((0.5 * avail) / ((workers + 1) * per_img)) + candidates = [bs for bs in candidates if bs <= max_bs_host] or [min(candidates)] + except ImportError: + pass + probe = _probe or _real_probe + last_err = None + for bs in candidates: + try: + peak = probe(bs) + except torch.cuda.OutOfMemoryError as e: # noqa: PERF203 + last_err = e + model.zero_grad(set_to_none=True) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + continue + if peak <= mem_frac * capacity: + return {"batch_size": int(bs), "peak_gb": peak / 1e9, + "capacity_gb": capacity / 1e9, "resolution": res, "n_gpus": n_gpu} + raise AssertionError( + f"tune_batch_size: no candidate in {candidates} fits at res {res} " + f"(capacity {capacity/1e9:.1f}GB, mem_frac {mem_frac}); last OOM: {last_err}") + + +def _assert_both_classes(df, name): + assert df is not None and len(df) > 0, f"preflight: {name}_df is empty" + labs = set(int(x) for x in df["label"].unique()) + assert labs == {0, 1}, f"preflight: {name}_df must have both classes, got {labs}" + + +def preflight_check(cfg, train_df, val_df, holdout_df=None, model=None, device="cpu", + n_smoke=4, resolver=None) -> dict: + """Validate data + run a forward-pass smoke before training. Raises on any failure. + + Checks, cheapest first: + 1. train/val have BOTH classes and are non-empty. + 2. holdout (if given) has both classes AND is img_id-disjoint from train (leak guard). + 3. a few train image paths actually decode to RGB. + 4. the model builds with THIS cfg; fusion params are trainable when n_intermediate_layers>1 + (catches the swallowed-config class of bug); trainable-param count is > 0. + 5. one real forward pass on `device` yields finite logits of the right shape, sigmoid in + [0,1] (catches GPU/kernel incompatibility — the P100 conv2d death — and NaN configs). + Returns a report dict of the verified facts. + """ + report = {} + # 1. class balance / non-empty + _assert_both_classes(train_df, "train") + _assert_both_classes(val_df, "val") + report["train_n"] = int(len(train_df)) + report["val_n"] = int(len(val_df)) + report["train_labels"] = {int(k): int(v) for k, v in train_df["label"].value_counts().items()} + + # 2. holdout sanity + leakage guard + if holdout_df is not None and len(holdout_df): + _assert_both_classes(holdout_df, "holdout") + if "img_id" in train_df.columns and "img_id" in holdout_df.columns: + overlap = set(train_df["img_id"]) & set(holdout_df["img_id"]) + assert not overlap, f"preflight: {len(overlap)} img_id(s) leak between train and holdout" + report["holdout_n"] = int(len(holdout_df)) + + # 3. image decode smoke — when the corpus is sharded, sample from >=3 distinct, + # NON-ADJACENT shards (spec: sharded-corpus-pipeline) so a bad/truncated shard + # fails here in seconds, not mid-epoch. + if "shard_name" in train_df.columns and resolver is not None: + shard_names = sorted(train_df["shard_name"].unique()) + step = max(1, len(shard_names) // 3) # non-adjacent spread + picks = shard_names[::step][:max(3, n_smoke)] + rows = [train_df[train_df["shard_name"] == s].iloc[0] for s in picks] + imgs = [] + for r in rows: + try: + imgs.append(resolver.resolve(r)) + except Exception as e: # noqa: BLE001 + raise AssertionError( + f"preflight: shard member failed to decode: " + f"{r['shard_name']}/{r['member_name']} ({e})") + report["shards_smoked"] = len(picks) + else: + paths = train_df["path"].head(n_smoke).tolist() + imgs = [] + for p in paths: + try: + imgs.append(Image.open(p).convert("RGB")) + except Exception as e: # noqa: BLE001 + raise AssertionError(f"preflight: train image failed to decode: {p} ({e})") + assert imgs, "preflight: no train images to smoke-test" + + # 4. model build + config sanity + if model is None: + from model.detector import Detector + model = Detector(cfg) + model.to(device).eval() + n_inter = int(getattr(cfg, "n_intermediate_layers", 1)) + if n_inter > 1: + fusion = getattr(model.backbone, "_fusion", None) + assert fusion is not None, ( + f"preflight: n_intermediate_layers={n_inter} but backbone built no fusion module") + n_fuse = sum(p.numel() for p in fusion.parameters() if p.requires_grad) + assert n_fuse > 0, "preflight: fusion params are not trainable (swallowed config?)" + report["fusion_trainable_params"] = int(n_fuse) + if getattr(cfg, "use_freq_branch", False): + assert getattr(model, "_freq", None) is not None, ( + "preflight: use_freq_branch=True but freq branch not built") + n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) + assert n_train > 0, "preflight: model has no trainable parameters" + report["trainable_params"] = int(n_train) + + # 5. real forward-pass smoke on the target device + x = torch.stack([preprocess_image(im, cfg.resolution) for im in imgs]).to(device) + with torch.no_grad(): + out = model(x) + logit = out["image_logit"] + assert tuple(logit.shape) == (len(imgs),), ( + f"preflight: image_logit shape {tuple(logit.shape)} != ({len(imgs)},)") + assert torch.isfinite(logit).all(), "preflight: non-finite logits from forward smoke" + probs = torch.sigmoid(logit) + assert bool(((probs >= 0) & (probs <= 1)).all()), "preflight: sigmoid probs out of [0,1]" + report["smoke_forward_ok"] = True + report["device"] = str(device) + + # 6. DataParallel smoke — the DP path only activates with >1 CUDA device, so it is NEVER + # exercised by CPU CI; its first real run is on the GPU host. Do a real DP forward+backward here so a + # DP incompatibility (dict-output gather, frozen-param grads, batch too small to split) fails in + # SECONDS instead of deep in training. Only runs when DP would actually engage. + if (getattr(cfg, "data_parallel", False) and torch.cuda.is_available() + and torch.cuda.device_count() > 1): + dp = torch.nn.DataParallel(model) + dp.train() + xb = x.repeat((max(2, torch.cuda.device_count()), 1, 1, 1)) # ensure batch splits across GPUs + out_dp = dp(xb) + loss = out_dp["image_logit"].sum() + loss.backward() # grads must flow to the head + grad_ok = any(p.grad is not None for p in model.parameters() if p.requires_grad) + assert grad_ok, "preflight: DataParallel produced no gradients on trainable params" + model.zero_grad(set_to_none=True) + report["data_parallel_smoke_ok"] = True + report["n_gpus"] = int(torch.cuda.device_count()) + return report diff --git a/src/train/watchdog.py b/src/train/watchdog.py new file mode 100644 index 0000000000000000000000000000000000000000..b7cd50534b9dfd0a2148f8e2b7da3fb8eb492662 --- /dev/null +++ b/src/train/watchdog.py @@ -0,0 +1,120 @@ +# src/train/watchdog.py +"""Separate-thread telemetry watchdog (user-directed). + +Inline logging has two structural flaws we kept paying for: (a) it only fires when +the main loop reaches a log point — a silent phase (serial eval, a hang) logs +NOTHING, exactly when visibility matters most; (b) sampling GPU/CPU stats inside +the hot loop measures the loop's own log-step, biasing perf characteristics. + +The watchdog is a daemon thread that samples on ITS OWN clock every `interval` +seconds regardless of what the main thread is doing: GPU util/mem (via NVML when +available — measured utilization, not torch allocator state), CPU%, RAM, plus a +caller-updatable `phase` string and arbitrary progress counters. Emits to stdout +and (optionally) a wandb-style logger under `sys/*` and `watchdog/*` keys. + +Usage: + wd = Watchdog(logger=logger, interval=15).start() + wd.set_phase("train epoch 0") + wd.update(step=120, cells_done=4000, cells_total=132962) + ... + wd.stop() +Threads never crash the run: every sample is exception-guarded. +""" +import threading +import time + + +class Watchdog: + def __init__(self, logger=None, interval=15.0, print_every=2): + self.logger = logger + self.interval = float(interval) + self.print_every = int(print_every) # print 1 of every N samples (wandb gets all) + self._state = {"phase": "init"} + self._lock = threading.Lock() + self._stop = threading.Event() + self._thread = None + self._n = 0 + + # ---- caller API ---- + def set_phase(self, phase): + with self._lock: + self._state["phase"] = str(phase) + self._state.pop("progress_frac", None) # phase change resets progress + + def update(self, **counters): + with self._lock: + self._state.update(counters) + ct, cd = self._state.get("cells_total"), self._state.get("cells_done") + if ct: + self._state["progress_frac"] = float(cd or 0) / float(ct) + + def start(self): + self._thread = threading.Thread(target=self._run, daemon=True, + name="track5-watchdog") + self._thread.start() + return self + + def stop(self): + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=self.interval + 1) + + # ---- sampling ---- + def _sample(self): + s = {} + try: + import psutil + vm = psutil.virtual_memory() + s["sys/cpu_percent"] = psutil.cpu_percent() + s["sys/ram_percent"] = vm.percent + except Exception: # noqa: BLE001 + pass + try: + import torch + if torch.cuda.is_available(): + for d in range(torch.cuda.device_count()): + sfx = "" if d == 0 else f"_gpu{d}" + s[f"sys/gpu_mem_allocated_gb{sfx}"] = torch.cuda.memory_allocated(d) / 1e9 + try: # NVML: MEASURED util, not allocator state + import pynvml + pynvml.nvmlInit() + for d in range(torch.cuda.device_count()): + h = pynvml.nvmlDeviceGetHandleByIndex(d) + u = pynvml.nvmlDeviceGetUtilizationRates(h) + sfx = "" if d == 0 else f"_gpu{d}" + s[f"sys/gpu_util_pct{sfx}"] = u.gpu + s[f"sys/gpu_mem_used_gb{sfx}"] = \ + pynvml.nvmlDeviceGetMemoryInfo(h).used / 1e9 + except Exception: # noqa: BLE001 - NVML optional + pass + except Exception: # noqa: BLE001 + pass + return s + + def _run(self): + while not self._stop.wait(self.interval): + self._n += 1 + try: + s = self._sample() + with self._lock: + st = dict(self._state) + s["watchdog/phase"] = st.get("phase", "?") + for k, v in st.items(): + if isinstance(v, (int, float)): + s[f"watchdog/{k}"] = v + if self.logger is not None: + try: + self.logger.log(s) + except Exception: # noqa: BLE001 + pass + if self._n % self.print_every == 0: + gpu = s.get("sys/gpu_util_pct") + frac = st.get("progress_frac") + print(f"[watchdog] {st.get('phase','?')}" + + (f" {100*frac:.0f}%" if frac is not None else "") + + (f" | gpu {gpu}%" if gpu is not None else "") + + f" | cpu {s.get('sys/cpu_percent','?')}%" + + f" | ram {s.get('sys/ram_percent','?')}%", + flush=True) + except Exception: # noqa: BLE001 - the watchdog NEVER kills the run + pass