eval_coco_map.py

"""COCO val2017 mAP eval for split-tower detection heads on frozen EUPE-ViT-B.

Default output is hard NMS (the standard reported in the FCOS / YOLO / DETR
literature). Pass `--soft-nms` to additionally report a soft-NMS comparison
pass on the same forward outputs.

Auto-detects architecture flags from the checkpoint's state_dict:
hidden width, n_std / n_dw tower depths, per_scale_bias, reg_on_raw.

Usage:
    python eval_coco_map.py --picker <ckpt.pth>
    python eval_coco_map.py --picker <ckpt.pth> --soft-nms
"""
import argparse
import os
import re
import sys
import time

os.environ.setdefault("COCO_TEXT_EMBED_PATH", "/mnt/d/detection-heads/_diag_text_embed_vitb32.pt")

import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from train_split_tower_5scale import SplitTowerHead, make_locations, STRIDES

DEVICE = "cuda"
COCO_ROOT = os.environ["ARENA_COCO_ROOT"]
VAL_CACHE = os.environ["ARENA_VAL_CACHE"]
RESOLUTION = 640
H = RESOLUTION // 16
NUM_CLASSES = 80
SCORE_THRESH = 0.05
MAX_PER_IMAGE = 100


def hard_nms_per_class(boxes, scores, labels, iou_thresh=0.5):
    """Per-class hard NMS via torchvision.ops.nms. Returns indices into input."""
    from torchvision.ops import nms
    keep = []
    for c in labels.unique():
        m = labels == c
        idx = m.nonzero(as_tuple=True)[0]
        k = nms(boxes[m], scores[m], iou_thresh)
        keep.append(idx[k])
    return torch.cat(keep) if keep else torch.tensor([], dtype=torch.long, device=boxes.device)


def soft_nms_linear(boxes, scores, iou_thresh=0.5, min_score=0.001):
    """GPU-vectorized linear-decay soft NMS (Bodla et al. 2017).
    Iteratively picks the max-score box, multiplicatively decays overlapping
    scores by (1 - IoU), repeats until the max falls below min_score."""
    if boxes.numel() == 0:
        return torch.tensor([], dtype=torch.long, device=boxes.device), scores.clone()
    N = boxes.shape[0]
    scores = scores.clone()
    areas = (boxes[:, 2] - boxes[:, 0]).clamp(min=0) * (boxes[:, 3] - boxes[:, 1]).clamp(min=0)
    active = torch.ones(N, dtype=torch.bool, device=boxes.device)
    kept_idx, kept_scores = [], []
    while active.any():
        masked = torch.where(active, scores, torch.full_like(scores, -float("inf")))
        i = int(masked.argmax())
        if scores[i] < min_score:
            break
        kept_idx.append(i); kept_scores.append(scores[i].item())
        active[i] = False
        box_i = boxes[i]
        x1 = torch.maximum(box_i[0], boxes[:, 0]); y1 = torch.maximum(box_i[1], boxes[:, 1])
        x2 = torch.minimum(box_i[2], boxes[:, 2]); y2 = torch.minimum(box_i[3], boxes[:, 3])
        inter = (x2 - x1).clamp(min=0) * (y2 - y1).clamp(min=0)
        iou = inter / (areas[i] + areas - inter).clamp(min=1e-6)
        decay = torch.where(iou > iou_thresh, 1 - iou, torch.ones_like(iou))
        scores = scores * decay
    if not kept_idx:
        return (torch.tensor([], dtype=torch.long, device=boxes.device),
                torch.tensor([], device=boxes.device))
    return (torch.tensor(kept_idx, dtype=torch.long, device=boxes.device),
            torch.tensor(kept_scores, device=boxes.device))


def soft_nms_per_class(boxes, scores, labels, iou_thresh=0.5):
    keep_idx, keep_scores = [], []
    for c in labels.unique():
        m = labels == c
        idx = m.nonzero(as_tuple=True)[0]
        k, rescored = soft_nms_linear(boxes[m], scores[m], iou_thresh)
        keep_idx.append(idx[k]); keep_scores.append(rescored)
    if not keep_idx:
        return (torch.tensor([], dtype=torch.long, device=boxes.device),
                torch.tensor([], device=boxes.device))
    return torch.cat(keep_idx), torch.cat(keep_scores)


