"""dm-bench evaluation harness (eval_version 1.0). Per page, for a submitted solution {fragment id -> 2x3 affine, fragment px -> page px}: error e(G, i) max over the fragment's convex-hull points p of |G A_pred_i p - A_gt_i p| (worst pixel displacement) direct_acc best global rigid alignment G (hypotheses from each placed fragment, refit by Kabsch on inliers until stable); share of present fragments with e <= tau. Fewer than 2 inliers -> 0, so a single-fragment submission scores exactly 0; random ones score 0 unless two fragments line up by chance. neighbor_acc GT adjacent pairs (shared edge >= MIN_SHARED_PX) whose two fragments fit one rigid G (pair Kabsch fit, either anchor, or the page's global G) within tau. Symmetric; perfect => 1. perfect direct_acc == 1 adjacency P/R/F1 pairs derived from the placed masks (within ADJ_GAP_PX) Hit@1/Hit@5/MRR only if the solution lists ranked partner candidates Poses must be rigid (finite, R^T R ~ I, det > 0); anything else counts as unplaced and is reported as invalid. Missing solution files score 0. Tolerance tau = TAU_FRAC * page width; also reported at TAU_CURVE. """ from __future__ import annotations import math import os from dataclasses import dataclass, field import numpy as np from scipy.ndimage import binary_dilation from scipy.spatial import ConvexHull, QhullError from . import canon EVAL_VERSION = "1.0" SOLUTION_SCHEMA = "dm-bench/solution/1" TAU_FRAC = 0.01 TAU_CURVE = (0.0025, 0.005, 0.01, 0.02) LO_ITERS = 5 ORTHO_TOL = 1e-3 BOOTSTRAP_B = 10000 BOOTSTRAP_SEED = 20260927 def min_shared_px(page_w: int) -> int: """GT pairs sharing less edge than 2*tau carry no matchable seam.""" return int(math.ceil(2 * TAU_FRAC * page_w)) def adj_gap_px(max_erosion: int) -> int: return 2 + 2 * max_erosion # --------------------------------------------------------------------------- # Rigid geometry # --------------------------------------------------------------------------- def to3(a) -> np.ndarray: m = np.eye(3) m[:2] = np.asarray(a, dtype=float) return m def rigid_inverse(m: np.ndarray) -> np.ndarray: """Closed form (R^T, -R^T t); avoids BLAS kernel differences.""" r, t = m[:2, :2], m[:2, 2] out = np.eye(3) out[:2, :2] = r.T out[:2, 2] = -r.T @ t return out def is_rigid(a) -> bool: try: m = np.asarray(a, dtype=float) except (TypeError, ValueError): return False if m.shape != (2, 3) or not np.isfinite(m).all(): return False r = m[:, :2] return bool(np.linalg.norm(r.T @ r - np.eye(2)) <= ORTHO_TOL and (r[0, 0] * r[1, 1] - r[0, 1] * r[1, 0]) > 0) def kabsch2d(src: np.ndarray, dst: np.ndarray) -> np.ndarray: """Best rigid (rotation + translation) map src -> dst, closed form, no SVD.""" cs, cd = src.mean(axis=0), dst.mean(axis=0) s, d = src - cs, dst - cd angle = math.atan2(float((s[:, 0] * d[:, 1] - s[:, 1] * d[:, 0]).sum()), float((s * d).sum())) c, n = math.cos(angle), math.sin(angle) m = np.eye(3) m[:2, :2] = [[c, -n], [n, c]] m[:2, 2] = cd - m[:2, :2] @ cs return m def hull_points(alpha: np.ndarray) -> np.ndarray: """Convex-hull vertices of opaque pixel centres, (k, 2) in fragment px.""" ys, xs = np.nonzero(alpha) pts = np.stack([xs, ys], axis=1).astype(float) if len(pts) < 3: return pts try: return pts[ConvexHull(pts).vertices] except QhullError: # degenerate (collinear) masks return pts[[0, -1]] def _apply(m: np.ndarray, pts: np.ndarray) -> np.ndarray: return pts @ m[:2, :2].T + m[:2, 2] # --------------------------------------------------------------------------- # Per-page scoring # --------------------------------------------------------------------------- @dataclass class PageInputs: page_id: str doc: int canvas: tuple[int, int] gt: dict[str, np.ndarray] # id -> 3x3 hulls: dict[str, np.ndarray] # id -> (k, 2) alphas: dict[str, np.ndarray] # id -> (h, w) uint8 adjacency: list[tuple[str, str, int]] max_erosion: int = 0 @dataclass class PageScore: page_id: str doc: int n_fragments: int placed: int = 0 invalid: int = 0 direct_acc: float = 0.0 neighbor_acc: float = 0.0 perfect: float = 0.0 adj_precision: float = 0.0 adj_recall: float = 0.0 adj_f1: float = 0.0 curve_direct: list[float] = field(default_factory=list) curve_neighbor: list[float] = field(default_factory=list) hit1: float | None = None hit5: float | None = None mrr: float | None = None class _Points: """All hull points of the placed fragments, predicted and GT, stacked once.""" def __init__(self, pred: dict, page: PageInputs): self.ids = sorted(pred) self.pred = {f: _apply(pred[f], page.hulls[f]) for f in self.ids} self.gt = {f: _apply(page.gt[f], page.hulls[f]) for f in self.ids} if self.ids: self.p = np.concatenate([self.pred[f] for f in self.ids]) self.q = np.concatenate([self.gt[f] for f in self.ids]) sizes = [len(self.pred[f]) for f in self.ids] self.starts = np.concatenate([[0], np.cumsum(sizes)[:-1]]).astype(int) def errors(self, g: np.ndarray) -> np.ndarray: """Per-fragment worst displacement under global map g (order = self.ids).""" d = np.linalg.norm(_apply(g, self.p) - self.q, axis=1) return np.maximum.reduceat(d, self.starts) def inliers(self, g: np.ndarray, tau: float) -> list[str]: return [f for f, e in zip(self.ids, self.errors(g)) if e <= tau] def refit(self, ids: list[str]) -> np.ndarray: return kabsch2d(np.concatenate([self.pred[f] for f in ids]), np.concatenate([self.gt[f] for f in ids])) def best_alignment(pred: dict, page: PageInputs, tau: float) -> tuple[np.ndarray | None, int]: """Deterministic LO-RANSAC: every single-fragment hypothesis, refit by Kabsch on its inliers while the inlier set grows.""" if len(pred) < 2: return None, 0 pts = _Points(pred, page) best_g, best_n = None, 0 for k in pts.ids: g = page.gt[k] @ rigid_inverse(pred[k]) inl = pts.inliers(g, tau) for _ in range(LO_ITERS): if len(inl) < 2: break g2 = pts.refit(inl) inl2 = pts.inliers(g2, tau) if len(inl2) <= len(inl): if len(inl2) == len(inl): g = g2 break g, inl = g2, inl2 if len(inl) > best_n: best_g, best_n = g, len(inl) return (best_g, best_n) if best_n >= 2 else (None, 0) def _pair_error(pred: dict, page: PageInputs, g: np.ndarray, ids) -> float: return max(float(np.max(np.linalg.norm(_apply(g @ pred[f], page.hulls[f]) - _apply(page.gt[f], page.hulls[f]), axis=1))) for f in ids) def _pair_ok(pred: dict, page: PageInputs, i: str, j: str, g_global, tau: float) -> bool: fit = kabsch2d(np.concatenate([_apply(pred[f], page.hulls[f]) for f in (i, j)]), np.concatenate([_apply(page.gt[f], page.hulls[f]) for f in (i, j)])) cands = [fit, page.gt[i] @ rigid_inverse(pred[i]), page.gt[j] @ rigid_inverse(pred[j])] if g_global is not None: cands.append(g_global) return any(_pair_error(pred, page, g, (i, j)) <= tau for g in cands) def _placed_adjacency(pred: dict, page: PageInputs, g: np.ndarray) -> set[tuple[str, str]]: """Neighbour pairs of the placed masks, rasterised in the page frame.""" w, h = page.canvas pad = 64 canvas = np.full((h + 2 * pad, w + 2 * pad), -1, dtype=np.int32) ids = sorted(pred) for k, fid in enumerate(ids): ys, xs = np.nonzero(page.alphas[fid]) pp = np.rint(_apply(g @ pred[fid], np.stack([xs, ys], axis=1).astype(float))).astype(int) + pad ok = (pp[:, 0] >= 0) & (pp[:, 0] < canvas.shape[1]) & (pp[:, 1] >= 0) & (pp[:, 1] < canvas.shape[0]) canvas[pp[ok, 1], pp[ok, 0]] = k gap = adj_gap_px(page.max_erosion) struct = np.ones((2 * gap + 1, 2 * gap + 1), dtype=bool) pairs = set() for k, fid in enumerate(ids): ys, xs = np.nonzero(canvas == k) if ys.size == 0: continue # Dilate only a window around the fragment: whole-canvas dilation # per fragment dominated runtime. y0, y1 = max(ys.min() - gap, 0), min(ys.max() + gap + 1, canvas.shape[0]) x0, x1 = max(xs.min() - gap, 0), min(xs.max() + gap + 1, canvas.shape[1]) win = canvas[y0:y1, x0:x1] mask = win == k ring = binary_dilation(mask, structure=struct) & (win >= 0) & ~mask for other in np.unique(win[ring]): a, b = sorted((fid, ids[int(other)])) pairs.add((a, b)) return pairs def score_page(page: PageInputs, solution: dict | None) -> PageScore: n = len(page.gt) score = PageScore(page.page_id, page.doc, n) if not solution: score.curve_direct = [0.0] * len(TAU_CURVE) score.curve_neighbor = [0.0] * len(TAU_CURVE) return score pred, invalid = {}, 0 entries = solution.get("fragments") for fid, entry in (entries if isinstance(entries, dict) else {}).items(): if fid in page.gt and isinstance(entry, dict) and is_rigid(entry.get("affine")): pred[fid] = to3(entry["affine"]) else: invalid += 1 score.placed, score.invalid = len(pred), invalid w = page.canvas[0] min_shared = min_shared_px(w) gt_pairs = [(a, b) for a, b, s in page.adjacency if s >= min_shared] for t_frac in TAU_CURVE: tau = t_frac * w g, n_in = best_alignment(pred, page, tau) direct = n_in / n if n else 0.0 good = sum(1 for a, b in gt_pairs if a in pred and b in pred and _pair_ok(pred, page, a, b, g, tau)) neighbor = good / len(gt_pairs) if gt_pairs else 0.0 score.curve_direct.append(round(direct, 6)) score.curve_neighbor.append(round(neighbor, 6)) if t_frac == TAU_FRAC: score.direct_acc, score.neighbor_acc = direct, neighbor score.perfect = 1.0 if n_in == n else 0.0 _adjacency_scores(score, pred, page, g, gt_pairs) _candidate_scores(score, solution.get("candidates"), gt_pairs) return score def _adjacency_scores(score: PageScore, pred, page, g, gt_pairs) -> None: # Without a valid alignment no two fragments sit correctly relative to each # other, so contacts between the placed masks are meaningless. if len(pred) < 2 or g is None: return placed = _placed_adjacency(pred, page, g) all_gt = {(a, b) for a, b, _ in page.adjacency} filt = set(gt_pairs) tp_all = len(placed & all_gt) score.adj_precision = tp_all / len(placed) if placed else 0.0 score.adj_recall = len(placed & filt) / len(filt) if filt else 0.0 p, r = score.adj_precision, score.adj_recall score.adj_f1 = 2 * p * r / (p + r) if p + r else 0.0 def _candidate_scores(score: PageScore, candidates, gt_pairs) -> None: if not isinstance(candidates, dict): return nbrs: dict[str, set[str]] = {} for a, b in gt_pairs: nbrs.setdefault(a, set()).add(b) nbrs.setdefault(b, set()).add(a) hits1 = hits5 = rr = 0.0 for fid, true in nbrs.items(): raw = candidates.get(fid) raw = raw if isinstance(raw, (list, tuple)) else [] ranked = [c[0] if isinstance(c, (list, tuple)) and c else c for c in raw] ranked = [c for c in ranked if isinstance(c, str)] # malformed entries never match rank = next((k + 1 for k, c in enumerate(ranked) if c in true), None) hits1 += rank == 1 hits5 += rank is not None and rank <= 5 rr += 1.0 / rank if rank else 0.0 m = len(nbrs) or 1 score.hit1, score.hit5, score.mrr = hits1 / m, hits5 / m, rr / m # --------------------------------------------------------------------------- # Loading + aggregation # --------------------------------------------------------------------------- def _read_json(path: str): with open(path, encoding="utf-8") as fh: return canon.loads(fh.read()) def load_page(puzzle_dir: str, answer: dict) -> PageInputs: from PIL import Image puzzle = _read_json(os.path.join(puzzle_dir, "puzzle.json")) alphas, hulls = {}, {} for f in puzzle["fragments"]: with Image.open(os.path.join(puzzle_dir, f["file"])) as im: alpha = np.array(im.getchannel("A")) alphas[f["id"]] = alpha hulls[f["id"]] = hull_points(alpha) gt = {fid: to3(v["affine"]) for fid, v in answer["fragments"].items()} max_ero = max((v.get("erosion_px", 0) for v in answer["fragments"].values()), default=0) return PageInputs( page_id=answer["page_id"], doc=int(answer["doc"]), canvas=tuple(puzzle["canvas"]), gt=gt, hulls=hulls, alphas=alphas, adjacency=[(a, b, int(s)) for a, b, s in answer["adjacency"]], max_erosion=max_ero, ) def load_solution(path: str) -> dict | None: if not os.path.exists(path): return None try: sol = _read_json(path) except ValueError: # includes JSONDecodeError return None return sol if isinstance(sol, dict) else None METRICS = ("direct_acc", "neighbor_acc", "perfect", "adj_precision", "adj_recall", "adj_f1") def aggregate(scores: list[PageScore]) -> dict: """Mean per metric with a document-cluster bootstrap 95% CI (fixed seed).""" scores = sorted(scores, key=lambda s: s.page_id) docs = sorted({s.doc for s in scores}) by_doc = {d: [s for s in scores if s.doc == d] for d in docs} rng = np.random.Generator(np.random.PCG64(BOOTSTRAP_SEED)) draws = rng.integers(0, len(docs), size=(BOOTSTRAP_B, len(docs))) if docs else None out: dict = {"pages": len(scores), "docs": len(docs)} for m in METRICS: per_doc = np.array([[sum(getattr(s, m) for s in by_doc[d]), len(by_doc[d])] for d in docs], dtype=float) mean = float(np.mean([getattr(s, m) for s in scores])) if scores else 0.0 if draws is not None: sums = per_doc[draws, 0].sum(axis=1) / per_doc[draws, 1].sum(axis=1) lo, hi = np.percentile(sums, [2.5, 97.5]) else: lo = hi = 0.0 out[m] = {"mean": round(mean, 4), "ci95": [round(float(lo), 4), round(float(hi), 4)]} for key in ("curve_direct", "curve_neighbor"): arr = np.array([getattr(s, key) for s in scores]) if scores else np.zeros((1, len(TAU_CURVE))) out[key] = {"tau_frac": list(TAU_CURVE), "mean": [round(float(v), 4) for v in arr.mean(axis=0)], "auc": round(float(arr.mean()), 4)} cand = [s for s in scores if s.mrr is not None] if cand: out["hit1"] = round(float(np.mean([s.hit1 for s in cand])), 4) out["hit5"] = round(float(np.mean([s.hit5 for s in cand])), 4) out["mrr"] = round(float(np.mean([s.mrr for s in cand])), 4) out["invalid_poses"] = int(sum(s.invalid for s in scores)) return out def split_answers(release: str, tier: str, split: str, secret_hex: str | None = None) -> dict[str, dict]: """page_id -> answer. Held-out `test` answers are regenerated from the secret.""" if split != "test": adir = os.path.join(release, "answers", tier, split) return {name[:-5]: _read_json(os.path.join(adir, name)) for name in sorted(os.listdir(adir))} if not secret_hex: raise ValueError("scoring the test split requires DM_BENCH_TEST_SECRET") from . import build spec = _read_json(os.path.join(release, "benchmark.json")) n_docs = spec["splits"][tier]["test"]["docs"] entropy = build.secret_entropy(secret_hex) out = {} for doc in range(n_docs): for built in build.build_doc(entropy, tier, "test", doc): out[built.page_id] = canon.loads(canon.dumps(built.answer)) return out def evaluate_split(release: str, solutions: str, tier: str, split: str, secret_hex: str | None = None) -> tuple[list[PageScore], dict]: answers = split_answers(release, tier, split, secret_hex) scores = [] for page_id in sorted(answers): page = load_page(os.path.join(release, "puzzles", tier, split, page_id), answers[page_id]) sol = load_solution(os.path.join(solutions, tier, split, f"{page_id}.json")) scores.append(score_page(page, sol)) return scores, aggregate(scores)