rpx_benchmark.vqa.sampling

 1"""Deterministic, diversity-preferring row selection shared by the
 2acceptance-manifest and benchmark-plan builders."""
 3
 4from __future__ import annotations
 5
 6import hashlib
 7from typing import Any, Callable
 8
 9from .canonical import semantic_key
10
11
12def rank(seed: int, identity: str) -> str:
13    return hashlib.sha256(f"{seed}:{identity}".encode()).hexdigest()
14
15
16def row_identity(row: dict[str, Any]) -> str:
17    return "\x1f".join(
18        str(row.get(key, ""))
19        for key in ("scene_id", "kind", "phase", "frame", "type", "question", "sample_id")
20    )
21
22
23def is_centered(row: dict[str, Any]) -> bool:
24    """True if the answer bbox centroid is in the middle 50% of both axes.
25    In-context rows already carry a precomputed answer_is_centered column
26    (built with this exact rule at generation time); normal rows do not, so
27    it is computed here from answer_bbox/img_w/img_h."""
28    if row.get("answer_is_centered") is not None:
29        return bool(row["answer_is_centered"])
30    bbox = row.get("answer_bbox")
31    if bbox is None:
32        return False
33    x0, y0, x1, y1 = bbox
34    cx = (x0 + x1) / (2 * row["img_w"])
35    cy = (y0 + y1) / (2 * row["img_h"])
36    return 0.25 <= cx <= 0.75 and 0.25 <= cy <= 0.75
37
38
39class DeterministicSelector:
40    """Greedy, seeded selection across many (cell -> quota) groups, sharing
41    one running state so "prefer distinct scenes/frames" and "no semantic
42    alias duplicate" hold across the whole output, not just within one cell.
43    """
44
45    def __init__(self, seed: int) -> None:
46        self.seed = seed
47        self.used_scene_frame: set[tuple[str, str]] = set()
48        self.used_semantic_keys: set[tuple] = set()
49        self.used_row_ids: set[str] = set()
50
51    def _rank_key(self, row: dict[str, Any]) -> tuple:
52        scene_frame = (str(row["scene_id"]), str(row["frame"]))
53        # Rows whose (scene, frame) is already used sort after rows that
54        # would introduce a new one -- "prefer distinct scenes/frames" as a
55        # soft, deterministic tiebreak, not a hard requirement.
56        reused = scene_frame in self.used_scene_frame
57        return (reused, rank(self.seed, row_identity(row)))
58
59    def select_cell(self, candidates: list[dict[str, Any]], quota: int) -> list[dict[str, Any]]:
60        ordered = sorted(candidates, key=self._rank_key)
61        picked: list[dict[str, Any]] = []
62        for row in ordered:
63            if len(picked) >= quota:
64                break
65            row_id = row_identity(row)
66            if row_id in self.used_row_ids:
67                continue
68            key = semantic_key(row)
69            if key in self.used_semantic_keys:
70                continue
71            picked.append(row)
72            self.used_row_ids.add(row_id)
73            self.used_semantic_keys.add(key)
74            self.used_scene_frame.add((str(row["scene_id"]), str(row["frame"])))
75        return picked
76
77    def select(
78        self,
79        candidates_by_cell: dict[Any, list[dict[str, Any]]],
80        quota_of: Callable[[Any], int],
81    ) -> dict[Any, list[dict[str, Any]]]:
82        """Process cells in a fixed, seed-independent order (sorted by cell
83        key) so the shared used_* state is applied deterministically."""
84        result: dict[Any, list[dict[str, Any]]] = {}
85        for cell in sorted(candidates_by_cell, key=str):
86            result[cell] = self.select_cell(candidates_by_cell[cell], quota_of(cell))
87        return result