rpx_benchmark.vqa.contract

  1"""Canonical, model-independent contract for RPX VQA benchmark rows.
  2
  3Supports both one-image ("normal") and two-image ("in-context") samples.
  4Normal rows carry a single target image and are unchanged from the original
  5single-image contract, including the strict locator round-trip check in
  6``VQASample.validate``. In-context rows additionally carry an ordered
  7reference image (Image 1) ahead of the target image (Image 2); the answer
  8bbox always belongs to the target image's own ``img_w``/``img_h``, never the
  9reference's. In-context locators are taken verbatim from the in-context
 10Parquets (``target_image``/``reference_image``) rather than recomputed by
 11``image_locator`` -- those Parquets pin an immutable Hub revision that is
 12unrelated to the consolidated revision the normal contract uses, and
 13reconstructing them would silently disconnect a sample
 14from the exact asset it was generated against.
 15"""
 16
 17from __future__ import annotations
 18
 19import hashlib
 20import json
 21import re
 22from dataclasses import dataclass
 23from pathlib import Path
 24from typing import Any, Iterable
 25
 26from ..exceptions import ManifestError
 27from ..hub import DEFAULT_REVISION
 28
 29SCHEMA_VERSION = "rpx-vqa-1.0"
 30
 31# Pinned immutable Hub revision baked into every in-context Parquet locator
 32# (data/mask_annotation/visual_grounding_gt/pilot/out/release_candidate_v3).
 33# Never replace this with "main" or any other mutable ref.
 34IN_CONTEXT_REVISION = "93e31d378f1f98eca18a7aa01a2279c9f332440c"
 35DATASET_REVISION = DEFAULT_REVISION
 36
 37BINARY_TYPES = frozenset({"spatial_lr_binary", "spatial_ud_binary"})
 38GENERAL_BBOX_TYPES = frozenset(
 39    {
 40        "attr_single_color",
 41        "attr_single_material",
 42        "attr_single_function",
 43        "attr_composition",
 44        "attr_odd_one_out",
 45    }
 46)
 47SPATIAL_BBOX_TYPES = frozenset(
 48    {"spatial_lr_extreme", "depth_closest", "spatial_farthest"}
 49)
 50# In-context (two-image) types. Note attr_odd_one_out has no normal-contract
 51# bbox counterpart in GENERAL_BBOX_TYPES above (the single-image odd-one-out
 52# question is answered as an object name, not a bbox); the in-context task
 53# always answers with a bbox, so it is only ever in this in-context set.
 54INCONTEXT_GENERAL_BBOX_TYPES = frozenset(
 55    {
 56        "inctx_attr_single_color",
 57        "inctx_attr_single_material",
 58        "inctx_attr_single_function",
 59        "inctx_attr_composition",
 60        "inctx_attr_odd_one_out",
 61    }
 62)
 63INCONTEXT_SPATIAL_BBOX_TYPES = frozenset({"inctx_spatial_farthest"})
 64INCONTEXT_BBOX_TYPES = INCONTEXT_GENERAL_BBOX_TYPES | INCONTEXT_SPATIAL_BBOX_TYPES
 65BBOX_TYPES = GENERAL_BBOX_TYPES | SPATIAL_BBOX_TYPES | INCONTEXT_BBOX_TYPES
 66ATTRIBUTE_TYPES: frozenset[str] = frozenset()
 67TASK_TYPES = BINARY_TYPES | BBOX_TYPES | ATTRIBUTE_TYPES
 68
 69_REQUIRED_LOCATOR_KEYS = frozenset({"repo_id", "revision", "shard", "member"})
 70_FLOAT_PATH_SEGMENT_RE = re.compile(r"/(\d+)\.0(?=/)")
 71
 72
 73def normalize_shard(shard: str) -> str:
 74    """Fix a known upstream defect: some in-context ``target_image`` shard
 75    paths serialize an integer MOS phase as a float path segment (e.g.
 76    ``scenes/scene001/0.0/rgb.tar``), which does not exist on the Hub (the
 77    real path is ``scenes/scene001/0/rgb.tar``). This is purely a URL-
 78    construction normalization -- the stored locator dict itself is never
 79    mutated, so provenance/hashing over the raw manifest is unaffected."""
 80    return _FLOAT_PATH_SEGMENT_RE.sub(r"/\1", shard)
 81
 82
 83def stable_sample_id(row: dict[str, Any]) -> str:
 84    """Return an ID stable across file moves and parquet rewrites."""
 85    identity = "\x1f".join(
 86        str(row.get(key, ""))
 87        for key in ("scene_id", "kind", "phase", "frame", "type", "question")
 88    )
 89    return "vqa_" + hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20]
 90
 91
 92def image_locator(row: dict[str, Any]) -> dict[str, str]:
 93    """Build a portable pointer to one RGB member in an uncompressed Hub tar."""
 94    kind = str(row["kind"])
 95    if kind == "ego":
 96        partition = "ego"
 97    elif kind == "mos":
 98        partition = str(row["phase"])
 99    else:
