#!/usr/bin/env python3 """AIP-AD experiments on Real-IAD (multi-view, 5 cameras). Stage 'score' : PatchCore-style per-view scoring -> results/scores/.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)