def autodetect_and_load(ckpt_path):
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt["head"] if isinstance(ckpt, dict) and "head" in ckpt else ckpt
    per_scale_bias = (sd["cls_bias"].shape == (5, 80))
    reg_on_raw = any("reg_stem" in k for k in sd.keys())
    hidden = sd["stem.weight"].shape[0]
    tower_idxs = sorted({int(m.group(1)) for k in sd if (m := re.match(r"cls_tower\.(\d+)\.", k))})
    n_std = sum(1 for i in tower_idxs if f"cls_tower.{i}.conv.weight" in sd)
    n_dw = len(tower_idxs) - n_std
    return sd, hidden, n_std, n_dw, per_scale_bias, reg_on_raw


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--picker", required=True, help="Path to picker .pth checkpoint")
    parser.add_argument("--soft-nms", action="store_true",
                        help="Also run a soft-NMS pass and report both. Default is hard NMS only.")
    args = parser.parse_args()

    t_start = time.time()
    print(f"[{time.time()-t_start:6.1f}s] START  pid={os.getpid()}", flush=True)
    print(f"[{time.time()-t_start:6.1f}s] picker={args.picker}  soft_nms={args.soft_nms}", flush=True)

    print(f"[{time.time()-t_start:6.1f}s] Inspecting checkpoint ...", flush=True)
    sd, hidden, n_std, n_dw, per_scale_bias, reg_on_raw = autodetect_and_load(args.picker)
    print(f"[{time.time()-t_start:6.1f}s] Arch: hidden={hidden} n_std={n_std} n_dw={n_dw} "
          f"per_scale_bias={per_scale_bias} reg_on_raw={reg_on_raw}", flush=True)

    print(f"[{time.time()-t_start:6.1f}s] Constructing head ...", flush=True)
    head = SplitTowerHead(hidden=hidden, n_std_layers=n_std, n_dw_layers=n_dw, n_scales=4,
                          reg_on_raw=reg_on_raw, per_scale_bias=per_scale_bias).to(DEVICE)
    head.load_state_dict(sd); head.eval()
    print(f"[{time.time()-t_start:6.1f}s] Head loaded: {sum(p.numel() for p in head.parameters()):,} params", flush=True)

    print(f"[{time.time()-t_start:6.1f}s] Loading val cache ...", flush=True)
    val = torch.load(VAL_CACHE, map_location="cpu", weights_only=False)
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval
    coco_gt = COCO(os.path.join(COCO_ROOT, "annotations", "instances_val2017.json"))
    cat_ids = sorted(coco_gt.getCatIds())
    idx_to_cat = {i: c for i, c in enumerate(cat_ids)}
    print(f"[{time.time()-t_start:6.1f}s] Val: {len(val)} items, {len(cat_ids)} categories", flush=True)

    feat_sizes = [(H*2, H*2), (H, H), (H//2, H//2), (H//4, H//4), (H//8, H//8)]
    all_locs = torch.cat(make_locations(feat_sizes, STRIDES, torch.device(DEVICE)))

    print(f"[{time.time()-t_start:6.1f}s] ==================== INFERENCE START ====================", flush=True)
    results_hard = []
    results_soft = [] if args.soft_nms else None
    t0 = time.time()
    PRINT_EVERY = 500

    with torch.no_grad():
        for idx, item in enumerate(val):
            spatial = item["spatial"].unsqueeze(0).float().to(DEVICE)
            img_id = int(item["img_id"]); scale = item["scale"]
            cls_l, reg_l, ctr_l = head(spatial)

            cls_s = torch.cat([c.permute(0,2,3,1).reshape(-1, NUM_CLASSES) for c in cls_l]).sigmoid()
            reg_s = torch.cat([r.permute(0,2,3,1).reshape(-1, 4) for r in reg_l])
            ctr_s = torch.cat([c.permute(0,2,3,1).reshape(-1) for c in ctr_l]).sigmoid()
            scores_all = cls_s * ctr_s.unsqueeze(1)

            mask = scores_all > SCORE_THRESH
            if not mask.any():
                if (idx+1) % PRINT_EVERY == 0:
                    elapsed = time.time()-t0; rate = (idx+1)/elapsed
                    print(f"[{time.time()-t_start:6.1f}s]   img {idx+1}/{len(val)} | "
                          f"hard={len(results_hard):,}"
                          + (f" soft={len(results_soft):,}" if args.soft_nms else "")
                          + f" | {rate:.1f} img/s | ETA {(len(val)-idx-1)/rate:.0f}s", flush=True)
                continue

            loc_idx, cls_idx = mask.nonzero(as_tuple=True)
            sc = scores_all[loc_idx, cls_idx]
            xy = all_locs[loc_idx]; rr = reg_s[loc_idx]
            x1 = (xy[:,0]-rr[:,0]).clamp(0, RESOLUTION); y1 = (xy[:,1]-rr[:,1]).clamp(0, RESOLUTION)
            x2 = (xy[:,0]+rr[:,2]).clamp(0, RESOLUTION); y2 = (xy[:,1]+rr[:,3]).clamp(0, RESOLUTION)
            boxes = torch.stack([x1, y1, x2, y2], dim=-1)

            keep_hard = hard_nms_per_class(boxes, sc, cls_idx, iou_thresh=0.5)
            _append_dets(results_hard, boxes, sc[keep_hard], keep_hard, cls_idx, scale, idx_to_cat, img_id)

            if args.soft_nms:
                keep_soft, sc_soft = soft_nms_per_class(boxes, sc, cls_idx, iou_thresh=0.5)
                _append_dets(results_soft, boxes, sc_soft, keep_soft, cls_idx, scale, idx_to_cat, img_id)

            if (idx+1) % PRINT_EVERY == 0:
                elapsed = time.time()-t0; rate = (idx+1)/elapsed
                msg = (f"[{time.time()-t_start:6.1f}s]   img {idx+1}/{len(val)} | "
                       f"hard={len(results_hard):,}")
                if args.soft_nms:
                    msg += f" soft={len(results_soft):,}"
                msg += f" | n_pos={mask.sum().item()} | {rate:.1f} img/s | ETA {(len(val)-idx-1)/rate:.0f}s"
                print(msg, flush=True)

    print(f"\n[{time.time()-t_start:6.1f}s] Inference done in {time.time()-t0:.0f}s", flush=True)

    passes = [("HARD NMS", results_hard)]
    if args.soft_nms:
        passes.append(("SOFT NMS", results_soft))
    for label, results in passes:
        print(f"\n{'='*40}\n{label}\n{'='*40}", flush=True)
        coco_dt = coco_gt.loadRes(results)
        ev = COCOeval(coco_gt, coco_dt, "bbox")
        ev.params.imgIds = sorted(coco_gt.getImgIds())[:len(val)]
        ev.evaluate(); ev.accumulate(); ev.summarize()


def _append_dets(sink, boxes, scores_kept, keep, cls_idx, scale, idx_to_cat, img_id):
    if keep.numel() == 0:
        return
    if scores_kept.numel() > MAX_PER_IMAGE:
        top = scores_kept.topk(MAX_PER_IMAGE)
        scores_kept = top.values; keep = keep[top.indices]
    bx = boxes[keep] / scale
    w = (bx[:,2]-bx[:,0]).clamp(min=0); h = (bx[:,3]-bx[:,1]).clamp(min=0)
    for i in range(keep.numel()):
        s = scores_kept[i].item()
        if s < SCORE_THRESH:
            continue
        sink.append({"image_id": img_id,
                     "category_id": idx_to_cat[cls_idx[keep[i]].item()],
                     "bbox": [bx[i,0].item(), bx[i,1].item(), w[i].item(), h[i].item()],
                     "score": s})


if __name__ == "__main__":
    main()