100        raise ManifestError(f"unsupported VQA capture kind: {kind!r}", hint="use 'mos' or 'ego'")
101    scene = str(row["scene_id"])
102    frame = str(row["frame"])
103    return {
104        "repo_id": "IRVLUTD/RPX",
105        "revision": DATASET_REVISION,
106        "shard": f"scenes/{scene}/{partition}/rgb.tar",
107        "member": f"rgb/{frame}.webp",
108    }
109
110
111def _bbox_tuple(values: Any) -> tuple[int, int, int, int]:
112    if len(values) != 4:
113        raise ManifestError("Bounding boxes must contain exactly four coordinates")
114    return int(values[0]), int(values[1]), int(values[2]), int(values[3])
115
116
117@dataclass(frozen=True)
118class VQASample:
119    sample_id: str
120    scene_id: str
121    kind: str
122    phase: int | None
123    frame: str
124    image: dict[str, str]
125    img_w: int
126    img_h: int
127    question_type: str
128    question: str
129    answer: str
130    answer_bbox: tuple[int, int, int, int] | None = None
131    evidence: dict[str, Any] | None = None
132    schema_version: str = SCHEMA_VERSION
133    # In-context (two-image) fields. None on every normal row.
134    reference_image: dict[str, str] | None = None
135    reference_crop_bbox: tuple[int, int, int, int] | None = None
136    reference_crop_sha256: str | None = None
137    source_sample_id: str | None = None
138
139    @property
140    def is_in_context(self) -> bool:
141        return self.question_type in INCONTEXT_BBOX_TYPES
142
143    @property
144    def images(self) -> tuple[dict[str, str], ...]:
145        """Ordered locators to send to a model: [reference, target] for
146        in-context samples, [target] alone for normal samples."""
147        if self.is_in_context:
148            assert self.reference_image is not None
149            return (self.reference_image, self.image)
150        return (self.image,)
151
152    @classmethod
153    def from_dict(cls, row: dict[str, Any]) -> "VQASample":
154        bbox = row.get("answer_bbox")
155        evidence = row.get("evidence")
156        if isinstance(evidence, str):
157            evidence = json.loads(evidence)
158        question_type = str(row.get("question_type") or row["type"])
159        in_context = question_type in INCONTEXT_BBOX_TYPES
160        # In-context Parquet rows store their target/reference locators under
161        # target_image/reference_image; a round-tripped manifest row (written
162        # by write_manifest, reloaded by load_manifest) stores them under the
163        # canonical image/reference_image keys. Accept either spelling.
164        target_locator = row.get("image") or row.get("target_image")
165        if target_locator is None:
166            target_locator = image_locator(row)
167        target_locator = {key: str(target_locator[key]) for key in _REQUIRED_LOCATOR_KEYS}
168        reference_locator = row.get("reference_image")
169        reference_crop_bbox = row.get("reference_crop_bbox")
170        if reference_locator is not None:
171            if reference_crop_bbox is None and "crop_bbox" in reference_locator:
172                # The raw in-context Parquet's reference_image struct also
173                # carries a redundant crop_bbox key; prefer the top-level
174                # reference_crop_bbox column when present (it always is in
175                # the current release), falling back to the nested copy.
176                reference_crop_bbox = reference_locator["crop_bbox"]
177            reference_locator = {key: str(reference_locator[key]) for key in _REQUIRED_LOCATOR_KEYS}
178        raw_phase = row.get("phase")
179        is_null_phase = raw_phase is None or (isinstance(raw_phase, float) and raw_phase != raw_phase)
180        sample = cls(
181            sample_id=str(row.get("sample_id") or stable_sample_id(row)),
182            scene_id=str(row["scene_id"]),
183            kind=str(row["kind"]),
184            phase=None if raw_phase is None or is_null_phase else int(raw_phase),
185            frame=str(row["frame"]),
186            image=dict(target_locator),
187            img_w=int(row["img_w"]),
188            img_h=int(row["img_h"]),
189            question_type=question_type,
190            question=str(row["question"]),
191            answer=str(row["answer"]),
192            answer_bbox=None if bbox is None else _bbox_tuple(bbox),
193            evidence=evidence,
194            schema_version=str(row.get("schema_version", SCHEMA_VERSION)),
195            reference_image=dict(reference_locator) if reference_locator is not None else None,
196            reference_crop_bbox=(
197                None if reference_crop_bbox is None else _bbox_tuple(reference_crop_bbox)
198            ),
199            reference_crop_sha256=(
200                None
201                if row.get("reference_crop_sha256") is None
202                else str(row["reference_crop_sha256"])
203            ),
204            source_sample_id=(
205                None if row.get("source_sample_id") is None else str(row["source_sample_id"])
206            ),
207        )
208        if in_context and reference_locator is None:
209            raise ManifestError(
210                f"in-context VQA type {question_type!r} is missing reference_image",
211                hint="regenerate from the in-context Parquet, not the single-image one",
212            )
213        sample.validate()
214        return sample
215
216    def validate(self) -> None:
217        if self.schema_version not in {SCHEMA_VERSION, "incontext_v1"}:
218            raise ManifestError(
219                f"unsupported VQA schema_version: {self.schema_version}",
220                hint=f"regenerate with schema {SCHEMA_VERSION} or incontext_v1",
221            )
222        if self.question_type not in TASK_TYPES:
223            raise ManifestError(
224                f"out-of-scope VQA type: {self.question_type}", hint="filter to the frozen task set"
225            )
226        if self.kind not in {"mos", "ego"}:
227            raise ManifestError(f"invalid VQA kind: {self.kind}", hint="use 'mos' or 'ego'")
228        if self.kind == "mos" and self.phase not in {0, 1, 2}:
229            raise ManifestError("MOS sample has an invalid phase", hint="use phase 0, 1, or 2")
230        if self.kind == "ego" and self.phase is not None:
231            raise ManifestError("ego sample has a phase", hint="set phase to null for ego")
232        if self.question_type in BINARY_TYPES and self.answer not in {"yes", "no"}:
233            raise ManifestError("invalid binary answer", hint="use exactly 'yes' or 'no'")
234        if self.question_type in BBOX_TYPES and self.answer_bbox is None:
235            raise ManifestError("bbox task is missing answer_bbox", hint="regenerate its ground truth")
236        if self.answer_bbox is not None:
237            # Always checked against the TARGET image's own img_w/img_h --
238            # for in-context rows that is Image 2, never Image 1 (reference).
239            x0, y0, x1, y1 = self.answer_bbox
240            if not (0 <= x0 <= x1 < self.img_w and 0 <= y0 <= y1 < self.img_h):
241                raise ManifestError(
242                    f"invalid inclusive xyxy bbox: {self.answer_bbox}",
243                    hint="keep all coordinates ordered and inside the original image",
244                )
245        if self.is_in_context:
246            self._validate_in_context_locators()
247        else:
248            expected = image_locator(
249                {
250                    "scene_id": self.scene_id,
251                    "kind": self.kind,
252                    "phase": self.phase,
253                    "frame": self.frame,
254                }
255            )
256            if self.image != expected:
257                raise ManifestError(
258                    f"non-canonical VQA image locator: {self.image!r}",
259                    hint=f"use {expected!r}",
260                )
261
262    def _validate_in_context_locators(self) -> None:
263        # In-context locators are taken verbatim from the in-context Parquet
264        # (see module docstring) -- never recomputed via image_locator. What
265        # IS enforced: both locators are structurally complete and pinned to
266        # an immutable (non-"main"/non-branch-name) revision, and the
267        # reference crop has the bbox+hash needed to reconstruct Image 1.
268        if self.reference_image is None:
269            raise ManifestError(
270                "in-context sample is missing reference_image", hint="regenerate from the in-context Parquet"
271            )
272        for label, locator in (("target_image", self.image), ("reference_image", self.reference_image)):
273            missing = _REQUIRED_LOCATOR_KEYS - locator.keys()
274            if missing:
275                raise ManifestError(
276                    f"{label} locator is missing keys: {sorted(missing)}",
277                    hint=f"expected keys {sorted(_REQUIRED_LOCATOR_KEYS)}",
278                )
279            revision = str(locator["revision"])
280            if revision in {"main", "master"} or not re.fullmatch(r"[0-9a-f]{40}", revision):
281                raise ManifestError(
282                    f"{label} is not pinned to an immutable revision: {revision!r}",
283                    hint="use the exact revision stored in the in-context Parquet",
284                )
285        if self.reference_crop_bbox is None:
286            raise ManifestError(
287                "in-context sample is missing reference_crop_bbox",
288                hint="regenerate from the in-context Parquet",
289            )
290        if self.reference_crop_sha256 is None or not re.fullmatch(
291            r"[0-9a-f]{64}", self.reference_crop_sha256
292        ):
293            raise ManifestError(
294                f"invalid reference_crop_sha256: {self.reference_crop_sha256!r}",
295                hint="expected a 64-char lowercase hex sha256",
296            )
297
298    def to_dict(self) -> dict[str, Any]:
299        return {
300            "schema_version": self.schema_version,
301            "sample_id": self.sample_id,
302            "scene_id": self.scene_id,
303            "kind": self.kind,
304            "phase": self.phase,
305            "frame": self.frame,
306            "image": self.image,
307            "img_w": self.img_w,
308            "img_h": self.img_h,
309            "question_type": self.question_type,
310            "question": self.question,
311            "answer": self.answer,
312            "answer_bbox": list(self.answer_bbox) if self.answer_bbox is not None else None,
313            "bbox_format": "xyxy_pixel_inclusive" if self.answer_bbox is not None else None,
314            "evidence": self.evidence,
315            "reference_image": self.reference_image,
316            "reference_crop_bbox": (
317                list(self.reference_crop_bbox) if self.reference_crop_bbox is not None else None
318            ),
319            "reference_crop_sha256": self.reference_crop_sha256,
320            "source_sample_id": self.source_sample_id,
321        }
322
323
324def load_manifest(path: str | Path) -> list[VQASample]:
325    records: list[VQASample] = []
326    with Path(path).open(encoding="utf-8") as handle:
327        for line_number, line in enumerate(handle, 1):
328            if not line.strip():
329                continue
330            try:
331                records.append(VQASample.from_dict(json.loads(line)))
332            except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
333                raise ManifestError(
334                    f"invalid VQA manifest row at {path}:{line_number}: {exc}",
335                    hint="regenerate the manifest from the source parquets",
336                ) from exc
337    if not records:
338        raise ManifestError(f"empty VQA manifest: {path}", hint="provide at least one sample")
339    ids = [record.sample_id for record in records]
340    if len(ids) != len(set(ids)):
341        raise ManifestError(
342            "duplicate sample_id in VQA manifest", hint="deduplicate source questions"
343        )
344    return records
345
346
347def write_manifest(samples: Iterable[VQASample], path: str | Path) -> None:
348    target = Path(path)
349    target.parent.mkdir(parents=True, exist_ok=True)
350    with target.open("w", encoding="utf-8") as handle:
351        for sample in samples:
352            handle.write(json.dumps(sample.to_dict(), sort_keys=True, separators=(",", ":")) + "\n")