Upload benchmark.py with huggingface_hub

benchmark.py

@@ -28,6 +28,7 @@ from __future__ import annotations
 
 import argparse
 import json
+import os
 from pathlib import Path
 
 import numpy as np
@@ -36,6 +37,11 @@ from PIL import Image
 from label_schema import CLASSES, IGNORE
 from provenance import sha256_file
 
+# Public baselines must load without credentials. A stale/invalid HF token in the
+# environment makes huggingface_hub return 401 for public repos, so disable the
+# implicit token outright (and pass token=False at the call sites).
+os.environ.setdefault("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1")
+
 BENCHMARK_VERSION = "1.0"
 TAXONOMY_VERSION = "1.0"
 REDUCED_CLASSES = ("other", "wall", "window")
@@ -140,13 +146,12 @@ def evaluate_painting_vision(checkpoint_path, data, split, size, limit, device):
     from torch.utils.data import DataLoader
 
     from label_schema import mask_to_labels
-    from models import DEFAULT_ARCH, build_model
+    from models import DEFAULT_ARCH, build_from_checkpoint
     from train import MEAN, STD, WallDataset, letterbox
 
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
     arch = str(checkpoint.get("arch", DEFAULT_ARCH))
-    model = build_model(arch, len(CLASSES), pretrained_backbone=False)
-    model.load_state_dict(checkpoint["model"])
+    model = build_from_checkpoint(checkpoint, len(CLASSES))
     model.to(device).eval()
 
     dataset = WallDataset(data, split, size)
@@ -156,7 +161,7 @@ def evaluate_painting_vision(checkpoint_path, data, split, size, limit, device):
     coverage_errors = []
     count = 0
     with torch.inference_mode():
-        for x, y, _ in loader:
+        for x, y, *_ in loader:
             prediction = model(x.to(device))["semantic"].argmax(1).cpu().numpy()
             for p, t in zip(prediction, y.numpy()):
                 full_cm += confusion(p, t, len(CLASSES))
@@ -187,8 +192,8 @@ def evaluate_baseline(name, data, split, limit, device):
 
     if name == "segformer-ade":
         from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
-        processor = AutoImageProcessor.from_pretrained(ADE_BASELINE)
-        model = SegformerForSemanticSegmentation.from_pretrained(ADE_BASELINE).to(device).eval()
+        processor = AutoImageProcessor.from_pretrained(ADE_BASELINE, token=False)
+        model = SegformerForSemanticSegmentation.from_pretrained(ADE_BASELINE, token=False).to(device).eval()
         lut = build_label_lut(model.config.id2label)
 
         def predict(image):
@@ -270,10 +275,14 @@ def main():
         print("no --checkpoint given: skipping the Painting Vision row (train one first)\n")
 
     if args.baseline != "none":
-        import torch
-        device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
-        results[f"baseline:{args.baseline}"] = evaluate_baseline(args.baseline, args.data, args.split,
-                                                                 args.limit, device)
+        try:
+            import torch
+            device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
+            results[f"baseline:{args.baseline}"] = evaluate_baseline(args.baseline, args.data, args.split,
+                                                                     args.limit, device)
+        except Exception as exc:  # an optional comparison must not invalidate our own result
+            results[f"baseline:{args.baseline}"] = {"error": f"{type(exc).__name__}: {exc}"}
+            print(f"baseline {args.baseline!r} could not run and was skipped: {type(exc).__name__}: {exc}\n")
 
     results["floor:all-other"] = evaluate_floor(args.data, args.split, args.limit)