""" Build and evaluate one analytical detection head variant. Usage: python analytical_one.py --name baseline python analytical_one.py --name whitened --transform zca python analytical_one.py --name spatial3x3 --spatial mean3x3 """ import argparse import json import math import os import sys import time import torch import torch.nn.functional as F import numpy as np SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, SCRIPT_DIR) CACHE_DIR = os.environ.get("ARENA_CACHE_DIR", "feature_cache") COCO_ROOT = os.environ.get("ARENA_COCO_ROOT", "coco") VAL_CACHE = os.environ.get("ARENA_VAL_CACHE", "val_cache/val.pt") STATS_DIR = os.path.join(SCRIPT_DIR, "analytical_stats_cache") RESULTS_DIR = os.path.join(SCRIPT_DIR, "analytical_variants") RESOLUTION = 640 NUM_CLASSES = 80 def cofiber_decompose(f, n_scales): cofibers = [] residual = f for _ in range(n_scales - 1): omega = F.avg_pool2d(residual, 2) sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False) cofibers.append(residual - sigma_omega) residual = omega cofibers.append(residual) return cofibers def make_locations(sizes, strides): locs = [] for (h, w), s in zip(sizes, strides): ys = (torch.arange(h, dtype=torch.float32) + 0.5) * s xs = (torch.arange(w, dtype=torch.float32) + 0.5) * s gy, gx = torch.meshgrid(ys, xs, indexing="ij") locs.append(torch.stack([gx.flatten(), gy.flatten()], -1)) return locs def assign_targets(loc, boxes, labels, stride, sr): n = loc.shape[0] if boxes.numel() == 0: return torch.full((n,), -1, dtype=torch.long), torch.zeros(n, 4), torch.zeros(n) areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) l = loc[:, None, 0] - boxes[None, :, 0] t = loc[:, None, 1] - boxes[None, :, 1] r = boxes[None, :, 2] - loc[:, None, 0] b = boxes[None, :, 3] - loc[:, None, 1] ltrb = torch.stack([l, t, r, b], -1) in_box = ltrb.min(-1).values > 0 cx = (boxes[:, 0] + boxes[:, 2]) / 2 cy = (boxes[:, 1] + boxes[:, 3]) / 2 rad = stride * 1.5 in_center = ((loc[:, None, 0] >= cx - rad) & (loc[:, None, 0] <= cx + rad) & (loc[:, None, 1] >= cy - rad) & (loc[:, None, 1] <= cy + rad)) max_d = ltrb.max(-1).values in_level = (max_d >= sr[0]) & (max_d <= sr[1]) pos = in_box & in_center & in_level a = areas[None, :].expand_as(pos).clone() a[~pos] = float("inf") matched = a.argmin(1) is_pos = a.gather(1, matched[:, None]).squeeze(1) < float("inf") ct = torch.full((n,), -1, dtype=torch.long) ct[is_pos] = labels[matched[is_pos]] rt = torch.zeros(n, 4) if is_pos.any(): rt[is_pos] = ltrb[torch.arange(n)[is_pos], matched[is_pos]] ctrt = torch.zeros(n) if is_pos.any(): lp, tp, rp, bp = rt[is_pos].unbind(-1) ctrt[is_pos] = torch.sqrt( (torch.minimum(lp, rp) / torch.maximum(lp, rp).clamp(min=1e-6)) * (torch.minimum(tp, bp) / torch.maximum(tp, bp).clamp(min=1e-6))) return ct, rt, ctrt # ============================================================ # Feature transforms # ============================================================ def transform_layernorm(f, C): return F.layer_norm(f, [C]) def transform_raw(f, C): return f def transform_l2norm(f, C): return F.normalize(f, p=2, dim=-1) def transform_power05(f, C): """Signed power normalization: sign(x) * |x|^0.5, then L2 normalize.""" out = f.sign() * f.abs().sqrt() return F.normalize(out, p=2, dim=-1) def transform_power025(f, C): """Stronger compression: sign(x) * |x|^0.25.""" out = f.sign() * f.abs().pow(0.25) return F.normalize(out, p=2, dim=-1) TRANSFORMS = { "layernorm": transform_layernorm, "raw": transform_raw, "l2norm": transform_l2norm, "power05": transform_power05, "power025": transform_power025, } # ============================================================ # Target encodings # ============================================================ def encode_log_ltrb(ltrb): valid = (ltrb > 0).all(1) out = torch.zeros_like(ltrb) if valid.any(): out[valid] = torch.log(ltrb[valid]) return out, valid def encode_sqrt_ltrb(ltrb): valid = (ltrb > 0).all(1) out = torch.zeros_like(ltrb) if valid.any(): out[valid] = torch.sqrt(ltrb[valid]) return out, valid def encode_ltrb(ltrb): valid = (ltrb > 0).all(1) return ltrb, valid def encode_corners(ltrb): valid = (ltrb > 0).all(1) return torch.stack([-ltrb[:, 0], -ltrb[:, 1], ltrb[:, 2], ltrb[:, 3]], 1), valid def encode_center_size(ltrb): valid = (ltrb > 0).all(1) cx_off = (ltrb[:, 2] - ltrb[:, 0]) / 2 cy_off = (ltrb[:, 3] - ltrb[:, 1]) / 2 w = ltrb[:, 0] + ltrb[:, 2] h = ltrb[:, 1] + ltrb[:, 3] return torch.stack([cx_off, cy_off, w, h], 1), valid ENCODINGS = { "log_ltrb": encode_log_ltrb, "sqrt_ltrb": encode_sqrt_ltrb, "ltrb": encode_ltrb, "corners": encode_corners, "center_size": encode_center_size, } # ============================================================ # Spatial context modes # ============================================================ def spatial_none(f_grid, B, H, W, C): """No spatial context. Per-token features only.""" return f_grid.reshape(-1, C) def spatial_mean3x3(f_grid, B, H, W, C): """Replace each token with the mean of its 3x3 neighborhood.""" f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) pooled = F.avg_pool2d(f_4d, 3, stride=1, padding=1) return pooled.permute(0, 2, 3, 1).reshape(-1, C) def spatial_cat_mean3x3(f_grid, B, H, W, C): """Concatenate: [center_token, mean_of_3x3_neighborhood]. 2*C dims.""" f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) pooled = F.avg_pool2d(f_4d, 3, stride=1, padding=1) center = f_grid.reshape(-1, C) neighbor_mean = pooled.permute(0, 2, 3, 1).reshape(-1, C) return torch.cat([center, neighbor_mean], dim=1) def spatial_diff_neighbors(f_grid, B, H, W, C): """Center token + (center - neighbor_mean). Emphasizes local contrast.""" f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) pooled = F.avg_pool2d(f_4d, 3, stride=1, padding=1) center = f_grid.reshape(-1, C) diff = center - pooled.permute(0, 2, 3, 1).reshape(-1, C) return torch.cat([center, diff], dim=1) def spatial_hv_neighbors(f_grid, B, H, W, C): """Center + horizontal mean + vertical mean. 3*C dims.""" f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) h_pool = F.avg_pool2d(f_4d, (1, 3), stride=1, padding=(0, 1)) v_pool = F.avg_pool2d(f_4d, (3, 1), stride=1, padding=(1, 0)) center = f_grid.reshape(-1, C) h_mean = h_pool.permute(0, 2, 3, 1).reshape(-1, C) v_mean = v_pool.permute(0, 2, 3, 1).reshape(-1, C) return torch.cat([center, h_mean, v_mean], dim=1) def spatial_sheaf_h1(f_grid, B, H, W, C): """Sheaf H^1: directional edge differences (4 cardinal Cech 1-cocycles). At each location, compute feature[here] - feature[neighbor] for all 4 cardinal directions. This is the Cech 1-cocycle representative on the spatial grid. It captures gluing obstructions — exactly where local feature sections fail to extend consistently. Object boundaries are such obstructions. Output: [center, d_up, d_down, d_left, d_right] = 5*C dims. """ f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) # (B, C, H, W) # Shift in each direction and subtract d_up = f_4d - F.pad(f_4d[:, :, 1:, :], (0, 0, 0, 1)) # diff with token above d_down = f_4d - F.pad(f_4d[:, :, :-1, :], (0, 0, 1, 0)) # diff with token below d_left = f_4d - F.pad(f_4d[:, :, :, 1:], (0, 1, 0, 0)) # diff with token left d_right = f_4d - F.pad(f_4d[:, :, :, :-1], (1, 0, 0, 0)) # diff with token right center = f_grid.reshape(-1, C) du = d_up.permute(0, 2, 3, 1).reshape(-1, C) dd = d_down.permute(0, 2, 3, 1).reshape(-1, C) dl = d_left.permute(0, 2, 3, 1).reshape(-1, C) dr = d_right.permute(0, 2, 3, 1).reshape(-1, C) return torch.cat([center, du, dd, dl, dr], dim=1) def spatial_sheaf_h1_compact(f_grid, B, H, W, C): """Sheaf H^1 compact: center + L1 norm of directional cocycles. Instead of raw directional differences (5*C dims), compute the L1 magnitude of each directional cocycle per channel. This gives a scalar "boundary strength" per channel per direction. Output: [center, |d_up|+|d_down|, |d_left|+|d_right|] = 3*C dims. Vertical and horizontal boundary strengths. """ f_4d = f_grid.reshape(B, H, W, C).permute(0, 3, 1, 2) d_up = f_4d - F.pad(f_4d[:, :, 1:, :], (0, 0, 0, 1)) d_down = f_4d - F.pad(f_4d[:, :, :-1, :], (0, 0, 1, 0)) d_left = f_4d - F.pad(f_4d[:, :, :, 1:], (0, 1, 0, 0)) d_right = f_4d - F.pad(f_4d[:, :, :, :-1], (1, 0, 0, 0)) center = f_grid.reshape(-1, C) v_boundary = (d_up.abs() + d_down.abs()).permute(0, 2, 3, 1).reshape(-1, C) h_boundary = (d_left.abs() + d_right.abs()).permute(0, 2, 3, 1).reshape(-1, C) return torch.cat([center, v_boundary, h_boundary], dim=1) SPATIAL = { "none": spatial_none, "mean3x3": spatial_mean3x3, "cat_mean3x3": spatial_cat_mean3x3, "diff_neighbors": spatial_diff_neighbors, "hv_neighbors": spatial_hv_neighbors, "sheaf_h1": spatial_sheaf_h1, "sheaf_h1_compact": spatial_sheaf_h1_compact, } # ============================================================ # Core: accumulate, solve, evaluate # ============================================================ def accumulate(n_images, transform_name, encoding_name, spatial_name, per_scale=False, neg_ratio=0.0): """Accumulate XtX/XtY from cached features.""" manifest = json.load(open(os.path.join(CACHE_DIR, "manifest.json"))) n_shards = manifest["n_shards"] strides = [16, 32, 64] H = RESOLUTION // 16 sizes = [(H, H), (H // 2, H // 2), (H // 4, H // 4)] sr = [(-1, 128), (128, 256), (256, float("inf"))] locs = make_locations(sizes, strides) transform_fn = TRANSFORMS[transform_name] encode_fn = ENCODINGS[encoding_name] spatial_fn = SPATIAL[spatial_name] # Determine feature dim test_f = torch.randn(1, 768) test_t = transform_fn(test_f, 768) test_s = spatial_fn(test_t.unsqueeze(0), 1, 1, 1, test_t.shape[-1]) feat_dim = test_s.shape[-1] n_scales = 3 if not per_scale else 1 scale_range = range(3) if not per_scale else [0] cls_XtX = torch.zeros(feat_dim + 1, feat_dim + 1) cls_XtY = torch.zeros(feat_dim + 1, NUM_CLASSES) reg_XtX = torch.zeros(feat_dim + 1, feat_dim + 1) reg_XtY = torch.zeros(feat_dim + 1, 4) ctr_XtX = torch.zeros(feat_dim + 1, feat_dim + 1) ctr_XtY = torch.zeros(feat_dim + 1, 1) n_pos = 0 seen = 0 t0 = time.time() for si in range(n_shards): if seen >= n_images: break shard = torch.load(os.path.join(CACHE_DIR, f"shard_{si:04d}.pt"), map_location="cpu", weights_only=False) for item in shard: if seen >= n_images: break sp = item["spatial"].unsqueeze(0).float() boxes = item["boxes"] labels = item["labels"] cofibers = cofiber_decompose(sp, 3) for sci in range(3): cof = cofibers[sci] B, C, Hc, Wc = cof.shape f_raw = cof.permute(0, 2, 3, 1).reshape(-1, C) f = transform_fn(f_raw, C) f_spatial = spatial_fn(f.unsqueeze(0) if f.dim() == 2 else f, B, Hc, Wc, f.shape[-1] if f.dim() == 2 else C) ct, rt, ctrt = assign_targets(locs[sci], boxes, labels, strides[sci], sr[sci]) pos = ct >= 0 if not pos.any(): continue fp = f_spatial[pos] fa = torch.cat([fp, torch.ones(fp.shape[0], 1)], 1) y_cls = torch.zeros(fp.shape[0], NUM_CLASSES) y_cls[torch.arange(fp.shape[0]), ct[pos]] = 1.0 cls_XtX += fa.T @ fa cls_XtY += fa.T @ y_cls reg_y, valid = encode_fn(rt[pos]) if valid.any(): fr = fa[valid] reg_XtX += fr.T @ fr reg_XtY += fr.T @ reg_y[valid] ctr_XtX += fa.T @ fa ctr_XtY += fa.T @ ctrt[pos].unsqueeze(1) n_pos += pos.sum().item() # Negative samples for classification (target = all zeros) if neg_ratio > 0: neg = ct < 0 n_neg_want = int(pos.sum().item() * neg_ratio) if neg.any() and n_neg_want > 0: neg_idx = neg.nonzero(as_tuple=True)[0] if len(neg_idx) > n_neg_want: neg_idx = neg_idx[torch.randperm(len(neg_idx))[:n_neg_want]] fn = f_spatial[neg_idx] fn_aug = torch.cat([fn, torch.ones(fn.shape[0], 1)], 1) cls_XtX += fn_aug.T @ fn_aug cls_XtY += fn_aug.T @ torch.zeros(fn.shape[0], NUM_CLASSES) seen += 1 del shard if (si + 1) % 5 == 0: elapsed = time.time() - t0 print(f" shard {si+1}: {seen} imgs, {n_pos} pos, {elapsed:.0f}s", flush=True) return {"cls_XtX": cls_XtX, "cls_XtY": cls_XtY, "reg_XtX": reg_XtX, "reg_XtY": reg_XtY, "ctr_XtX": ctr_XtX, "ctr_XtY": ctr_XtY, "n_pos": n_pos, "feat_dim": feat_dim, "n_images": seen, "elapsed": time.time() - t0} def solve(stats, lam): fd = stats["feat_dim"] n = stats["n_pos"] I = torch.eye(fd + 1) cls_W = torch.linalg.solve(stats["cls_XtX"] + lam * I * n, stats["cls_XtY"]) reg_W = torch.linalg.solve(stats["reg_XtX"] + lam * I * n, stats["reg_XtY"]) ctr_W = torch.linalg.solve(stats["ctr_XtX"] + lam * I * n, stats["ctr_XtY"]) return {"cls_w": cls_W[:fd].T, "cls_b": cls_W[fd], "reg_w": reg_W[:fd].T, "reg_b": reg_W[fd], "ctr_w": ctr_W[:fd].T, "ctr_b": ctr_W[fd], "feat_dim": fd} def evaluate(head, val_path, transform_name, spatial_name, encoding_name="log_ltrb", n_images=500): """Evaluate on val set. CPU-only. Returns metrics dict.""" val = torch.load(val_path, map_location="cpu", weights_only=False) encode_fn = ENCODINGS[encoding_name] # Load COCO GT from pycocotools.coco import COCO ann_file = os.path.join(COCO_ROOT, "annotations", "instances_val2017.json") coco = COCO(ann_file) cat_ids = sorted(coco.getCatIds()) cat_to_idx = {c: i for i, c in enumerate(cat_ids)} transform_fn = TRANSFORMS[transform_name] spatial_fn = SPATIAL[spatial_name] strides = [16, 32, 64] H = RESOLUTION // 16 sizes = [(H, H), (H // 2, H // 2), (H // 4, H // 4)] sr = [(-1, 128), (128, 256), (256, float("inf"))] locs = make_locations(sizes, strides) correct = 0 n_pos_total = 0 n_det = 0 true_det = 0 reg_errors = [] for idx in range(min(n_images, len(val))): item = val[idx] spatial = item["spatial"].unsqueeze(0).float() img_id = item["img_id"] scale = item["scale"] ann_ids = coco.getAnnIds(imgIds=int(img_id), iscrowd=False) anns = coco.loadAnns(ann_ids) boxes = [] labels = [] for ann in anns: x, y, w, h = ann["bbox"] if w < 1 or h < 1: continue boxes.append([x * scale, y * scale, (x + w) * scale, (y + h) * scale]) labels.append(cat_to_idx[ann["category_id"]]) boxes_t = torch.tensor(boxes, dtype=torch.float32) if boxes else torch.zeros(0, 4) labels_t = torch.tensor(labels, dtype=torch.long) if labels else torch.zeros(0, dtype=torch.long) cofibers = cofiber_decompose(spatial, 3) for sci, cof in enumerate(cofibers): B, C, Hc, Wc = cof.shape f_raw = cof.permute(0, 2, 3, 1).reshape(-1, C) f = transform_fn(f_raw, C) f_s = spatial_fn(f.unsqueeze(0), B, Hc, Wc, f.shape[-1]) ct, rt, _ = assign_targets(locs[sci], boxes_t, labels_t, strides[sci], sr[sci]) pos = ct >= 0 # Classification scores = f_s @ head["cls_w"].T + head["cls_b"] pred_cls = scores.argmax(1) pred_conf = scores.sigmoid().max(1).values if pos.any(): correct += (pred_cls[pos] == ct[pos]).sum().item() n_pos_total += pos.sum().item() # Detection count det = pred_conf > 0.3 n_det += det.sum().item() true_det += (det & pos).sum().item() # Regression quality (in encoded target space — comparable across encodings) if pos.any(): pred_reg = f_s[pos] @ head["reg_w"].T + head["reg_b"] gt_ltrb = rt[pos] valid = (gt_ltrb > 0).all(1) if valid.any(): gt_encoded, _ = encode_fn(gt_ltrb[valid]) pred_encoded = pred_reg[valid] mse = ((pred_encoded - gt_encoded) ** 2).mean(1) # Convert to quality: 1 / (1 + mse), bounded in [0, 1] quality = (1.0 / (1.0 + mse)).tolist() reg_errors.extend(quality) cls_acc = correct / max(n_pos_total, 1) precision = true_det / max(n_det, 1) reg_quality = sum(reg_errors) / max(len(reg_errors), 1) # mean quality in [0, 1] n_params = (head["cls_w"].numel() + head["cls_b"].numel() + head["reg_w"].numel() + head["reg_b"].numel() + head["ctr_w"].numel() + head["ctr_b"].numel()) return { "cls_accuracy": round(cls_acc, 4), "precision": round(precision, 4), "reg_quality": round(reg_quality, 4), "n_detections": n_det, "n_positives": n_pos_total, "n_params": n_params, "composite": round(cls_acc * 0.5 + precision * 0.25 + reg_quality * 0.25, 4), } def main(): parser = argparse.ArgumentParser() parser.add_argument("--name", required=True, help="Variant name") parser.add_argument("--transform", default="layernorm", choices=list(TRANSFORMS.keys())) parser.add_argument("--encoding", default="log_ltrb", choices=list(ENCODINGS.keys())) parser.add_argument("--spatial", default="none", choices=list(SPATIAL.keys())) parser.add_argument("--lam", type=float, default=1e-3) parser.add_argument("--n-train", type=int, default=10000) parser.add_argument("--n-eval", type=int, default=500) parser.add_argument("--neg-ratio", type=float, default=0.0, help="Ratio of negative to positive samples for classification (0=positives only)") parser.add_argument("--notes", default="", help="Why this variant exists") args = parser.parse_args() os.makedirs(RESULTS_DIR, exist_ok=True) os.makedirs(STATS_DIR, exist_ok=True) print(f"{'='*60}") print(f"Variant: {args.name}") print(f" transform={args.transform} encoding={args.encoding} spatial={args.spatial} lam={args.lam}") if args.notes: print(f" rationale: {args.notes}") print(f"{'='*60}", flush=True) # Check for cached stats neg_tag = f"_neg{args.neg_ratio}" if args.neg_ratio > 0 else "" cache_key = f"stats_s3_{args.transform}_{args.encoding}_{args.spatial}_{args.n_train}{neg_tag}" cache_path = os.path.join(STATS_DIR, f"{cache_key}.pt") if os.path.isfile(cache_path): print(f" Loading cached stats: {cache_key}", flush=True) stats = torch.load(cache_path, map_location="cpu", weights_only=False) else: print(f" Accumulating...", flush=True) stats = accumulate(args.n_train, args.transform, args.encoding, args.spatial, neg_ratio=args.neg_ratio) torch.save(stats, cache_path) print(f" Cached: {cache_path}", flush=True) print(f" {stats['n_pos']} positives, feat_dim={stats['feat_dim']}", flush=True) # Solve t0 = time.time() head = solve(stats, args.lam) solve_time = time.time() - t0 print(f" Solved in {solve_time*1000:.0f}ms", flush=True) # Evaluate print(f" Evaluating ({args.n_eval} images)...", flush=True) t0 = time.time() metrics = evaluate(head, VAL_CACHE, args.transform, args.spatial, args.encoding, args.n_eval) eval_time = time.time() - t0 print(f"\n Results:") print(f" cls_accuracy: {metrics['cls_accuracy']}") print(f" precision: {metrics['precision']}") print(f" reg_quality: {metrics['reg_quality']}") print(f" composite: {metrics['composite']}") print(f" params: {metrics['n_params']}") print(f" eval time: {eval_time:.1f}s") # Save result = { "name": args.name, "config": {"transform": args.transform, "encoding": args.encoding, "spatial": args.spatial, "lam": args.lam, "n_train": args.n_train, "n_eval": args.n_eval}, "notes": args.notes, "metrics": metrics, "solve_time_ms": round(solve_time * 1000), "eval_time_s": round(eval_time, 1), } result_path = os.path.join(RESULTS_DIR, f"{args.name}.json") with open(result_path, "w") as f: json.dump(result, f, indent=2) print(f"\n Saved: {result_path}") if __name__ == "__main__": main()