Add arena runner for NYU Depth V2 screening with RMSE eval

arena.py

@@ -0,0 +1,249 @@
+"""
+Depth Head Arena — train and evaluate heads on NYU Depth V2.
+
+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")
+NYU_ROOT = os.environ.get("ARENA_NYU_ROOT", "/home/zootest/datasets/NYU")
+CACHE_DIR = os.environ.get("ARENA_CACHE_DIR", "./arena_cache")
+RESOLUTION = 416
+MIN_DEPTH = 0.001
+MAX_DEPTH = 10.0
+
+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, data_root, split, n_images, cache_path):
+    """Cache backbone features + depth maps."""
+    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)),
+    ])
+
+    # NYU structure: images in sync/ subdirectories
+    img_dir = os.path.join(data_root, split, "images")
+    depth_dir = os.path.join(data_root, split, "depth")
+
+    if not os.path.isdir(img_dir):
+        # Try flat structure
+        img_dir = os.path.join(data_root, split)
+        depth_dir = os.path.join(data_root, split)
+
+    fnames = sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg') or f.endswith('.png')])[:n_images]
+    cached = []
+
+    print(f"  Caching {len(fnames)} {split} 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 depth
+        depth_fname = fname.replace('.jpg', '.png')
+        depth_path = os.path.join(depth_dir, depth_fname)
+        if os.path.isfile(depth_path):
+            depth = np.array(Image.open(depth_path)).astype(np.float32) / 1000.0
+            depth = torch.from_numpy(depth).float()
+            depth = F.interpolate(depth.unsqueeze(0).unsqueeze(0),
+                                  size=(RESOLUTION, RESOLUTION), mode="bilinear",
+                                  align_corners=False)[0, 0]
+        else:
+            depth = torch.zeros(RESOLUTION, RESOLUTION)
+
+        cached.append({"spatial": spatial, "depth": depth})
+        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_rmse(pred, gt, min_depth=MIN_DEPTH, max_depth=MAX_DEPTH):
+    mask = (gt > min_depth) & (gt < max_depth)
+    if mask.sum() == 0:
+        return float("nan")
+    return float(torch.sqrt(((pred[mask] - gt[mask]) ** 2).mean()))
+
+
+def silog_loss(pred, target, variance_focus=0.85):
+    pred = pred.flatten()
+    target = target.flatten()
+    mask = target > MIN_DEPTH
+    pred = pred[mask].clamp(min=1e-6)
+    target = target[mask].clamp(min=1e-6)
+    if pred.numel() == 0:
+        return torch.tensor(0.0, device=pred.device)
+    d = torch.log(pred) - torch.log(target)
+    return torch.sqrt((d ** 2).mean() - variance_focus * (d.mean() ** 2) + 1e-8)
+
+
+def run_candidate(head_name, train_data, val_data, steps=2000, lr=3e-4, 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-3)
+    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()
+        depth_gt = train_data[idx]["depth"].unsqueeze(0).cuda()
+
+        pred = head(spatial)
+        pred_up = F.interpolate(pred, size=depth_gt.shape[1:], mode="bilinear", align_corners=False)
+        loss = silog_loss(pred_up.squeeze(1), depth_gt)
+
+        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()
+    rmses = []
+    with torch.no_grad():
+        for item in val_data:
+            spatial = item["spatial"].unsqueeze(0).float().cuda()
+            depth_gt = item["depth"]
+            pred = head(spatial)
+            pred_up = F.interpolate(pred, size=depth_gt.shape, mode="bilinear", align_corners=False)
+            pred_np = pred_up[0, 0].cpu()
+            rmse = compute_rmse(pred_np, depth_gt)
+            if not np.isnan(rmse):
+                rmses.append(rmse)
+
+    mean_rmse = float(np.mean(rmses)) if rmses else float("nan")
+
+    result = {
+        "name": head_name,
+        "params_M": n_params,
+        "loss_end": np.mean(losses[-10:]) if losses else float("nan"),
+        "rmse": mean_rmse,
+        "train_time_s": time.time() - t0,
+    }
+    print(f"  loss: {result['loss_end']:.4f}, RMSE: {mean_rmse:.4f}", 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("Depth 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, "nyu_train.pt")
+    val_cache = os.path.join(CACHE_DIR, "nyu_val.pt")
+    cache_features(backbone, NYU_ROOT, "train", args.n_train, train_cache)
+    cache_features(backbone, NYU_ROOT, "test", 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)
+            import traceback
+            traceback.print_exc()
+
+    print(f"\n{'='*60}")
+    print(f"{'Name':<25} {'Params':>7} {'RMSE':>8} {'Loss':>7}")
+    print("-" * 50)
+    for r in sorted(results, key=lambda x: x["rmse"] if not np.isnan(x["rmse"]) else 999):
+        print(f"{r['name']:<25} {r['params_M']:>6.2f}M {r['rmse']:>8.4f} {r['loss_end']:>7.3f}")
+
+    out = os.path.join(CACHE_DIR, "depth_results.json")
+    with open(out, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"\nSaved: {out}", flush=True)
+
+
+if __name__ == "__main__":
+    main()