rpx_benchmark.schemas

  1"""Strict, typed manifest schemas for the RPX benchmark.
  2
  3This module defines Pydantic v2 models that mirror the JSON manifest
  4format consumed by :class:`rpx_benchmark.loader.RPXDataset`. It is an
  5**additive, opt-in** validation layer: the loader's hand-rolled checks
  6remain the default path and continue to work without Pydantic
  7installed. Import this module when you want:
  8
  9- Field-level validation errors (e.g. "samples[3].boxes: Field
 10  required") instead of the loader's coarser ``KeyError`` surface.
 11- JSON Schema export for third-party tooling or docs generation, via
 12  :func:`dump_json_schema`.
 13- Pre-upload validation in release automation that catches malformed
 14  shards before they reach the HuggingFace repo.
 15
 16Stability
 17---------
 18The model classes here are part of the public API at the
 19``rpx_benchmark.schemas`` surface. Field additions are non-breaking;
 20renames or removals require a major version bump. ``TaskType`` and
 21the enum values re-used here are defined in :mod:`rpx_benchmark.api`.
 22
 23Design
 24------
 25One :class:`Manifest` model holds the top-level fields (``task``,
 26``root``, ``samples``). Per-task sample schemas live in
 27:data:`SAMPLE_MODELS` and are dispatched on the manifest's
 28``task`` field by :func:`validate_manifest`. This keeps the model
 29class count small while still producing precise errors per task.
 30"""
 31
 32from __future__ import annotations
 33
 34from pathlib import Path
 35from typing import Any, Dict, List, Type, Union
 36
 37try:
 38    from pydantic import BaseModel, ConfigDict, Field, ValidationError
 39except ImportError as e:  # pragma: no cover — fail at import time with a clear hint
 40    raise ImportError(
 41        "rpx_benchmark.schemas requires pydantic>=2.5. "
 42        "Install with: pip install 'rpx-benchmark[schemas]'"
 43    ) from e
 44
 45from .api import Difficulty, Phase, TaskType
 46from .exceptions import ManifestError
 47
 48__all__ = [
 49    "BaseSampleEntry",
 50    "DepthSampleEntry",
 51    "DetectionSampleEntry",
 52    "SegmentationSampleEntry",
 53    "TrackingSampleEntry",
 54    "RelativePoseSampleEntry",
 55    "VisualGroundingSampleEntry",
 56    "SparseDepthSampleEntry",
 57    "KeypointMatchingSampleEntry",
 58    "Manifest",
 59    "SAMPLE_MODELS",
 60    "validate_manifest",
 61    "dump_json_schema",
 62]
 63
 64
 65# --------------------------------------------------------------------------- #
 66# Base configuration
 67# --------------------------------------------------------------------------- #
 68
 69
 70class _StrictBase(BaseModel):
 71    """Base for all manifest models.
 72
 73    - ``extra='allow'`` keeps forward-compatibility: unknown keys on a
 74      sample do not fail validation (they are preserved in metadata by
 75      the loader).
 76    - ``str_strip_whitespace`` normalises pasted paths.
 77    - ``frozen=False`` — we allow downstream tooling to mutate fields
 78      (e.g. rewriting relative paths to absolute).
 79    """
 80
 81    model_config = ConfigDict(
 82        extra="allow",
 83        str_strip_whitespace=True,
 84        populate_by_name=True,
 85    )
 86
 87
 88# --------------------------------------------------------------------------- #
 89# Sample schemas (one per task)
 90# --------------------------------------------------------------------------- #
 91
 92
 93class BaseSampleEntry(_StrictBase):
 94    """Fields common to every per-task sample entry."""
 95
 96    id: str = Field(..., description="Unique sample id, typically '{scene}_{phase}_{frame}'.")
 97    rgb: str = Field(..., description="Relative or absolute path to the RGB image.")
 98    scene: str | None = None
 99    phase: Phase | int | None = Field(
100        default=None,
101        description="Capture phase label ('clutter'/'interaction'/'clean') or 0/1/2.",
102    )
103    difficulty: Difficulty | None = None
104    pose: str | None = Field(
105        default=None,
106        description="Path to a T265 NPZ pose file (camera-to-world SE(3)).",
107    )
108    fisheye_left: str | None = None
109    fisheye_right: str | None = None
110    rgb_b: str | None = Field(
111        default=None,
112        description="Secondary RGB frame path (pair tasks: keypoint matching, pose).",
113    )
114    metadata: Dict[str, Any] | None = None
115
116
117class DepthSampleEntry(BaseSampleEntry):
118    depth: str = Field(..., description="Path to a 16-bit PNG depth map in millimetres.")
119
120
121class VideoDepthSampleEntry(BaseSampleEntry):
122    """Per-clip sample entry for the Video Depth task.
123
124    Differs from :class:`DepthSampleEntry` in that ``rgb`` and ``depth``
125    fields refer to entire frame sequences:
126
127    * ``rgb`` carries the *first* frame's path (kept for compatibility
128      with :class:`BaseSampleEntry`'s required field); the full clip
129      is in ``rgb_seq``.
130    * ``rgb_seq``, ``depth_seq`` are lists of relative-or-absolute
131      paths in temporal order.
132    * ``frame_indices`` records which original phase-relative indices
133      these paths correspond to (for adapters that downsample, this
134      lets the metric runner reconcile predictions to GT).
135    """
136
137    rgb_seq: List[str] = Field(
138        ...,
139        description="Ordered list of RGB frame paths for one (scene, phase) clip.",
140    )
141    depth_seq: List[str] = Field(
142        ...,
143        description="Ordered list of 16-bit PNG depth maps in millimetres.",
144    )
145    frame_indices: List[int] = Field(
146        ...,
147        description="Phase-relative integer indices for each frame in the clip.",
148    )
149
150
151class DetectionSampleEntry(BaseSampleEntry):
152    boxes: str = Field(..., description="Path to a JSON list of {bbox: [x1,y1,x2,y2], label: str}.")
153
154
155class SegmentationSampleEntry(BaseSampleEntry):
156    mask: str = Field(..., description="Path to a single-channel PNG with per-pixel instance IDs.")
157
158
159class TrackingSampleEntry(BaseSampleEntry):
160    tracks: str = Field(..., description="Path to a JSON list of tracklets.")
161
162
163class RelativePoseSampleEntry(BaseSampleEntry):
164    pose_a: str
165    pose_b: str
166    rgb_b: str = Field(..., description="Secondary RGB frame (required for pose pairs).")
167
168
169class VisualGroundingSampleEntry(BaseSampleEntry):
170    text: str
171    boxes: str | List[List[float]] | None = None
172    labels: List[str] | None = None
173
174
175class SparseDepthSampleEntry(BaseSampleEntry):
176    coordinates: str | List[List[float]]
177    depths: str | List[float]
178
179
180
181
182class KeypointMatchingSampleEntry(BaseSampleEntry):
183    points0: str | List[List[float]]
184    points1: str | List[List[float]]
185    visibility: str | List[bool] | None = None
186    rgb_b: str = Field(..., description="Secondary RGB frame (required for keypoint matching).")
187
188
189#: Dispatch table: ``TaskType → SampleEntry`` model class.
190SAMPLE_MODELS: Dict[TaskType, Type[BaseSampleEntry]] = {
191    TaskType.MONOCULAR_DEPTH: DepthSampleEntry,
192    TaskType.VIDEO_DEPTH: VideoDepthSampleEntry,
193    TaskType.OBJECT_DETECTION: DetectionSampleEntry,
194    TaskType.OPEN_VOCAB_DETECTION: DetectionSampleEntry,
195    TaskType.OBJECT_SEGMENTATION: SegmentationSampleEntry,
196    TaskType.OBJECT_TRACKING: TrackingSampleEntry,
197    TaskType.RELATIVE_CAMERA_POSE: RelativePoseSampleEntry,
198    TaskType.VISUAL_GROUNDING: VisualGroundingSampleEntry,
199    TaskType.SPARSE_DEPTH: SparseDepthSampleEntry,
200    TaskType.KEYPOINT_MATCHING: KeypointMatchingSampleEntry,
201}
202
203
204# --------------------------------------------------------------------------- #
205# Top-level manifest
206# --------------------------------------------------------------------------- #
207
208AnySampleEntry = Union[
209    DepthSampleEntry,
210    DetectionSampleEntry,
211    SegmentationSampleEntry,
212    TrackingSampleEntry,
213    RelativePoseSampleEntry,
214    VisualGroundingSampleEntry,
215    SparseDepthSampleEntry,
216    KeypointMatchingSampleEntry,
217]
218
219
220class Manifest(_StrictBase):
221    """Typed manifest describing a slice of the RPX dataset.
222
223    Attributes
224    ----------
225    task : TaskType
226        Task this manifest targets. Determines which
227        :class:`BaseSampleEntry` subclass each sample is validated
228        against.
229    root : str, optional
230        Root the sample paths resolve against. Absent manifests must
231        carry absolute paths in every sample.
232    samples : list of SampleEntry
233        Per-sample records. Parsed against the task-specific model in
234        :data:`SAMPLE_MODELS`.
235    scenes : list, optional
236        Alternate representation enumerating ``(scene, phase)`` pairs.
237        Preserved verbatim so hub download logic can read it without
238        round-tripping through per-sample records.
239    """
240
241    task: TaskType
242    root: str | None = None
243    samples: List[Dict[str, Any]] = Field(default_factory=list)
244    scenes: List[Dict[str, Any]] | None = None
245
246
247# --------------------------------------------------------------------------- #
248# Public helpers
249# --------------------------------------------------------------------------- #
250
251
252def validate_manifest(source: str | Path | Dict[str, Any]) -> Manifest:
253    """Validate ``source`` against the :class:`Manifest` schema.
254
255    Parameters
256    ----------
257    source : str, Path, or dict
258        Path to a manifest JSON, or an already-parsed dict.
259
260    Returns
261    -------
262    Manifest
263        The validated model.
264
265    Raises
266    ------
267    ManifestError
268        If the JSON is malformed or fails any schema validation. The
269        wrapped ``ValidationError`` message is surfaced in the
270        exception's ``details`` dict for programmatic access.
271    """
272    import json
273
274    if isinstance(source, (str, Path)):
275        path = Path(source)
276        if not path.is_file():
277            raise ManifestError(f"Manifest file not found: {path}")
278        try:
279            with path.open("r", encoding="utf-8") as f:
280                data = json.load(f)
281        except json.JSONDecodeError as e:
282            raise ManifestError(f"Manifest at {path} is not valid JSON: {e}") from e
283    else:
284        data = source
285
286    try:
287        manifest = Manifest.model_validate(data)
288    except ValidationError as e:
289        raise ManifestError(
290            "Manifest failed schema validation.",
291            hint="Run `rpx validate-manifest <path>` to see field-level errors.",
292            details={"pydantic_errors": e.errors()},
293        ) from e
294
295    # Per-sample validation against the task-specific schema. Kept
296    # outside the model_validator so per-sample errors surface as a
297    # single ManifestError with the full pydantic error list rather
298    # than being collapsed into one ValueError.
299    sample_model = SAMPLE_MODELS[manifest.task]
300    sample_errors: list[Dict[str, Any]] = []
301    rebuilt: List[Dict[str, Any]] = []
302    for idx, raw in enumerate(manifest.samples):
303        try:
304            rebuilt.append(sample_model.model_validate(raw).model_dump(exclude_none=True))
305        except ValidationError as e:
306            for err in e.errors():
307                # Prefix each error's location with the sample index so
308                # downstream tooling can point at the offending entry.
309                err = dict(err)
310                err["loc"] = ("samples", idx, *err.get("loc", ()))
311                sample_errors.append(err)
312    if sample_errors:
313        raise ManifestError(
314            f"{len(sample_errors)} sample(s) failed {sample_model.__name__} validation.",
315            hint="Each 'loc' entry points at the offending sample index and field.",
316            details={"pydantic_errors": sample_errors},
317        )
318    manifest.samples = rebuilt
319    return manifest
320
321
322def dump_json_schema(task: TaskType | None = None) -> Dict[str, Any]:
323    """Return the JSON Schema for a manifest (optionally task-scoped).
324
325    When ``task`` is ``None``, returns the generic
326    :class:`Manifest` schema. When a task is passed, returns the
327    per-task :class:`BaseSampleEntry` schema — useful for emitting
328    per-task documentation tables.
329
330    Examples
331    --------
332    >>> from rpx_benchmark.schemas import dump_json_schema
333    >>> from rpx_benchmark.api import TaskType
334    >>> schema = dump_json_schema(TaskType.MONOCULAR_DEPTH)
335    >>> "depth" in schema["properties"]
336    True
337    """
338    if task is None:
339        return Manifest.model_json_schema()
340    return SAMPLE_MODELS[task].model_json_schema()