Blame for run_experiments.py

#!/usr/bin/env python3
"""AIP-AD experiments on Real-IAD (multi-view, 5 cameras).

Stage 'score'  : PatchCore-style per-view scoring -> results/scores/<cat>.npz
Stage 'analyze': fusion tables + GRPO arbiter + continual-learning curves
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from src.data import load_samples, subsample_samples, split_calibration  # noqa: E402
from src.scorer import PatchScorer  # noqa: E402

ROOT = os.path.dirname(os.path.abspath(__file__))
DATA = "/home/z/my-project/data"
IMG_ROOT = os.path.join(DATA, "realiad_256")
JSON_ROOT = os.path.join(DATA, "mirror/realiad_jsons")
RESULTS = os.path.join(ROOT, "results")


def stage_score(categories, n_ok, n_ng, n_train, seed, force=False):
    import torch
    torch.set_num_threads(2)
    os.makedirs(os.path.join(RESULTS, "scores"), exist_ok=True)
    for cat in categories:
        out = os.path.join(RESULTS, "scores", f"{cat}.npz")
        if os.path.exists(out) and not force:
            print(f"[skip] {cat} already scored")
            continue
        t0 = time.time()
        train = load_samples(cat, IMG_ROOT, JSON_ROOT, "train")
        train_cap = subsample_samples(train, n_ok=10**6, n_ng=0, seed=seed)
        paths = [p for s in train_cap[:n_train] for p in s.image_paths]
        scorer = PatchScorer(seed=seed)
        scorer.build_bank(paths, patches_per_image=120, bank_size=2048)
        scorer.save(os.path.join(RESULTS, "scores", f"{cat}_bank.npy"))

        test = load_samples(cat, IMG_ROOT, JSON_ROOT, "test")
        test_sub = subsample_samples(test, n_ok=n_ok, n_ng=n_ng, seed=seed)
        calib, ev = split_calibration(test_sub, calib_frac=0.25, seed=seed)
        rows = []
        for split_name, samples in (("calib", calib), ("eval", ev)):
            flat_paths, sample_idx, view_idx = [], [], []
            for si, s in enumerate(samples):
                for vi, p in enumerate(s.image_paths):
                    flat_paths.append(p)
                    sample_idx.append(si)
                    view_idx.append(vi)
            sc, _ = scorer.score_images(flat_paths)
            n = len(samples)
            vs = np.full((n, 5), np.nan, dtype=np.float32)
            vl = np.zeros((n, 5), dtype=np.int8)
            for k, (si, vi) in enumerate(zip(sample_idx, view_idx)):
                vs[si, vi] = sc[k]
                vl[si, vi] = samples[si].label
            for si, s in enumerate(samples):
                rows.append(dict(split=split_name, category=cat, sample_id=s.sample_id,
                                 label=int(s.label), anomaly_class=s.anomaly_class))
            np.savez_compressed(out + f".{split_name}.tmp.npz", view_scores=vs, view_labels=vl)
            os.replace(out + f".{split_name}.tmp.npz", out + f".{split_name}.npz")
        with open(out + ".meta.json", "w") as f:
            json.dump({"rows": rows, "n_train_imgs": len(paths), "seed": seed,
                       "n_test_ok": n_ok, "n_test_ng": n_ng}, f)
        print(f"[done] {cat}: {len(rows)} samples in {time.time()-t0:.0f}s")


def stage_analyze(categories, out_json):
    from src.arbiter import CalibNormalizer, GRPOArbiter, OnlineHedge
    from src.evalmetrics import auroc
    from src.fusion import STRATEGY_NAMES, fuse_all

    data = {}
    acl_path = os.path.join(RESULTS, "agent_calibration.json")
    acl_stats = json.load(open(acl_path)) if os.path.exists(acl_path) else {}

    def acl_apply(vs, cat):
        """Agent Calibration Layer: per-camera quantile standardization."""
        if not acl_stats or cat not in acl_stats:
            return vs
        out = vs.copy()
        for v in range(vs.shape[1]):
            st = acl_stats[cat][f"C{v+1}"]
            out[:, v] = (vs[:, v] - st["median"]) / st["iqr"]
        return out

    for cat in categories:
        cal = np.load(os.path.join(RESULTS, "scores", f"{cat}.npz.calib.npz"))
        ev = np.load(os.path.join(RESULTS, "scores", f"{cat}.npz.eval.npz"))
        data[cat] = dict(
            calib=cal, eval=ev,
            calib_acl=acl_apply(cal["view_scores"].astype(np.float64), cat),
            eval_acl=acl_apply(ev["view_scores"].astype(np.float64), cat),
        )
        assert np.isfinite(cal["view_scores"]).all() and np.isfinite(ev["view_scores"]).all()

    # ---- fixed strategy table (eval split) ----
    table = {}
    img_auc = {}
    for cat in categories:
        ev = data[cat]["eval"]
        vs, vl, y = ev["view_scores"].astype(np.float64), ev["view_labels"], ev["view_labels"].max(1)
        fused = fuse_all(vs)
        for name, s in fused.items():
            table.setdefault(name, {})[cat] = auroc(s, y)
        img_auc[cat] = auroc(vs.reshape(-1), np.repeat(y, vs.shape[1]))
    for name in table:
        table[name]["macro"] = float(np.nanmean([table[name][c] for c in categories]))

    # ---- calibration-split strategy table (used for hyperparameter selection) ----
    calib_table = {}
    for cat in categories:
        cal = data[cat]["calib"]
        vs = cal["view_scores"].astype(np.float64)
        y = cal["view_labels"].max(1)
        fused = fuse_all(vs)
        for name, s in fused.items():
            calib_table.setdefault(name, {})[cat] = auroc(s, y)
    for name in calib_table:
        calib_table[name]["macro"] = float(np.nanmean([calib_table[name][c] for c in categories]))

    # ---- ACL (calibrated agents) fixed-strategy table (eval split) ----
    acl_table = {}
    for cat in categories:
        vs = data[cat]["eval_acl"]
        y = data[cat]["eval"]["view_labels"].max(1)
        fused = fuse_all(vs)
        for name, s in fused.items():
            acl_table.setdefault(name, {})[cat] = auroc(s, y)
    for name in acl_table:
        acl_table[name]["macro"] = float(np.nanmean([acl_table[name][c] for c in categories]))

    # ---- GRPO arbiter (train on pooled calibration, eval) ----
    norm = CalibNormalizer()
    cal_scores_by_cat, cal_labels_by_cat = {}, {}
    for cat in categories:
        vs = data[cat]["calib_acl"]
        y = data[cat]["calib"]["view_labels"].max(1)
        fused = fuse_all(vs)
        cal_scores_by_cat[cat] = {n: fused[n] for n in STRATEGY_NAMES}
        cal_labels_by_cat[cat] = y
    norm.fit(cal_scores_by_cat, cal_labels_by_cat)

    vs_all, y_all, c_all = [], [], []
    for cat in categories:
        cal = data[cat]["calib"]
        vs_all.append(data[cat]["calib_acl"])
        y_all.append(cal["view_labels"].max(1))
        c_all.append(np.full(len(cal["view_scores"]), cat))
    vs_all, y_all, c_all = np.concatenate(vs_all), np.concatenate(y_all), np.concatenate(c_all)

    arb = GRPOArbiter(seed=0)
    arb.fit(vs_all, y_all, c_all, norm, epochs=50, log=print)

    eval_scores, eval_y, eval_c, eval_policy = [], [], [], []
    for cat in categories:
        ev = data[cat]["eval"]
        vs = data[cat]["eval_acl"]
        y = ev["view_labels"].max(1)
        c = np.full(len(vs), cat)
        s, P, _ = arb.fused_scores(vs, c, norm)
        eval_scores.append(s)
        eval_y.append(y)
        eval_c.append(c)
        eval_policy.append(P.argmax(1))
    eval_scores, eval_y, eval_c = np.concatenate(eval_scores), np.concatenate(eval_y), np.concatenate(eval_c)
    eval_policy = np.concatenate(eval_policy)
    arb_per = {cat: auroc(eval_scores[eval_c == cat], eval_y[eval_c == cat]) for cat in categories}
    arb_per["macro"] = float(np.nanmean(list(arb_per.values())))
    # hard gating: per-sample argmax strategy (policy-gated mixture = soft gating)
    hard_scores = []
    for cat in categories:
        ev = data[cat]["eval"]
        vs = data[cat]["eval_acl"]
        c = np.full(len(vs), cat)
        s, P, S = arb.fused_scores(vs, c, norm)
        hard = S[np.arange(len(s)), P.argmax(1)]
        hard_scores.append(hard)
    hard_scores = np.concatenate(hard_scores)
    arb_hard = {cat: auroc(hard_scores[eval_c == cat], eval_y[eval_c == cat]) for cat in categories}
    arb_hard["macro"] = float(np.nanmean(list(arb_hard.values())))
    from collections import Counter
    action_use = dict(Counter(STRATEGY_NAMES[i] for i in eval_policy))

    # ---- continual learning: stream categories, incremental GRPO ----
    stream_orders = [sorted(categories), list(reversed(sorted(categories)))]
    continual = []
    for oi, order in enumerate(stream_orders):
        arb2 = GRPOArbiter(seed=0)
        seen = []
        curve = []
        for cat in order:
            seen.append(cat)
            cal = data[cat]["calib"]
            vs = data[cat]["calib_acl"]
            y = cal["view_labels"].max(1)
            c = np.full(len(vs), cat)
            # replay a bit of previously seen categories (rehearsal)
            if len(seen) > 1:
                prev = [data[pc]["calib_acl"] for pc in seen[:-1]]
                rp_vs = np.concatenate(prev)
                rp_y = np.concatenate([data[pc]["calib"]["view_labels"].max(1) for pc in seen[:-1]])
                rp_c = np.concatenate([np.full(len(p), pc) for p, pc in zip(prev, seen[:-1])])
                sel = np.concatenate([np.where(rp_c == pc)[0][:400] for pc in seen[:-1]])
                vs = np.concatenate([vs, rp_vs[sel]])
                y = np.concatenate([y, rp_y[sel]])
                c = np.concatenate([c, rp_c[sel]])
            arb2.fit(vs, y, c, norm, epochs=15, log=lambda *a, **k: None)
            macro = []
            for cc in categories:
                s, _, _ = arb2.fused_scores(data[cc]["eval_acl"],
                                            np.full(len(data[cc]["eval_acl"]), cc), norm)
                yy = data[cc]["eval"]["view_labels"].max(1)
                macro.append(auroc(s, yy))
            curve.append({"after": cat, "seen": seen.copy(),
                          "macro_sauroc_all6": float(np.nanmean(macro)),
                          "macro_sauroc_seen": float(np.nanmean([macro[categories.index(c2)] for c2 in seen]))})
        continual.append({"order": order, "curve": curve})

    # ---- online hedge on eval stream (operator feedback, delayed labels) ----
    hedge = OnlineHedge()
    hedge_curve = []
    running_y, running_s = [], []
    for cat in categories:
        ev = data[cat]["eval"]
        vs = data[cat]["eval_acl"]
        y = ev["view_labels"].max(1)
        fused = fuse_all(vs)
        Snorm = np.stack([norm.normalize(cat, n, fused[n]) for n in STRATEGY_NAMES], axis=1)
        mix = Snorm @ hedge.w
        running_y.append(y)
        running_s.append(mix)
        yy = np.concatenate(running_y)
        ss = np.concatenate(running_s)
        hedge_curve.append({"after": cat, "macro_sauroc_all": auroc(ss, yy)})
        # operator labels for this category arrive after the batch: reward =
        # each strategy's AUROC on the revealed batch (group-standardized inside)
        R = np.array([auroc(Snorm[:, k], y) for k in range(Snorm.shape[1])])
        hedge.update(R)

    results = {
        "fixed_strategies": table,
        "calib_strategies": calib_table,
        "acl_strategies": acl_table,
        "image_auroc_maxpool": img_auc,
        "grpo_arbiter": {"per_category": arb_per, "hard_gating": arb_hard,
                         "action_usage": action_use,
                         "train_history": arb.history},
        "continual": continual,
        "online_hedge": hedge_curve,
        "acl_stats_source": "results/agent_calibration.json (normal-train, 150 samples/category)",
    }
    with open(out_json, "w") as f:
        json.dump(results, f, indent=1)
    print(json.dumps({"calib_macro": {n: round(calib_table[n]["macro"], 4) for n in calib_table},
                      "fixed_macro": {n: round(table[n]["macro"], 4) for n in table},
                      "acl_macro": {n: round(acl_table[n]["macro"], 4) for n in acl_table},
                      "arbiter_macro": round(arb_per["macro"], 4),
                      "arbiter_hard_macro": round(arb_hard["macro"], 4)}, indent=1))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--stage", choices=["score", "analyze"], required=True)
    ap.add_argument("--categories", nargs="+",
                    default=["usb", "pcb", "mint", "eraser", "bottle_cap", "zipper"])
    ap.add_argument("--n-ok", type=int, default=150)
    ap.add_argument("--n-ng", type=int, default=150)
    ap.add_argument("--n-train", type=int, default=700)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default=os.path.join(RESULTS, "results.json"))
    args = ap.parse_args()
    if args.stage == "score":
        stage_score(args.categories, args.n_ok, args.n_ng, args.n_train, args.seed)
    else:
        stage_analyze(args.categories, args.out)