Upload s23dr_2026/eval.py with huggingface_hub

s23dr_2026/eval.py

@@ -0,0 +1,124 @@
+import argparse
+from pathlib import Path
+
+import hoho2025.metric_helper
+import numpy as np
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+
+from .model import get_model, load_checkpoint_compat
+from .inference import predict_wireframe_v2, predict_wireframe_v2_tta
+from .scene import Scene
+from .utils import set_random_seed, start_debug
+
+# Scenes with known pose/annotation misalignment, grouped by severity.
+# Used to produce a clean validation2 split from the public validation set.
+_MISALIGNED_FILE = Path(__file__).resolve().parent.parent / "assets" / "misaligned.txt"
+_SEVERITY_ORDER = ("extreme", "severe", "medium")
+
+
+def load_validation2(stream: bool = False, misalignment_severity: str = "medium") -> object:
+    """Load the validation split filtered to validation2 (misalignment-cleaned subset)."""
+    dataset = load_dataset("usm3d/hoho22k_2026_trainval", trust_remote_code=True, streaming=stream)
+    return _drop_misaligned(dataset["validation"], misalignment_severity)
+
+
+def _drop_misaligned(data, severity: str = "medium"):
+    """Drop scenes with pose/annotation misalignment up to the given severity level."""
+    min_idx = _SEVERITY_ORDER.index(severity)
+    drop: set[str] = set()
+    section = None
+    for line in _MISALIGNED_FILE.read_text().splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            section = line.lstrip("#").strip().lower() if line.startswith("#") else section
+            continue
+        if section in _SEVERITY_ORDER and _SEVERITY_ORDER.index(section) <= min_idx:
+            drop.add(line)
+    print(f"Dropping {len(drop)} misaligned scenes (severity <= {severity})")
+    return data.filter(lambda s: s["order_id"] not in drop)
+
+
+def run_eval(model, dataset, device, args) -> list[tuple[str, object]]:
+    results = []
+    for data in tqdm(dataset, desc="Evaluating"):
+        scene = Scene(data)
+        if args.tta_rotations > 1:
+            verts, edges = predict_wireframe_v2_tta(
+                scene, model, device,
+                pt_type=args.pt_type,
+                num_points=args.num_points,
+                score_threshold=args.threshold,
+                merge_distance_threshold=args.merge_distance_threshold,
+
+                n_rotations=args.tta_rotations,
+            )
+        else:
+            verts, edges = predict_wireframe_v2(
+                scene, model, device,
+                pt_type=args.pt_type,
+                num_points=args.num_points,
+                score_threshold=args.threshold,
+
+                merge_distance_threshold=args.merge_distance_threshold,
+                merge_method=args.merge_method,
+                feature_alpha=args.feature_alpha,
+            )
+        result = hoho2025.metric_helper.hss(verts, edges, scene.verts, scene.edges)
+        results.append((scene.order_id, result))
+        print(f"{scene.order_id}: HSS={result.hss:.4f} F1={result.f1:.4f} IoU={result.iou:.4f}")
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Evaluate WireframeDETR on a dataset split.")
+    parser.add_argument("--split", required=True, help="Dataset split (train / validation / validation2).")
+    parser.add_argument("--checkpoint", type=Path, required=True, help="Trained model checkpoint.")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--seed", type=int, help="Shuffle seed.")
+    parser.add_argument("--stream", action="store_true", help="Stream dataset instead of downloading.")
+    parser.add_argument("-d", "--debug", action="store_true")
+    parser.add_argument("--num_points", "-np", type=int, default=7168)
+    parser.add_argument("--threshold", "-t", type=float, default=0.9, help="Edge confidence threshold.")
+    parser.add_argument("--merge_distance_threshold", type=float, default=0.5)
+    parser.add_argument("--pt_type", type=str, default="colmap", choices=["monodepth", "colmap"])
+    parser.add_argument("--merge_method", default="distance", choices=["distance", "feature", "auto"])
+    parser.add_argument("--feature_alpha", type=float, default=0.95)
+    parser.add_argument("--tta_rotations", type=int, default=1,
+                        help="Y-axis TTA rotations (1=off, 4=0/90/180/270deg).")
+    parser.add_argument("--misalignment_severity", default="medium", choices=list(_SEVERITY_ORDER),
+                        help="Max misalignment severity to drop when split=validation2.")
+    parser.add_argument("--results", type=Path, help="CSV path for per-scene scores.")
+    args = parser.parse_args()
+
+    if args.debug:
+        start_debug()
+
+    if args.split == "validation2":
+        dataset = load_validation2(stream=args.stream, misalignment_severity=args.misalignment_severity)
+    else:
+        dataset = load_dataset("usm3d/hoho22k_2026_trainval", trust_remote_code=True, streaming=args.stream)[args.split]
+
+    set_random_seed(args.seed or 0)
+    if args.seed is not None:
+        dataset = dataset.shuffle(seed=args.seed)
+
+    checkpoint = torch.load(args.checkpoint, map_location=args.device)
+    model = get_model(checkpoint, num_classes=1)
+    load_checkpoint_compat(model, checkpoint)
+    model.to(args.device).eval()
+
+    results = run_eval(model, dataset, args.device, args)
+
+    mean_hss = np.mean([r.hss for _, r in results])
+    mean_f1 = np.mean([r.f1 for _, r in results])
+    mean_iou = np.mean([r.iou for _, r in results])
+    print(f"\nHSS={mean_hss:.4f}  F1={mean_f1:.4f}  IoU={mean_iou:.4f}")
+
+    if args.results:
+        with open(args.results, "w") as f:
+            f.write("order_id,hss,f1,iou\n")
+            for order_id, result in results:
+                f.write(f"{order_id},{result.hss:.4f},{result.f1:.4f},{result.iou:.4f}\n")
+        print(f"Saved to {args.results}")