ensemble/run_attack.py

"""Spark-side runner: attack each image against the TRAIN encoder ensemble, then
score decoy-truth margin (clean vs adv, adv-after-JPEG) on BOTH train and HELD-OUT
encoders to measure in-ensemble effect and cross-architecture transfer.

Saves adversarial PNGs (for later frontier eval on the Mac) + results.json.

  python -m ensemble.run_attack --manifest testset.csv --images testset \
      --out runs/demo --steps 200 --eps 12 --subset 4 [--limit 8] [--smoke]
"""

from __future__ import annotations

import argparse
import io
import json
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from ensemble.attack import AttackCfg, attack_image
from ensemble.encoders import (ABLATION_TRAIN, DEFAULT_HELDOUT, DEFAULT_TRAIN,
                               FULL_TRAIN, V02_HELDOUT, V02_TRAIN,
                               V021_ATTACK, V021_ATTACK_BASE, V021_ATTACK_GIANTS,
                               V021_HELDOUT, load_encoder)

# v0.1 (non-Imagenette) decoys — kept only for reference / --train full reproduction.
DECOYS_V01 = {
    "cassette player": "jellyfish", "tench": "steam locomotive", "church": "jellyfish",
    "chainsaw": "peacock", "english springer": "cassette player",
    "french horn": "jellyfish", "garbage truck": "sunflower", "gas pump": "banana",
    "golf ball": "volcano", "parachute": "octopus",
}

# v0.2 decoys are WITHIN Imagenette (reciprocal far pairs) so feature towers get an
# image-feature centroid target (M4) from the exemplar pool, and contrastive text +
# feature steering pull the SAME decoy direction. Semantically far within the 10 classes.
DECOYS = {
    "tench": "church", "church": "tench",
    "english springer": "gas pump", "gas pump": "english springer",
    "cassette player": "parachute", "parachute": "cassette player",
    "chain saw": "golf ball", "golf ball": "chain saw",
    "french horn": "garbage truck", "garbage truck": "french horn",
}


def log(m):
    print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)


def to_tensor(img: Image.Image) -> torch.Tensor:
    a = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(a).permute(2, 0, 1).unsqueeze(0).cuda()


def jpeg(img: Image.Image, q=85) -> Image.Image:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="JPEG", quality=q)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


def tensor_to_pil(x: torch.Tensor) -> Image.Image:
    a = (x.squeeze(0).clamp(0, 1).permute(1, 2, 0).cpu().float().numpy() * 255).round().astype(np.uint8)
    return Image.fromarray(a)


