rpx_benchmark.data.features

  1"""Canonical `datasets.Features` definitions for every RPX task.
  2
  3A shared schema per task means:
  4
  5- Shard-generation scripts write a known column layout.
  6- The :func:`row_to_sample` bridge can dispatch on task without
  7  second-guessing what columns exist.
  8- Third-party tooling (TypeScript SDK, parquet viewers, ...) can
  9  introspect expected columns from a single source of truth.
 10
 11The schemas intentionally keep **the raw modalities as HF native
 12``Image`` features** (PIL-decoded on read) rather than as numpy blobs:
 13``datasets`` streams these lazily from Parquet / Arrow, which is
 14orders of magnitude faster than dumping base64 arrays.
 15"""
 16
 17from __future__ import annotations
 18
 19from typing import Any, Dict
 20
 21try:
 22    from datasets import Features, Image, Sequence, Value
 23except ImportError as e:  # pragma: no cover — fail at import time
 24    raise ImportError(
 25        "rpx_benchmark.data.features requires the `datasets` library. "
 26        "Install with: pip install 'rpx-benchmark[hf-datasets]'"
 27    ) from e
 28
 29from ..api import TaskType
 30
 31# --------------------------------------------------------------------------- #
 32# Shared columns
 33# --------------------------------------------------------------------------- #
 34
 35
 36# Every RPX sample carries these top-level identity columns regardless
 37# of task. Grouping them here prevents drift across per-task schemas.
 38def _shared_columns() -> Dict[str, Any]:
 39    return {
 40        "id": Value("string"),
 41        "scene": Value("string"),
 42        "phase": Value("string"),  # "clutter" | "interaction" | "clean"
 43        "difficulty": Value("string"),  # "easy" | "medium" | "hard"
 44        "rgb": Image(decode=True),
 45        # 4×4 SE(3) camera-to-world, row-major float32. Flattened to a
 46        # fixed-length Sequence because datasets' Arrow backend stores
 47        # fixed-size arrays more efficiently than ragged ones.
 48        "camera_pose": Sequence(Value("float32"), length=16),
 49    }
 50
 51
 52# --------------------------------------------------------------------------- #
 53# Per-task Features
 54# --------------------------------------------------------------------------- #
 55
 56
 57def _depth_features() -> Features:
 58    return Features(
 59        {
 60            **_shared_columns(),
 61            # 16-bit PNG on the wire; PIL-decoded at read time so the
 62            # consumer can cast to float32 meters without re-encoding.
 63            "depth": Image(decode=True),
 64        }
 65    )
 66
 67
 68def _video_depth_features() -> Features:
 69    """Features for Video Depth.
 70
 71    Each row is a clip — not a frame — so the image columns become
 72    sequences. Tooling that consumes this schema must iterate the
 73    sequences pairwise (rgb_seq[t] aligns with depth_seq[t]).
 74    """
 75    return Features(
 76        {
 77            **_shared_columns(),
 78            # Per-clip arrays. Each list element is the same modality
 79            # in :func:`_depth_features` for one frame within the clip.
 80            "rgb_seq": Sequence(Image(decode=True)),
 81            "depth_seq": Sequence(Image(decode=True)),
 82            # Original phase-relative indices of the frames in this
 83            # clip (lets adapters that downsample report which frames
 84            # they kept). int32 on the wire.
 85            "frame_indices": Sequence(Value("int32")),
 86        }
 87    )
 88
 89
 90def _detection_features() -> Features:
 91    return Features(
 92        {
 93            **_shared_columns(),
 94            "boxes": Sequence(Sequence(Value("float32"), length=4)),
 95            "labels": Sequence(Value("string")),
 96        }
 97    )
 98
 99
100def _segmentation_features() -> Features:
101    return Features(
102        {
103            **_shared_columns(),
104            # Single-channel PNG where pixel values are instance IDs.
105            "mask": Image(decode=True),
106        }
107    )
108
109
110def _tracking_features() -> Features:
111    # Tracklets are variable-length; a struct-of-sequences layout keeps
112    # the Arrow schema cheap while preserving the tracklet grouping.
113    return Features(
114        {
115            **_shared_columns(),
116            "tracks": Sequence(
117                {
118                    "track_id": Value("string"),
119                    "boxes": Sequence(Sequence(Value("float32"), length=4)),
120                    "scores": Sequence(Value("float32")),
121                }
122            ),
123        }
124    )
125
126
127def _relative_pose_features() -> Features:
128    return Features(
129        {
130            **_shared_columns(),
131            "rgb_b": Image(decode=True),
132            "pose_a": Sequence(Value("float32"), length=16),
133            "pose_b": Sequence(Value("float32"), length=16),
134        }
135    )
136
137
138def _visual_grounding_features() -> Features:
139    return Features(
140        {
141            **_shared_columns(),
142            "text": Value("string"),
143            "boxes": Sequence(Sequence(Value("float32"), length=4)),
144            "labels": Sequence(Value("string")),
145        }
146    )
147
148
149def _sparse_depth_features() -> Features:
150    return Features(
151        {
152            **_shared_columns(),
153            "coordinates": Sequence(Sequence(Value("float32"), length=2)),
154            "depths": Sequence(Value("float32")),
155        }
156    )
157
158
159
160
161def _keypoint_features() -> Features:
162    return Features(
163        {
164            **_shared_columns(),
165            "rgb_b": Image(decode=True),
166            "points0": Sequence(Sequence(Value("float32"), length=2)),
167            "points1": Sequence(Sequence(Value("float32"), length=2)),
168            "visibility": Sequence(Value("bool")),
169        }
170    )
171
172
173#: Dispatch table: TaskType → HF `Features` definition.
174RPX_FEATURES: Dict[TaskType, Features] = {
175    TaskType.MONOCULAR_DEPTH: _depth_features(),
176    TaskType.VIDEO_DEPTH: _video_depth_features(),
177    TaskType.OBJECT_DETECTION: _detection_features(),
178    TaskType.OPEN_VOCAB_DETECTION: _detection_features(),
179    TaskType.OBJECT_SEGMENTATION: _segmentation_features(),
180    TaskType.OBJECT_TRACKING: _tracking_features(),
181    TaskType.RELATIVE_CAMERA_POSE: _relative_pose_features(),
182    TaskType.VISUAL_GROUNDING: _visual_grounding_features(),
183    TaskType.SPARSE_DEPTH: _sparse_depth_features(),
184    TaskType.KEYPOINT_MATCHING: _keypoint_features(),
185}
186
187
188def features_for_task(task: TaskType | str) -> Features:
189    """Return the canonical :class:`datasets.Features` for ``task``.
190
191    Parameters
192    ----------
193    task : TaskType or str
194
195    Returns
196    -------
197    datasets.Features
198    """
199    if isinstance(task, str):
200        task = TaskType(task)
201    return RPX_FEATURES[task]