""" Segmentation Head Arena — train and evaluate heads on ADE20K. Usage: python arena.py --head linear_probe --steps 2000 --batch 1 python arena.py --head all --steps 2000 --batch 1 python arena.py --list """ import argparse import json import math import os import sys import time import numpy as np import torch import torch.nn.functional as F from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision.transforms import v2 sys.path.insert(0, os.path.dirname(__file__)) BACKBONE_REPO = os.environ.get("ARENA_BACKBONE_REPO", "/home/zootest/EUPE") BACKBONE_WEIGHTS = os.environ.get("ARENA_BACKBONE_WEIGHTS", "/home/zootest/weights/eupe_vitb/EUPE-ViT-B.pt") BACKBONE_HUB_ENTRY = os.environ.get("ARENA_BACKBONE_ENTRY", "eupe_vitb16") ADE20K_ROOT = os.environ.get("ARENA_ADE20K_ROOT", "/home/zootest/datasets/ADE20K/ADEChallengeData2016") CACHE_DIR = os.environ.get("ARENA_CACHE_DIR", "./arena_cache") RESOLUTION = 512 NUM_CLASSES = 150 if BACKBONE_REPO not in sys.path: sys.path.insert(0, BACKBONE_REPO) from heads import REGISTRY, ALL_NAMES, get_head def cache_features(backbone, img_dir, ann_dir, n_images, cache_path): """Cache backbone features + segmentation labels.""" if os.path.isfile(cache_path): print(f" Cache exists: {cache_path}", flush=True) return normalize = v2.Compose([ v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ]) fnames = sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])[:n_images] cached = [] print(f" Caching {len(fnames)} images...", flush=True) for i, fname in enumerate(fnames): img = Image.open(os.path.join(img_dir, fname)).convert("RGB") img_resized = img.resize((RESOLUTION, RESOLUTION), Image.BILINEAR) x = normalize(img_resized).unsqueeze(0).cuda() with torch.no_grad(): with torch.autocast("cuda", dtype=torch.bfloat16): out = backbone.forward_features(x) patches = out["x_norm_patchtokens"].float() B, N, D = patches.shape h = w = int(N ** 0.5) spatial = patches[0].permute(1, 0).reshape(D, h, w).half().cpu() # Load annotation ann_fname = fname.replace('.jpg', '.png') ann_path = os.path.join(ann_dir, ann_fname) if os.path.isfile(ann_path): ann = np.array(Image.open(ann_path)) ann = ann - 1 # ADE20K: 0=bg(ignored), 1-150=classes -> 0-indexed ann[ann < 0] = 255 label = torch.from_numpy(ann).long() label = F.interpolate(label.unsqueeze(0).unsqueeze(0).float(), size=(RESOLUTION, RESOLUTION), mode="nearest")[0, 0].long() else: label = torch.full((RESOLUTION, RESOLUTION), 255, dtype=torch.long) cached.append({"spatial": spatial, "label": label}) if (i + 1) % 200 == 0: print(f" {i+1}/{len(fnames)}", flush=True) os.makedirs(os.path.dirname(cache_path), exist_ok=True) torch.save(cached, cache_path) print(f" Saved: {cache_path} ({len(cached)} images)", flush=True) def compute_miou(pred, gt, num_classes=150, ignore=255): ious = [] for c in range(num_classes): pc = pred == c gc = gt == c valid = gt != ignore inter = (pc & gc & valid).sum() union = ((pc | gc) & valid).sum() if union > 0: ious.append(float(inter) / float(union)) return float(np.mean(ious)) if ious else 0.0 def run_candidate(head_name, train_data, val_data, steps=2000, lr=1e-3, seed=42): torch.manual_seed(seed) torch.cuda.manual_seed(seed) head = get_head(head_name).cuda() n_params = sum(p.numel() for p in head.parameters()) / 1e6 print(f"\n{'='*60}") print(f" {head_name} ({n_params:.2f}M params)") print(f"{'='*60}", flush=True) optimizer = torch.optim.AdamW(head.parameters(), lr=lr, weight_decay=1e-4) head.train() n = len(train_data) losses = [] t0 = time.time() for step in range(steps): idx = torch.randint(0, n, (1,)).item() spatial = train_data[idx]["spatial"].unsqueeze(0).float().cuda() label = train_data[idx]["label"].unsqueeze(0).cuda() logits = head(spatial) logits_up = F.interpolate(logits, size=label.shape[1:], mode="bilinear", align_corners=False) loss = F.cross_entropy(logits_up, label, ignore_index=255) if torch.isnan(loss) or torch.isinf(loss): optimizer.zero_grad() continue optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(head.parameters(), 5.0) optimizer.step() losses.append(loss.item()) if (step + 1) % 500 == 0: avg = np.mean(losses[-100:]) print(f" step {step+1}/{steps} loss={avg:.4f} ({time.time()-t0:.1f}s)", flush=True) # Eval head.eval() all_pred, all_gt = [], [] with torch.no_grad(): for item in val_data: spatial = item["spatial"].unsqueeze(0).float().cuda() label = item["label"] logits = head(spatial) logits_up = F.interpolate(logits, size=label.shape, mode="bilinear", align_corners=False) pred = logits_up.argmax(dim=1)[0].cpu().numpy() all_pred.append(pred) all_gt.append(label.numpy()) all_pred = np.concatenate([p.flatten() for p in all_pred]) all_gt = np.concatenate([g.flatten() for g in all_gt]) miou = compute_miou(all_pred, all_gt, NUM_CLASSES) result = { "name": head_name, "params_M": n_params, "loss_end": np.mean(losses[-10:]) if losses else float("nan"), "miou": miou, "train_time_s": time.time() - t0, } print(f" loss: {result['loss_end']:.4f}, mIoU: {miou*100:.2f}%", flush=True) del head torch.cuda.empty_cache() return result def main(): parser = argparse.ArgumentParser() parser.add_argument("--head", default="all") parser.add_argument("--steps", type=int, default=2000) parser.add_argument("--list", action="store_true") parser.add_argument("--n-train", type=int, default=500) parser.add_argument("--n-val", type=int, default=100) args = parser.parse_args() if args.list: for name in ALL_NAMES: h = get_head(name) p = sum(v.numel() for v in h.parameters()) / 1e6 print(f" {name:<25} {p:.2f}M") return print("=" * 60) print("Segmentation Head Arena") print("=" * 60, flush=True) print("\nLoading backbone...", flush=True) backbone = torch.hub.load(BACKBONE_REPO, BACKBONE_HUB_ENTRY, source="local", weights=BACKBONE_WEIGHTS) backbone = backbone.cuda().eval() for p in backbone.parameters(): p.requires_grad = False print("\nCaching features...", flush=True) train_cache = os.path.join(CACHE_DIR, "ade20k_train.pt") val_cache = os.path.join(CACHE_DIR, "ade20k_val.pt") cache_features(backbone, os.path.join(ADE20K_ROOT, "images", "training"), os.path.join(ADE20K_ROOT, "annotations", "training"), args.n_train, train_cache) cache_features(backbone, os.path.join(ADE20K_ROOT, "images", "validation"), os.path.join(ADE20K_ROOT, "annotations", "validation"), args.n_val, val_cache) del backbone torch.cuda.empty_cache() train_data = torch.load(train_cache, map_location="cpu", weights_only=False) val_data = torch.load(val_cache, map_location="cpu", weights_only=False) print(f" train: {len(train_data)}, val: {len(val_data)}", flush=True) heads = ALL_NAMES if args.head == "all" else [h.strip() for h in args.head.split(",")] results = [] for name in heads: try: r = run_candidate(name, train_data, val_data, steps=args.steps) results.append(r) except Exception as e: print(f" ERROR: {name}: {e}", flush=True) print(f"\n{'='*60}") print(f"{'Name':<25} {'Params':>7} {'mIoU':>8} {'Loss':>7}") print("-" * 50) for r in sorted(results, key=lambda x: -x["miou"]): print(f"{r['name']:<25} {r['params_M']:>6.2f}M {r['miou']*100:>7.2f}% {r['loss_end']:>7.3f}") out = os.path.join(CACHE_DIR, "seg_results.json") with open(out, "w") as f: json.dump(results, f, indent=2) print(f"\nSaved: {out}", flush=True) if __name__ == "__main__": main()