@torch.no_grad()
def margin(enc, img: Image.Image, truth: str, decoy: str, centroids=None) -> float:
    """Decoy-minus-truth similarity margin (higher = more fooled toward decoy).

    Contrastive encoders use text-tower label embeddings; feature towers use the
    decoy/truth image-feature centroids (M4) when available.
    """
    x = to_tensor(img)
    f = enc.image_feat(x).float()
    if enc.kind == "contrastive":
        t = enc.text_feat([truth, decoy]).float()
        return float(((f @ t[1:2].T) - (f @ t[0:1].T)).item())
    cts = centroids.get(enc.name) if centroids else None
    if not cts or decoy not in cts or truth not in cts:
        return float("nan")
    c_d, c_t = cts[decoy].float(), cts[truth].float()
    return float(((f @ c_d.T) - (f @ c_t.T)).item())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--images", required=True, help="dir holding the image files")
    ap.add_argument("--out", default="runs/demo")
    ap.add_argument("--steps", type=int, default=200)
    ap.add_argument("--eps", type=float, default=12.0, help="L-inf in /255")
    ap.add_argument("--subset", type=int, default=4)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--train", default="default",
                    choices=["default", "ablation", "full", "v0.2", "v0.2.1",
                             "v0.2.1-base", "v0.2.1-giants"],
                    help="'ablation' adds 2 timm VLM towers; 'full' also adds the 3 "
                         "modern-LLM towers; 'v0.2' cross-arch split w/ M4 centroids; "
                         "'v0.2.1' ADOPTED = full 13-tower v0.1 ensemble + M4 + never-trained "
                         "cross-arch judges (C-RADIO, Qwen3-VL); 'v0.2.1-giants' adds the "
                         "ablation-rejected SigLIP2-giant + MetaCLIP2-H")
    ap.add_argument("--exemplars", default="examples/exemplars",
                    help="dir of per-class exemplar images for M4 feature-tower targets")
    ap.add_argument("--no-feature-target", action="store_true",
                    help="ablate M4: feature towers use repel-only (still SCORED via centroids)")
    ap.add_argument("--decoys", default="v0.2", choices=["v0.1", "v0.2"],
                    help="v0.2 = reciprocal Imagenette; v0.1 = original jellyfish-style decoys")
    # transfer/stealth levers (defaults off = plain momentum-PGD)
    ap.add_argument("--grad-norm", action="store_true", help="per-encoder unit-L2 grad agg")
    ap.add_argument("--vmi-n", type=int, default=0, help="VMI neighbor samples (0=off)")
    ap.add_argument("--vmi-beta", type=float, default=1.5)
    ap.add_argument("--max-per-family", type=int, default=99)
    ap.add_argument("--min-feature", type=int, default=0)
    ap.add_argument("--lpips-weight", type=float, default=0.0)
    ap.add_argument("--lpips-tau", type=float, default=0.0)
    ap.add_argument("--dct-keep", type=float, default=1.0)
    ap.add_argument("--metrics", action="store_true", help="record PSNR/SSIM/ΔE/LPIPS")
    ap.add_argument("--smoke", action="store_true")
    args = ap.parse_args()

    if args.smoke:
        args.steps, args.limit = 30, 2

    out = Path(args.out)
    (out / "adv").mkdir(parents=True, exist_ok=True)
    imgdir = Path(args.images)

    rows = []
    for line in Path(args.manifest).read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#"):
            p, t = line.split(",", 1)
            rows.append((Path(p).name, t.strip().lower()))
    if args.limit:
        step = max(1, len(rows) // args.limit)
        rows = rows[::step][:args.limit]

    train_names = {"full": FULL_TRAIN, "ablation": ABLATION_TRAIN,
                   "v0.2": V02_TRAIN, "v0.2.1": V021_ATTACK,
                   "v0.2.1-base": V021_ATTACK_BASE,
                   "v0.2.1-giants": V021_ATTACK_GIANTS}.get(args.train, DEFAULT_TRAIN)
    m4_trains = ("v0.2", "v0.2.1", "v0.2.1-base", "v0.2.1-giants")
    held_names = (V021_HELDOUT if args.train in ("v0.2.1", "v0.2.1-base", "v0.2.1-giants")
                  else V02_HELDOUT if args.train == "v0.2" else DEFAULT_HELDOUT)
    log(f"loading TRAIN encoders ({args.train}): {train_names}")
    train = [load_encoder(n) for n in train_names]
    log(f"loading HELD-OUT encoders: {held_names}")
    held = [load_encoder(n) for n in held_names]
    torch.cuda.synchronize()
    log(f"VRAM after load: {torch.cuda.memory_allocated()/1e9:.1f} GB")

    # M4: per-encoder class feature centroids for feature-tower decoy steering + scoring.
    centroids = None
    has_feature = any(e.kind != "contrastive" for e in train + held)
    if args.train in m4_trains and has_feature:
        from ensemble.targets import compute_centroids
        log(f"computing feature-tower centroids from exemplars: {args.exemplars}")
        centroids = compute_centroids(train + held, args.exemplars, log=log)
    # centroids drive SCORING always; M4 ablation only removes them from the attack.
    attack_centroids = None if args.no_feature_target else centroids

    cfg = AttackCfg(
        eps=args.eps / 255, steps=args.steps, subset=args.subset,
        grad_norm=args.grad_norm, vmi_n=args.vmi_n, vmi_beta=args.vmi_beta,
        max_per_family=args.max_per_family, min_feature=args.min_feature,
        lpips_weight=args.lpips_weight, lpips_tau=args.lpips_tau, dct_keep=args.dct_keep,
    )
    lpips_fn = None
    if args.lpips_weight > 0 or args.lpips_tau > 0 or args.metrics:
        from ensemble.perceptual import LPIPSPenalty
        lpips_fn = LPIPSPenalty(net="alex", device="cuda")
        log("LPIPS (alex) loaded")
    results = []
    for i, (fname, truth) in enumerate(rows):
        fp = imgdir / fname
        if not fp.exists():
            log(f"skip missing {fp}")
            continue
        img = Image.open(fp).convert("RGB")
        decoy_map = DECOYS_V01 if args.decoys == "v0.1" else DECOYS
        decoy = decoy_map.get(truth, "jellyfish")
        t0 = time.time()
        x_clean = to_tensor(img)
        adv = attack_image(x_clean, truth, decoy, train, cfg,
                           lpips_fn=lpips_fn, centroids=attack_centroids)
        adv_img = tensor_to_pil(adv)
        adv_img.save(out / "adv" / f"{fp.stem}.png")
        adv_jpeg = jpeg(adv_img, 85)

        rec = {"image": fname, "truth": truth, "decoy": decoy, "train": {}, "heldout": {}}
        if args.metrics:
            from ensemble.perceptual import stealth_metrics
            rec["stealth"] = stealth_metrics(adv, x_clean, lpips_fn)
        for grp, encs in (("train", train), ("heldout", held)):
            for enc in encs:
                # feature towers score via M4 centroids; skip only if no centroid target
                if enc.kind != "contrastive" and centroids is None:
                    continue
                mc = round(margin(enc, img, truth, decoy, centroids), 3)
                ma = round(margin(enc, adv_img, truth, decoy, centroids), 3)
                mj = round(margin(enc, adv_jpeg, truth, decoy, centroids), 3)
                if mc != mc or ma != ma or mj != mj:   # NaN -> no target, skip
                    continue
                rec[grp][enc.name] = {"clean": mc, "adv": ma, "adv_jpeg": mj}

        def flips(grp):  # margin crosses 0 (decoy now beats truth) after JPEG
            v = rec[grp]
            return sum(d["adv_jpeg"] > 0 and d["clean"] < 0 for d in v.values()), len(v)

        tf, tn = flips("train"); hf, hn = flips("heldout")
        rec["train_flip_jpeg"] = f"{tf}/{tn}"
        rec["heldout_flip_jpeg"] = f"{hf}/{hn}"
        results.append(rec)
        log(f"[{i+1}/{len(rows)}] {fname} {truth!r}->{decoy!r} "
            f"train_flip(jpeg)={tf}/{tn} heldout_flip(jpeg)={hf}/{hn} ({time.time()-t0:.0f}s)")

    (out / "results.json").write_text(json.dumps(results, indent=2))

    def agg(grp, key):
        vals = [d[key] for r in results for d in r[grp].values()]
        return sum(vals) / len(vals) if vals else 0.0

    def flip_total(grp):
        f = n = 0
        for r in results:
            for d in r[grp].values():
                n += 1; f += (d["adv_jpeg"] > 0 and d["clean"] < 0)
        return f, n

    print("\n============ ENSEMBLE ATTACK SUMMARY ============")
    print(f"images: {len(results)}  eps={args.eps}/255  steps={args.steps}  subset={args.subset}")
    print(f"levers: grad_norm={args.grad_norm} vmi_n={args.vmi_n} "
          f"max_per_family={args.max_per_family} min_feature={args.min_feature} "
          f"lpips_w={args.lpips_weight} lpips_tau={args.lpips_tau} dct_keep={args.dct_keep}")
    for grp in ("train", "heldout"):
        f, n = flip_total(grp)
        print(f"{grp:>7}: mean clean margin {agg(grp,'clean'):+.3f} -> adv {agg(grp,'adv'):+.3f} "
              f"-> adv+JPEG {agg(grp,'adv_jpeg'):+.3f} | decoy-beats-truth after JPEG {f}/{n}")
    if args.metrics and results:
        keys = [k for k in ("psnr", "ssim", "deltaE_p95", "lpips") if k in results[0].get("stealth", {})]
        ms = {k: sum(r["stealth"][k] for r in results) / len(results) for k in keys}
        print("stealth (mean): " + "  ".join(f"{k}={v:.3f}" for k, v in ms.items()))
    print("====================================================")


if __name__ == "__main__":
    main()