"""v3j = v3e (Fisher LDA cls + hard-neg + 15 anchors) + v3g (per-anchor refinement), dropping v3h's GMM classifier. Isolates per-anchor refinement from the GMM failure mode in v3i. """ import json import math import os import time import torch import torch.nn.functional as F from torchvision.ops import nms DEVICE = "cuda" if torch.cuda.is_available() else "cpu" TRAIN_CACHE = os.environ["ARENA_CACHE_DIR"] VAL_CACHE = os.environ["ARENA_VAL_CACHE"] COCO_ROOT = os.environ["ARENA_COCO_ROOT"] RESOLUTION = 640 PERSON_CLASS_IDX = 0 PERSON_COCO_CAT_ID = 1 N_SHARDS = 3 RIDGE_LAMBDA = 1.0 N_SCALE_BINS = 5 ASPECT_RATIOS = [0.4, 1.0, 2.5] K_ANCHORS = N_SCALE_BINS * len(ASPECT_RATIOS) H = RESOLUTION // 16 STRIDES = [8, 16, 32, 64, 128] LEVEL_SIZES = [(H * 2, H * 2), (H, H), (H // 2, H // 2), (H // 4, H // 4), (H // 8, H // 8)] LEVEL_COUNTS = [h * w for h, w in LEVEL_SIZES] LEVEL_OFFSETS = [sum(LEVEL_COUNTS[:i]) for i in range(len(LEVEL_COUNTS))] SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) V3D_HEAD = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "analytical_person_v3d", "head_v3d.pt") V3I_HEAD = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "analytical_person_v3i", "head_v3i.pt") OUT_DIR = os.path.join(SCRIPT_DIR, "heads", "cofiber_threshold", "analytical_person_v3j") os.makedirs(OUT_DIR, exist_ok=True) def cofiber_decompose(f, n_scales): cofibers = []; residual = f for _ in range(n_scales - 1): omega = F.avg_pool2d(residual, 2) up = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False) cofibers.append(residual - up); residual = omega cofibers.append(residual) return cofibers def feature_levels(spatial): f = spatial.unsqueeze(0) cs = cofiber_decompose(f, n_scales=4) p3 = F.interpolate(cs[0], scale_factor=2, mode="bilinear", align_corners=False) return {0: p3[0], 1: cs[0][0], 2: cs[1][0], 3: cs[2][0], 4: cs[3][0]} def main(): # Load v3d head for Fisher + anchors + anchor-classifier v3d = torch.load(V3D_HEAD, map_location=DEVICE, weights_only=False) fisher_w = v3d["fisher_w"].to(DEVICE) fisher_b_cal = v3d["fisher_b_cal"] anchors_t = v3d["anchors"].to(DEVICE) cls_W = [w.to(DEVICE) for w in v3d["cls_W"]] # Load v3i head for per-anchor refinement weights v3i = torch.load(V3I_HEAD, map_location=DEVICE, weights_only=False) per_anchor_W = v3i["per_anchor_W"].to(DEVICE) # (15, 770, 4) print(f"v3j: Fisher bias calibrated = {fisher_b_cal:.3f}") print(f" anchors: 15 (5 scale × 3 aspect)") print(f" per-anchor refiner: {tuple(per_anchor_W.shape)} = {per_anchor_W.numel()} params") # Full val print(f"\n{'=' * 72}\nFull val eval: v3j (Fisher LDA + 15 anchors + per-anchor refinement)\n{'=' * 72}") val = torch.load(VAL_CACHE, map_location="cpu", weights_only=False) from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval coco = COCO(os.path.join(COCO_ROOT, "annotations", "instances_val2017.json")) all_results = [] t0 = time.time() for vi, item in enumerate(val): spatial = item["spatial"].float().to(DEVICE) img_id = int(item["img_id"]) scale = item["scale"] feats = feature_levels(spatial) all_boxes = []; all_scores = [] for li in range(5): s = STRIDES[li]; h, w = LEVEL_SIZES[li] f_level = feats[li] flat_n = F.normalize(f_level.permute(1, 2, 0).reshape(-1, 768), p=2, dim=-1) # Fisher cls (discriminative, bayes-calibrated) cls_logit = flat_n @ fisher_w + fisher_b_cal cls_score = torch.sigmoid(cls_logit) ys = (torch.arange(h, device=DEVICE, dtype=torch.float32) + 0.5) * s xs = (torch.arange(w, device=DEVICE, dtype=torch.float32) + 0.5) * s gy, gx = torch.meshgrid(ys, xs, indexing="ij") cx = gx.flatten(); cy = gy.flatten() cx_n = (cx / RESOLUTION) * 2 - 1 cy_n = (cy / RESOLUTION) * 2 - 1 X = torch.cat([flat_n, torch.stack([cx_n, cy_n], dim=-1)], dim=-1) # Anchor classifier anchor_scores = torch.softmax(X @ cls_W[li], dim=-1) a_best, a_id = anchor_scores.max(dim=-1) # Per-anchor refinement W_selected = per_anchor_W[a_id] # (N, 770, 4) residual = torch.einsum("ni,nio->no", X, W_selected) anchor_wh = anchors_t[a_id] box_cx = cx + residual[:, 0] * s box_cy = cy + residual[:, 1] * s box_w = ((anchor_wh[:, 0] + 1.0) * residual[:, 2].exp() - 1.0).clamp(min=0) box_h = ((anchor_wh[:, 1] + 1.0) * residual[:, 3].exp() - 1.0).clamp(min=0) x1 = (box_cx - box_w / 2).clamp(0, RESOLUTION) y1 = (box_cy - box_h / 2).clamp(0, RESOLUTION) x2 = (box_cx + box_w / 2).clamp(0, RESOLUTION) y2 = (box_cy + box_h / 2).clamp(0, RESOLUTION) areas = (x2 - x1) * (y2 - y1) valid = areas > 0 final_score = cls_score * a_best keep = valid & (final_score > 0.005) if keep.any(): idx = keep.nonzero(as_tuple=True)[0] all_boxes.append(torch.stack([x1[idx], y1[idx], x2[idx], y2[idx]], dim=-1)) all_scores.append(final_score[idx]) if not all_boxes: continue boxes = torch.cat(all_boxes, dim=0); scores = torch.cat(all_scores, dim=0) if scores.numel() > 5000: top = scores.topk(5000); scores = top.values; boxes = boxes[top.indices] nms_keep = nms(boxes, scores, 0.5) boxes_f = boxes[nms_keep]; scores_f = scores[nms_keep] if scores_f.numel() > 100: top = scores_f.topk(100); scores_f = top.values; boxes_f = boxes_f[top.indices] for i in range(boxes_f.shape[0]): x1, y1, x2, y2 = boxes_f[i].cpu().tolist() x1 /= scale; y1 /= scale; x2 /= scale; y2 /= scale w = max(0, x2 - x1); h = max(0, y2 - y1) if w == 0 or h == 0: continue all_results.append({ "image_id": img_id, "category_id": PERSON_COCO_CAT_ID, "bbox": [x1, y1, w, h], "score": float(scores_f[i].item()), }) if (vi + 1) % 1000 == 0: print(f" {vi+1}/5000 ({time.time()-t0:.0f}s, {len(all_results)} dets)", flush=True) print(f"\n total detections: {len(all_results)}") if not all_results: return coco_dt = coco.loadRes(all_results) ev = COCOeval(coco, coco_dt, "bbox") ev.params.catIds = [PERSON_COCO_CAT_ID] ev.evaluate(); ev.accumulate(); ev.summarize() out = os.path.join(OUT_DIR, "v3j_full_val.json") with open(out, "w") as f: json.dump({ "head": "analytical_person_v3j", "AP_person": float(ev.stats[0]), "AP50_person": float(ev.stats[1]), "AP75_person": float(ev.stats[2]), "APs_person": float(ev.stats[3]), "APm_person": float(ev.stats[4]), "APl_person": float(ev.stats[5]), }, f, indent=2) print(f"\nSaved: {out}") if __name__ == "__main__": main()