rpx_benchmark.adapters
Formal adapter framework for bringing models into the RPX benchmark.
The contract is intentionally minimal: preprocess the RPX Sample into
whatever the model wants, invoke the model, postprocess the output into
the RPX prediction contract. Users composing their own adapter stack
only need to touch the model -- the input and output adapters shipped
here handle the plumbing for common model families.
See rpx_benchmark.adapters.base for the core types and
the repository script scripts/depth_models/hf_pipeline.py for a
Hugging Face depth-pipeline wrapper.
1"""Formal adapter framework for bringing models into the RPX benchmark. 2 3The contract is intentionally minimal: preprocess the RPX ``Sample`` into 4whatever the model wants, invoke the model, postprocess the output into 5the RPX prediction contract. Users composing their own adapter stack 6only need to touch the *model* -- the input and output adapters shipped 7here handle the plumbing for common model families. 8 9See :mod:`rpx_benchmark.adapters.base` for the core types and 10the repository script ``scripts/depth_models/hf_pipeline.py`` for a 11Hugging Face depth-pipeline wrapper. 12""" 13 14from .base import ( 15 BenchmarkableModel, 16 InputAdapter, 17 ModelInvoker, 18 OutputAdapter, 19 PreparedInput, 20 default_invoker, 21 make_numpy_depth_model, 22 make_numpy_detection_model, 23 make_numpy_grounding_model, 24 make_numpy_keypoint_model, 25 make_numpy_mask_model, 26 make_numpy_pose_model, 27 make_numpy_sparse_depth_model, 28 make_numpy_tracking_model, 29 make_numpy_video_depth_model, 30) 31from .batched_depth import BatchedDepthBenchmarkModel 32from .batched_multimodal import ( 33 BatchedRelativePoseBenchmarkModel, 34 BatchedSegmentationBenchmarkModel, 35 BatchedTaskBenchmarkModel, 36) 37 38__all__ = [ 39 "BatchedDepthBenchmarkModel", 40 "BatchedRelativePoseBenchmarkModel", 41 "BatchedSegmentationBenchmarkModel", 42 "BatchedTaskBenchmarkModel", 43 "BenchmarkableModel", 44 "InputAdapter", 45 "ModelInvoker", 46 "OutputAdapter", 47 "PreparedInput", 48 "default_invoker", 49 # Per-task numpy fast paths 50 "make_numpy_depth_model", 51 "make_numpy_detection_model", 52 "make_numpy_grounding_model", 53 "make_numpy_keypoint_model", 54 "make_numpy_mask_model", 55 "make_numpy_pose_model", 56 "make_numpy_sparse_depth_model", 57 "make_numpy_tracking_model", 58 "make_numpy_video_depth_model", 59]
78class BatchedDepthBenchmarkModel: 79 """Depth model wrapper that batches predictions across a sample list. 80 81 Parameters 82 ---------- 83 adapter 84 Callable: ``adapter(rgb_list_or_single) -> depth_list_or_single``. 85 Single-image: ``np.ndarray (H, W, 3) uint8`` → ``np.ndarray (H', W') float``. 86 Batched: ``list[np.ndarray]`` → ``list[np.ndarray]``. 87 name 88 Display name surfaced in result.json and Box paths. 89 save_dir 90 If provided, writes per-frame ``.npz`` predictions under 91 ``<save_dir>/<scene>/<phase>/<frame>.npz``. Created lazily on the 92 first save (no side effect at construction time). 93 native_alignment 94 Declared by the adapter (``"none"`` for metric models like ZoeDepth / 95 UniDepth V2 / Depth Pro; ``"ls_affine"`` for up-to-scale models 96 like Marigold / Lotus-2). The runner reads this to apply the 97 right alignment by default. Falls back to ``"none"`` if the 98 adapter doesn't declare it. The wrapper returns raw predictions; 99 the runner applies this mode once per complete scene/phase cell. 100 native_precision 101 Declared by the adapter (``"fp32"``, ``"fp16"``, ``"bf16"``). 102 Used by the runner to record the OperatingPoint precision 103 tag. Falls back to ``"fp32"`` if not declared. 104 allow_nonpositive_predictions 105 Permit finite, non-degenerate raw depth maps containing zero or 106 negative values. This is opt-in for protocols that preserve raw model 107 output and apply a documented evaluation-domain transform later. 108 """ 109 110 task = TaskType.MONOCULAR_DEPTH 111 112 def __init__( 113 self, 114 adapter: Any, 115 *, 116 name: str, 117 save_dir: Optional[str | Path] = None, 118 native_alignment: Optional[str] = None, 119 native_precision: Optional[str] = None, 120 resume_predictions: bool = False, 121 allow_nonpositive_predictions: bool = False, 122 ) -> None: 123 self._adapter = adapter 124 self.name = name 125 self._save_dir = Path(save_dir) if save_dir else None 126 if resume_predictions and self._save_dir is None: 127 from ..exceptions import ConfigError 128 129 raise ConfigError("resume_predictions requires save_dir") 130 self.resume_predictions = bool(resume_predictions) 131 self.allow_nonpositive_predictions = bool(allow_nonpositive_predictions) 132 self.resume_stats = { 133 "cache_hits": 0, 134 "inferred_new": 0, 135 "invalid_recomputed": 0, 136 } 137 self.last_cache_hits: list[bool] = [] 138 self._last_save_required: list[bool] = [] 139 # Read alignment off the adapter if it declared one, else fall 140 # back to the constructor arg, else "none" (metric assumption). 141 self.native_alignment: str = ( 142 native_alignment or getattr(adapter, "native_alignment", None) or "none" 143 ) 144 self.native_precision: str = ( 145 native_precision or getattr(adapter, "native_precision", None) or "fp32" 146 ) 147 self.depth_output_kind: str = "relative" if self.native_alignment != "none" else "metric" 148 # Profiler walker reaches the underlying nn.Module via this attr. 149 self.model = adapter 150 151 # No-op: the adapter loads its own weights at construction. 152 def setup(self) -> None: 153 return None 154 155 def predict(self, batch: Sequence[Sample]) -> List[DepthPrediction]: 156 cached: dict[int, np.ndarray] = {} 157 missing_indices: list[int] = [] 158 invalid_existing: set[int] = set() 159 for index, sample in enumerate(batch): 160 path = self._prediction_path(sample) 161 if self.resume_predictions and path is not None and path.exists(): 162 depth = self._load_valid_prediction( 163 path, 164 np.asarray(sample.rgb).shape[:2], 165 require_positive=not self.allow_nonpositive_predictions, 166 ) 167 if depth is not None: 168 cached[index] = depth 169 self.resume_stats["cache_hits"] += 1 170 continue 171 invalid_existing.add(index) 172 missing_indices.append(index) 173 174 inferred: list[np.ndarray] = [] 175 if missing_indices: 176 rgbs = [np.asarray(batch[i].rgb, dtype=np.uint8) for i in missing_indices] 177 raw_depths = self._adapter(rgbs) 178 if not isinstance(raw_depths, (list, tuple)): 179 raw_depths = [raw_depths] 180 inferred = [np.asarray(depth, dtype=np.float32) for depth in raw_depths] 181 if len(inferred) != len(missing_indices): 182 from ..exceptions import AdapterError 183 184 raise AdapterError( 185 f"adapter returned {len(inferred)} depths for {len(missing_indices)} inputs", 186 hint="Check the adapter's batched contract: it must return " 187 "one depth map per input RGB.", 188 ) 189 190 inferred_by_index = dict(zip(missing_indices, inferred, strict=True)) 191 preds: List[DepthPrediction] = [] 192 self.last_cache_hits = [] 193 self._last_save_required = [] 194 for index, sample in enumerate(batch): 195 is_cached = index in cached 196 d_raw = cached[index] if is_cached else inferred_by_index[index] 197 target_hw = np.asarray(sample.rgb).shape[:2] 198 if d_raw.shape != target_hw: 199 d_raw = _resize_bilinear_2d(d_raw, target_hw) 200 if not is_cached: 201 if index in invalid_existing: 202 self.resume_stats["invalid_recomputed"] += 1 203 else: 204 self.resume_stats["inferred_new"] += 1 205 206 # Return the raw prediction. Relative outputs are deliberately 207 # aligned later by BenchmarkRunner, after every frame in the 208 # same (scene, phase) cell is available for one pooled solve. 209 preds.append(DepthPrediction(depth_map=d_raw)) 210 self.last_cache_hits.append(is_cached) 211 self._last_save_required.append(not is_cached) 212 return preds 213 214 def persist_predictions( 215 self, 216 batch: Sequence[Sample], 217 predictions: Sequence[DepthPrediction], 218 ) -> None: 219 """Persist newly inferred predictions after the runner stops timing.""" 220 pending = [ 221 (sample, prediction.depth_map) 222 for sample, prediction, required in zip( 223 batch, 224 predictions, 225 self._last_save_required, 226 strict=True, 227 ) 228 if required 229 ] 230 if len(pending) <= 1: 231 for sample, depth in pending: 232 self._maybe_save(sample, depth) 233 return 234 235 # Compression is CPU-bound and each frame has a distinct atomic target. 236 # Parallelize it within the batch so a fast GPU is not followed by a 237 # serial chain of np.savez_compressed calls. 238 workers = min(len(pending), 8) 239 with ThreadPoolExecutor(max_workers=workers) as executor: 240 futures = [ 241 executor.submit(self._maybe_save, sample, depth) 242 for sample, depth in pending 243 ] 244 for future in futures: 245 future.result() 246 247 def _prediction_path(self, sample: Sample) -> Path | None: 248 if self._save_dir is None: 249 return None 250 # Prefer the manifest's metadata block (set by local_manifest.py 251 # and dataset_hub.split_manifests so the loader rides it through 252 # to Sample.metadata). Fall back to id-parsing. 253 meta = getattr(sample, "metadata", None) or {} 254 scene = meta.get("scene_id") 255 phase = meta.get("phase_idx") 256 frame = meta.get("frame") 257 if not (scene and frame is not None and phase is not None): 258 try: 259 scene, phase, frame = str(sample.id).split("__", 2) 260 except ValueError: 261 scene, phase, frame = "unknown", "0", str(sample.id) 262 return self._save_dir / str(scene) / str(phase) / f"{frame}.npz" 263 264 @staticmethod 265 def _load_valid_prediction( 266 path: Path, 267 target_hw: tuple[int, int], 268 *, 269 require_positive: bool = True, 270 ) -> np.ndarray | None: 271 try: 272 with np.load(path, allow_pickle=False) as payload: 273 if payload.files != ["depth"]: 274 return None 275 depth = np.asarray(payload["depth"]) 276 return BatchedDepthBenchmarkModel._validated_depth( 277 depth, 278 target_hw, 279 require_positive=require_positive, 280 ) 281 except (OSError, ValueError, KeyError, EOFError, zipfile.BadZipFile): 282 return None 283 284 @staticmethod 285 def _validated_depth( 286 depth: np.ndarray, 287 target_hw: tuple[int, int], 288 *, 289 require_positive: bool = True, 290 ) -> np.ndarray | None: 291 depth = np.asarray(depth) 292 if depth.ndim != 2 or depth.shape != target_hw: 293 return None 294 if depth.dtype != np.dtype(np.float32): 295 return None 296 if not np.isfinite(depth).all(): 297 return None 298 if require_positive and np.any(depth <= 0): 299 return None 300 if float(np.ptp(depth)) <= 1e-6: 301 return None 302 return depth 303 304 def _maybe_save(self, sample: Sample, depth: np.ndarray) -> None: 305 out = self._prediction_path(sample) 306 if out is None: 307 return 308 canonical = np.asarray(depth, dtype=np.float32) 309 target_hw = np.asarray(sample.rgb).shape[:2] 310 if ( 311 self._validated_depth( 312 canonical, 313 target_hw, 314 require_positive=not self.allow_nonpositive_predictions, 315 ) 316 is None 317 ): 318 from ..exceptions import AdapterError 319 320 raise AdapterError( 321 f"Refusing to save an invalid depth prediction for sample {sample.id!r}: " 322 f"expected a finite, non-degenerate float32 map with shape {target_hw}" 323 + ( 324 " and strictly positive values." 325 if not self.allow_nonpositive_predictions 326 else "." 327 ) 328 ) 329 out.parent.mkdir(parents=True, exist_ok=True) 330 temp_path: Path | None = None 331 try: 332 with tempfile.NamedTemporaryFile( 333 mode="w+b", 334 suffix=".npz.part", 335 dir=out.parent, 336 delete=False, 337 ) as handle: 338 temp_path = Path(handle.name) 339 np.savez_compressed(handle, depth=canonical) 340 handle.flush() 341 os.fsync(handle.fileno()) 342 os.replace(temp_path, out) 343 finally: 344 if temp_path is not None and temp_path.exists(): 345 temp_path.unlink()
Depth model wrapper that batches predictions across a sample list.
Parameters
adapter
Callable: adapter(rgb_list_or_single) -> depth_list_or_single.
Single-image: np.ndarray (H, W, 3) uint8 → np.ndarray (H', W') float.
Batched: list[np.ndarray] → list[np.ndarray].
name
Display name surfaced in result.json and Box paths.
save_dir
If provided, writes per-frame .npz predictions under
<save_dir>/<scene>/<phase>/<frame>.npz. Created lazily on the
first save (no side effect at construction time).
native_alignment
Declared by the adapter ("none" for metric models like ZoeDepth /
UniDepth V2 / Depth Pro; "ls_affine" for up-to-scale models
like Marigold / Lotus-2). The runner reads this to apply the
right alignment by default. Falls back to "none" if the
adapter doesn't declare it. The wrapper returns raw predictions;
the runner applies this mode once per complete scene/phase cell.
native_precision
Declared by the adapter ("fp32", "fp16", "bf16").
Used by the runner to record the OperatingPoint precision
tag. Falls back to "fp32" if not declared.
allow_nonpositive_predictions
Permit finite, non-degenerate raw depth maps containing zero or
negative values. This is opt-in for protocols that preserve raw model
output and apply a documented evaluation-domain transform later.
112 def __init__( 113 self, 114 adapter: Any, 115 *, 116 name: str, 117 save_dir: Optional[str | Path] = None, 118 native_alignment: Optional[str] = None, 119 native_precision: Optional[str] = None, 120 resume_predictions: bool = False, 121 allow_nonpositive_predictions: bool = False, 122 ) -> None: 123 self._adapter = adapter 124 self.name = name 125 self._save_dir = Path(save_dir) if save_dir else None 126 if resume_predictions and self._save_dir is None: 127 from ..exceptions import ConfigError 128 129 raise ConfigError("resume_predictions requires save_dir") 130 self.resume_predictions = bool(resume_predictions) 131 self.allow_nonpositive_predictions = bool(allow_nonpositive_predictions) 132 self.resume_stats = { 133 "cache_hits": 0, 134 "inferred_new": 0, 135 "invalid_recomputed": 0, 136 } 137 self.last_cache_hits: list[bool] = [] 138 self._last_save_required: list[bool] = [] 139 # Read alignment off the adapter if it declared one, else fall 140 # back to the constructor arg, else "none" (metric assumption). 141 self.native_alignment: str = ( 142 native_alignment or getattr(adapter, "native_alignment", None) or "none" 143 ) 144 self.native_precision: str = ( 145 native_precision or getattr(adapter, "native_precision", None) or "fp32" 146 ) 147 self.depth_output_kind: str = "relative" if self.native_alignment != "none" else "metric" 148 # Profiler walker reaches the underlying nn.Module via this attr. 149 self.model = adapter
155 def predict(self, batch: Sequence[Sample]) -> List[DepthPrediction]: 156 cached: dict[int, np.ndarray] = {} 157 missing_indices: list[int] = [] 158 invalid_existing: set[int] = set() 159 for index, sample in enumerate(batch): 160 path = self._prediction_path(sample) 161 if self.resume_predictions and path is not None and path.exists(): 162 depth = self._load_valid_prediction( 163 path, 164 np.asarray(sample.rgb).shape[:2], 165 require_positive=not self.allow_nonpositive_predictions, 166 ) 167 if depth is not None: 168 cached[index] = depth 169 self.resume_stats["cache_hits"] += 1 170 continue 171 invalid_existing.add(index) 172 missing_indices.append(index) 173 174 inferred: list[np.ndarray] = [] 175 if missing_indices: 176 rgbs = [np.asarray(batch[i].rgb, dtype=np.uint8) for i in missing_indices] 177 raw_depths = self._adapter(rgbs) 178 if not isinstance(raw_depths, (list, tuple)): 179 raw_depths = [raw_depths] 180 inferred = [np.asarray(depth, dtype=np.float32) for depth in raw_depths] 181 if len(inferred) != len(missing_indices): 182 from ..exceptions import AdapterError 183 184 raise AdapterError( 185 f"adapter returned {len(inferred)} depths for {len(missing_indices)} inputs", 186 hint="Check the adapter's batched contract: it must return " 187 "one depth map per input RGB.", 188 ) 189 190 inferred_by_index = dict(zip(missing_indices, inferred, strict=True)) 191 preds: List[DepthPrediction] = [] 192 self.last_cache_hits = [] 193 self._last_save_required = [] 194 for index, sample in enumerate(batch): 195 is_cached = index in cached 196 d_raw = cached[index] if is_cached else inferred_by_index[index] 197 target_hw = np.asarray(sample.rgb).shape[:2] 198 if d_raw.shape != target_hw: 199 d_raw = _resize_bilinear_2d(d_raw, target_hw) 200 if not is_cached: 201 if index in invalid_existing: 202 self.resume_stats["invalid_recomputed"] += 1 203 else: 204 self.resume_stats["inferred_new"] += 1 205 206 # Return the raw prediction. Relative outputs are deliberately 207 # aligned later by BenchmarkRunner, after every frame in the 208 # same (scene, phase) cell is available for one pooled solve. 209 preds.append(DepthPrediction(depth_map=d_raw)) 210 self.last_cache_hits.append(is_cached) 211 self._last_save_required.append(not is_cached) 212 return preds
214 def persist_predictions( 215 self, 216 batch: Sequence[Sample], 217 predictions: Sequence[DepthPrediction], 218 ) -> None: 219 """Persist newly inferred predictions after the runner stops timing.""" 220 pending = [ 221 (sample, prediction.depth_map) 222 for sample, prediction, required in zip( 223 batch, 224 predictions, 225 self._last_save_required, 226 strict=True, 227 ) 228 if required 229 ] 230 if len(pending) <= 1: 231 for sample, depth in pending: 232 self._maybe_save(sample, depth) 233 return 234 235 # Compression is CPU-bound and each frame has a distinct atomic target. 236 # Parallelize it within the batch so a fast GPU is not followed by a 237 # serial chain of np.savez_compressed calls. 238 workers = min(len(pending), 8) 239 with ThreadPoolExecutor(max_workers=workers) as executor: 240 futures = [ 241 executor.submit(self._maybe_save, sample, depth) 242 for sample, depth in pending 243 ] 244 for future in futures: 245 future.result()
Persist newly inferred predictions after the runner stops timing.
216class BatchedRelativePoseBenchmarkModel(BatchedTaskBenchmarkModel): 217 """Batched wrapper for paired-frame relative-pose models. 218 219 Adapter contract: ``adapter(list_of_dicts) -> list_of_pose_dicts``, 220 where each input dict is ``{"rgb_a": HxWx3 uint8, "rgb_b": HxWx3 uint8}`` 221 and each output is ``{"rotation": 3x3, "translation": 3}`` (or a tuple). 222 """ 223 224 task = TaskType.RELATIVE_CAMERA_POSE 225 226 def extract_inputs(self, sample: Sample) -> dict: 227 meta = getattr(sample, "metadata", None) or {} 228 rgb_a = np.asarray(sample.rgb, dtype=np.uint8) 229 rgb_b = meta.get("rgb_b") 230 if rgb_b is None: 231 from ..exceptions import AdapterError 232 233 raise AdapterError( 234 "Relative-pose adapter requires `rgb_b` in sample.metadata; " 235 "check the manifest writes `rgb_b` so the loader stashes it.", 236 ) 237 return {"rgb_a": rgb_a, "rgb_b": np.asarray(rgb_b, dtype=np.uint8)} 238 239 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 240 from ..api import RelativePosePrediction 241 242 if isinstance(model_output, dict): 243 rot = model_output["rotation"] 244 trans = model_output["translation"] 245 else: 246 rot, trans = model_output 247 return RelativePosePrediction( 248 rotation=np.asarray(rot, dtype=np.float64), 249 translation=np.asarray(trans, dtype=np.float64), 250 ) 251 252 #: CSV columns. ``sample_id`` is the canonical identity; both phases are 253 #: retained so cross-phase pairs cannot collide with intra-phase pairs. 254 _LOG_COLUMNS: tuple[str, ...] = ( 255 "sample_id", 256 "scene_id", 257 "phase", 258 "phase_b", 259 "frame_a", 260 "frame_b", 261 "R00", 262 "R01", 263 "R02", 264 "R10", 265 "R11", 266 "R12", 267 "R20", 268 "R21", 269 "R22", 270 "tx", 271 "ty", 272 "tz", 273 ) 274 275 def maybe_save(self, sample: Sample, model_output: Any) -> None: 276 """Append one row per pair to ``<save_dir>/predictions.csv``. 277 278 Rationale: paired-pose runs produce ~3 K rows per (model, split); 279 a single CSV is friendlier for Box and analytics than 3 K 280 per-pair ``.npz`` files. Columns: scene_id, phase, frame_a, 281 frame_b, the 9 rotation entries (row-major) and 3 translation 282 components. Header is written exactly once per file (on the 283 first append to an empty / non-existent log). Idempotent on 284 ``sample.id`` within a single process so 285 the runner's deployment-readiness double pass (warmup + measure) 286 doesn't duplicate rows. 287 """ 288 if self._save_dir is None: 289 return 290 import csv 291 292 meta = getattr(sample, "metadata", None) or {} 293 scene = meta.get("scene_id") or "unknown" 294 phase = meta.get("phase_idx") if meta.get("phase_idx") is not None else "0" 295 phase_b = meta.get("phase_idx_b") 296 if phase_b is None: 297 phase_b = phase 298 frame_a = meta.get("frame") or str(sample.id) 299 frame_b = meta.get("frame_b") or "?" 300 sample_id = str(sample.id) 301 302 seen = self.__dict__.setdefault("_pose_csv_seen", set()) 303 key = sample_id 304 if key in seen: 305 return 306 seen.add(key) 307 308 if isinstance(model_output, dict): 309 rot = np.asarray(model_output["rotation"], dtype=np.float64) 310 trans = np.asarray(model_output["translation"], dtype=np.float64) 311 else: 312 rot, trans = model_output 313 rot = np.asarray(rot, dtype=np.float64) 314 trans = np.asarray(trans, dtype=np.float64) 315 316 # Tolerate rotation given as a flat 9-vector or a quaternion. 317 rot = rot.reshape(-1) 318 if rot.size == 4: 319 # Quaternion (x, y, z, w) → 3×3 rotation. We don't depend on 320 # scipy here; build it inline from the unit-quaternion formula. 321 x, y, z, w = rot.tolist() 322 rot = np.asarray( 323 [ 324 1 - 2 * (y * y + z * z), 325 2 * (x * y - z * w), 326 2 * (x * z + y * w), 327 2 * (x * y + z * w), 328 1 - 2 * (x * x + z * z), 329 2 * (y * z - x * w), 330 2 * (x * z - y * w), 331 2 * (y * z + x * w), 332 1 - 2 * (x * x + y * y), 333 ], 334 dtype=np.float64, 335 ) 336 elif rot.size != 9: 337 from ..exceptions import AdapterError 338 339 raise AdapterError( 340 f"unexpected rotation shape: got {rot.size} entries; " 341 "expected 9 (3×3 matrix) or 4 (xyzw quaternion).", 342 ) 343 trans = trans.reshape(-1) 344 if trans.size != 3: 345 from ..exceptions import AdapterError 346 347 raise AdapterError( 348 f"unexpected translation shape: got {trans.size} entries; expected 3.", 349 ) 350 351 log_path = self._save_dir / "predictions.csv" 352 log_path.parent.mkdir(parents=True, exist_ok=True) 353 write_header = not log_path.exists() or log_path.stat().st_size == 0 354 with log_path.open("a", newline="") as f: 355 writer = csv.writer(f) 356 if write_header: 357 writer.writerow(self._LOG_COLUMNS) 358 writer.writerow( 359 [ 360 sample_id, 361 scene, 362 phase, 363 phase_b, 364 frame_a, 365 frame_b, 366 *(f"{v:.10g}" for v in rot.tolist()), 367 *(f"{v:.10g}" for v in trans.tolist()), 368 ] 369 )
Batched wrapper for paired-frame relative-pose models.
Adapter contract: adapter(list_of_dicts) -> list_of_pose_dicts,
where each input dict is {"rgb_a": HxWx3 uint8, "rgb_b": HxWx3 uint8}
and each output is {"rotation": 3x3, "translation": 3} (or a tuple).
226 def extract_inputs(self, sample: Sample) -> dict: 227 meta = getattr(sample, "metadata", None) or {} 228 rgb_a = np.asarray(sample.rgb, dtype=np.uint8) 229 rgb_b = meta.get("rgb_b") 230 if rgb_b is None: 231 from ..exceptions import AdapterError 232 233 raise AdapterError( 234 "Relative-pose adapter requires `rgb_b` in sample.metadata; " 235 "check the manifest writes `rgb_b` so the loader stashes it.", 236 ) 237 return {"rgb_a": rgb_a, "rgb_b": np.asarray(rgb_b, dtype=np.uint8)}
Subclass: return whatever the adapter expects per sample.
Examples:
- Depth:
np.asarray(sample.rgb, dtype=np.uint8) - Pose:
{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]} - Stereo depth:
{"left": metadata["fisheye_left"], "right": metadata["fisheye_right"]}
239 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 240 from ..api import RelativePosePrediction 241 242 if isinstance(model_output, dict): 243 rot = model_output["rotation"] 244 trans = model_output["translation"] 245 else: 246 rot, trans = model_output 247 return RelativePosePrediction( 248 rotation=np.asarray(rot, dtype=np.float64), 249 translation=np.asarray(trans, dtype=np.float64), 250 )
Subclass: wrap one model output as the task's Prediction dataclass.
275 def maybe_save(self, sample: Sample, model_output: Any) -> None: 276 """Append one row per pair to ``<save_dir>/predictions.csv``. 277 278 Rationale: paired-pose runs produce ~3 K rows per (model, split); 279 a single CSV is friendlier for Box and analytics than 3 K 280 per-pair ``.npz`` files. Columns: scene_id, phase, frame_a, 281 frame_b, the 9 rotation entries (row-major) and 3 translation 282 components. Header is written exactly once per file (on the 283 first append to an empty / non-existent log). Idempotent on 284 ``sample.id`` within a single process so 285 the runner's deployment-readiness double pass (warmup + measure) 286 doesn't duplicate rows. 287 """ 288 if self._save_dir is None: 289 return 290 import csv 291 292 meta = getattr(sample, "metadata", None) or {} 293 scene = meta.get("scene_id") or "unknown" 294 phase = meta.get("phase_idx") if meta.get("phase_idx") is not None else "0" 295 phase_b = meta.get("phase_idx_b") 296 if phase_b is None: 297 phase_b = phase 298 frame_a = meta.get("frame") or str(sample.id) 299 frame_b = meta.get("frame_b") or "?" 300 sample_id = str(sample.id) 301 302 seen = self.__dict__.setdefault("_pose_csv_seen", set()) 303 key = sample_id 304 if key in seen: 305 return 306 seen.add(key) 307 308 if isinstance(model_output, dict): 309 rot = np.asarray(model_output["rotation"], dtype=np.float64) 310 trans = np.asarray(model_output["translation"], dtype=np.float64) 311 else: 312 rot, trans = model_output 313 rot = np.asarray(rot, dtype=np.float64) 314 trans = np.asarray(trans, dtype=np.float64) 315 316 # Tolerate rotation given as a flat 9-vector or a quaternion. 317 rot = rot.reshape(-1) 318 if rot.size == 4: 319 # Quaternion (x, y, z, w) → 3×3 rotation. We don't depend on 320 # scipy here; build it inline from the unit-quaternion formula. 321 x, y, z, w = rot.tolist() 322 rot = np.asarray( 323 [ 324 1 - 2 * (y * y + z * z), 325 2 * (x * y - z * w), 326 2 * (x * z + y * w), 327 2 * (x * y + z * w), 328 1 - 2 * (x * x + z * z), 329 2 * (y * z - x * w), 330 2 * (x * z - y * w), 331 2 * (y * z + x * w), 332 1 - 2 * (x * x + y * y), 333 ], 334 dtype=np.float64, 335 ) 336 elif rot.size != 9: 337 from ..exceptions import AdapterError 338 339 raise AdapterError( 340 f"unexpected rotation shape: got {rot.size} entries; " 341 "expected 9 (3×3 matrix) or 4 (xyzw quaternion).", 342 ) 343 trans = trans.reshape(-1) 344 if trans.size != 3: 345 from ..exceptions import AdapterError 346 347 raise AdapterError( 348 f"unexpected translation shape: got {trans.size} entries; expected 3.", 349 ) 350 351 log_path = self._save_dir / "predictions.csv" 352 log_path.parent.mkdir(parents=True, exist_ok=True) 353 write_header = not log_path.exists() or log_path.stat().st_size == 0 354 with log_path.open("a", newline="") as f: 355 writer = csv.writer(f) 356 if write_header: 357 writer.writerow(self._LOG_COLUMNS) 358 writer.writerow( 359 [ 360 sample_id, 361 scene, 362 phase, 363 phase_b, 364 frame_a, 365 frame_b, 366 *(f"{v:.10g}" for v in rot.tolist()), 367 *(f"{v:.10g}" for v in trans.tolist()), 368 ] 369 )
Append one row per pair to <save_dir>/predictions.csv.
Rationale: paired-pose runs produce ~3 K rows per (model, split);
a single CSV is friendlier for Box and analytics than 3 K
per-pair .npz files. Columns: scene_id, phase, frame_a,
frame_b, the 9 rotation entries (row-major) and 3 translation
components. Header is written exactly once per file (on the
first append to an empty / non-existent log). Idempotent on
sample.id within a single process so
the runner's deployment-readiness double pass (warmup + measure)
doesn't duplicate rows.
Inherited Members
168class BatchedSegmentationBenchmarkModel(BatchedTaskBenchmarkModel): 169 """Batched wrapper for object-segmentation models. 170 171 Adapter contract: ``adapter(list_of_rgb_uint8_HxWx3) -> list_of_int_mask_HxW``. 172 Output mask values are integer instance IDs (0 = background). If the 173 model returns a different shape than the input RGB, we nearest-neighbour 174 resize back to preserve integer IDs. 175 """ 176 177 task = TaskType.OBJECT_SEGMENTATION 178 179 def extract_inputs(self, sample: Sample) -> np.ndarray: 180 return np.asarray(sample.rgb, dtype=np.uint8) 181 182 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 183 from PIL import Image 184 185 from ..api import SegmentationPrediction 186 187 mask = np.asarray(model_output) 188 target_hw = np.asarray(sample.rgb).shape[:2] 189 if mask.shape != target_hw: 190 # Nearest-neighbour preserves integer IDs. 191 img = Image.fromarray(mask.astype(np.int32), mode="I") 192 img = img.resize((target_hw[1], target_hw[0]), Image.NEAREST) 193 mask = np.asarray(img, dtype=np.int32) 194 return SegmentationPrediction(mask=mask.astype(np.int32)) 195 196 def maybe_save(self, sample: Sample, model_output: Any) -> None: 197 if self._save_dir is None: 198 return 199 from PIL import Image 200 201 meta = getattr(sample, "metadata", None) or {} 202 scene = meta.get("scene_id") 203 phase = meta.get("phase_idx") 204 frame = meta.get("frame") 205 if not (scene and frame is not None and phase is not None): 206 try: 207 scene, phase, frame = str(sample.id).split("__", 2) 208 except ValueError: 209 scene, phase, frame = "unknown", "0", str(sample.id) 210 out = self._save_dir / str(scene) / str(phase) / f"{frame}.png" 211 out.parent.mkdir(parents=True, exist_ok=True) 212 # 16-bit PNG preserves up to 65 K instance ids; sufficient for SAM2. 213 Image.fromarray(np.asarray(model_output, dtype=np.uint16), mode="I;16").save(out)
Batched wrapper for object-segmentation models.
Adapter contract: adapter(list_of_rgb_uint8_HxWx3) -> list_of_int_mask_HxW.
Output mask values are integer instance IDs (0 = background). If the
model returns a different shape than the input RGB, we nearest-neighbour
resize back to preserve integer IDs.
179 def extract_inputs(self, sample: Sample) -> np.ndarray: 180 return np.asarray(sample.rgb, dtype=np.uint8)
Subclass: return whatever the adapter expects per sample.
Examples:
- Depth:
np.asarray(sample.rgb, dtype=np.uint8) - Pose:
{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]} - Stereo depth:
{"left": metadata["fisheye_left"], "right": metadata["fisheye_right"]}
182 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 183 from PIL import Image 184 185 from ..api import SegmentationPrediction 186 187 mask = np.asarray(model_output) 188 target_hw = np.asarray(sample.rgb).shape[:2] 189 if mask.shape != target_hw: 190 # Nearest-neighbour preserves integer IDs. 191 img = Image.fromarray(mask.astype(np.int32), mode="I") 192 img = img.resize((target_hw[1], target_hw[0]), Image.NEAREST) 193 mask = np.asarray(img, dtype=np.int32) 194 return SegmentationPrediction(mask=mask.astype(np.int32))
Subclass: wrap one model output as the task's Prediction dataclass.
196 def maybe_save(self, sample: Sample, model_output: Any) -> None: 197 if self._save_dir is None: 198 return 199 from PIL import Image 200 201 meta = getattr(sample, "metadata", None) or {} 202 scene = meta.get("scene_id") 203 phase = meta.get("phase_idx") 204 frame = meta.get("frame") 205 if not (scene and frame is not None and phase is not None): 206 try: 207 scene, phase, frame = str(sample.id).split("__", 2) 208 except ValueError: 209 scene, phase, frame = "unknown", "0", str(sample.id) 210 out = self._save_dir / str(scene) / str(phase) / f"{frame}.png" 211 out.parent.mkdir(parents=True, exist_ok=True) 212 # 16-bit PNG preserves up to 65 K instance ids; sufficient for SAM2. 213 Image.fromarray(np.asarray(model_output, dtype=np.uint16), mode="I;16").save(out)
Optional override — persist per-frame artefacts for analytics.
Inherited Members
73class BatchedTaskBenchmarkModel: 74 """Abstract base — true-batched-dispatch wrapper for any RPX task. 75 76 Subclasses set the class-level ``task`` (a :class:`TaskType`) and 77 override ``extract_inputs`` + ``wrap_output``. The harness contract 78 matches ``BenchmarkableModel`` (``setup`` + ``predict(batch)``). 79 80 Parameters 81 ---------- 82 adapter 83 The model callable. Receives a list of whatever ``extract_inputs`` 84 returns (one per sample). Returns a list of model outputs (one per 85 sample, same order). 86 name 87 Display name surfaced in result.json + Box paths. 88 save_dir 89 If set, ``maybe_save`` writes per-frame artefacts. Default: no save. 90 """ 91 92 #: Task this wrapper produces predictions for. Subclasses set this. 93 task: TaskType = TaskType.MONOCULAR_DEPTH # placeholder — must be overridden 94 95 def __init__( 96 self, 97 adapter: Any, 98 *, 99 name: str, 100 save_dir: Optional[str | Path] = None, 101 ) -> None: 102 self._adapter = adapter 103 self.name = name 104 self._save_dir = Path(save_dir) if save_dir else None 105 # Profiler walker reaches the underlying nn.Module via this attr. 106 self.model = adapter 107 108 def setup(self) -> None: 109 """No-op — the adapter loads its own weights at construction.""" 110 return None 111 112 # ──────────────────────── hooks ──────────────────────── 113 114 def extract_inputs(self, sample: Sample) -> Any: 115 """Subclass: return whatever the adapter expects per sample. 116 117 Examples: 118 - Depth: ``np.asarray(sample.rgb, dtype=np.uint8)`` 119 - Pose: ``{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]}`` 120 - Stereo depth: ``{"left": metadata["fisheye_left"], 121 "right": metadata["fisheye_right"]}`` 122 """ 123 raise NotImplementedError 124 125 def call_adapter(self, inputs: List[Any]) -> List[Any]: 126 """Subclass override point: how the adapter is invoked on a list. 127 128 Default: ``self._adapter(inputs)`` — works when the adapter 129 accepts a list directly (HF pipelines, our depth adapters). 130 Override for adapters with a different batched contract. 131 """ 132 out = self._adapter(inputs) 133 if not isinstance(out, (list, tuple)): 134 out = [out] 135 return list(out) 136 137 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 138 """Subclass: wrap one model output as the task's Prediction dataclass.""" 139 raise NotImplementedError 140 141 def maybe_save(self, sample: Sample, model_output: Any) -> None: 142 """Optional override — persist per-frame artefacts for analytics.""" 143 return None 144 145 # ──────────────────────── predict() ──────────────────────── 146 147 def predict(self, batch: Sequence[Sample]) -> List[Any]: 148 inputs = [self.extract_inputs(s) for s in batch] 149 outputs = self.call_adapter(inputs) 150 if len(outputs) != len(batch): 151 from ..exceptions import AdapterError 152 153 raise AdapterError( 154 f"adapter returned {len(outputs)} outputs for a batch of {len(batch)}", 155 hint="Check the adapter's batched contract: it must return " 156 "one output per input sample.", 157 ) 158 preds: List[Any] = [] 159 for sample, out in zip(batch, outputs, strict=False): 160 preds.append(self.wrap_output(out, sample)) 161 self.maybe_save(sample, out) 162 return preds
Abstract base — true-batched-dispatch wrapper for any RPX task.
Subclasses set the class-level task (a TaskType) and
override extract_inputs + wrap_output. The harness contract
matches BenchmarkableModel (setup + predict(batch)).
Parameters
adapter
The model callable. Receives a list of whatever extract_inputs
returns (one per sample). Returns a list of model outputs (one per
sample, same order).
name
Display name surfaced in result.json + Box paths.
save_dir
If set, maybe_save writes per-frame artefacts. Default: no save.
95 def __init__( 96 self, 97 adapter: Any, 98 *, 99 name: str, 100 save_dir: Optional[str | Path] = None, 101 ) -> None: 102 self._adapter = adapter 103 self.name = name 104 self._save_dir = Path(save_dir) if save_dir else None 105 # Profiler walker reaches the underlying nn.Module via this attr. 106 self.model = adapter
108 def setup(self) -> None: 109 """No-op — the adapter loads its own weights at construction.""" 110 return None
No-op — the adapter loads its own weights at construction.
114 def extract_inputs(self, sample: Sample) -> Any: 115 """Subclass: return whatever the adapter expects per sample. 116 117 Examples: 118 - Depth: ``np.asarray(sample.rgb, dtype=np.uint8)`` 119 - Pose: ``{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]}`` 120 - Stereo depth: ``{"left": metadata["fisheye_left"], 121 "right": metadata["fisheye_right"]}`` 122 """ 123 raise NotImplementedError
Subclass: return whatever the adapter expects per sample.
Examples:
- Depth:
np.asarray(sample.rgb, dtype=np.uint8) - Pose:
{"rgb_a": ..., "rgb_b": sample.metadata["rgb_b"]} - Stereo depth:
{"left": metadata["fisheye_left"], "right": metadata["fisheye_right"]}
125 def call_adapter(self, inputs: List[Any]) -> List[Any]: 126 """Subclass override point: how the adapter is invoked on a list. 127 128 Default: ``self._adapter(inputs)`` — works when the adapter 129 accepts a list directly (HF pipelines, our depth adapters). 130 Override for adapters with a different batched contract. 131 """ 132 out = self._adapter(inputs) 133 if not isinstance(out, (list, tuple)): 134 out = [out] 135 return list(out)
Subclass override point: how the adapter is invoked on a list.
Default: self._adapter(inputs) — works when the adapter
accepts a list directly (HF pipelines, our depth adapters).
Override for adapters with a different batched contract.
137 def wrap_output(self, model_output: Any, sample: Sample) -> Any: 138 """Subclass: wrap one model output as the task's Prediction dataclass.""" 139 raise NotImplementedError
Subclass: wrap one model output as the task's Prediction dataclass.
141 def maybe_save(self, sample: Sample, model_output: Any) -> None: 142 """Optional override — persist per-frame artefacts for analytics.""" 143 return None
Optional override — persist per-frame artefacts for analytics.
147 def predict(self, batch: Sequence[Sample]) -> List[Any]: 148 inputs = [self.extract_inputs(s) for s in batch] 149 outputs = self.call_adapter(inputs) 150 if len(outputs) != len(batch): 151 from ..exceptions import AdapterError 152 153 raise AdapterError( 154 f"adapter returned {len(outputs)} outputs for a batch of {len(batch)}", 155 hint="Check the adapter's batched contract: it must return " 156 "one output per input sample.", 157 ) 158 preds: List[Any] = [] 159 for sample, out in zip(batch, outputs, strict=False): 160 preds.append(self.wrap_output(out, sample)) 161 self.maybe_save(sample, out) 162 return preds
131@dataclass 132class BenchmarkableModel(BenchmarkModel): 133 """Compose an input adapter, a model, and an output adapter. 134 135 This is the canonical way to plug a model into the RPX benchmark 136 harness. The ``BenchmarkRunner`` only ever sees the 137 :class:`BenchmarkModel` contract (``setup``, ``predict``); all the 138 model-family-specific logic lives in the adapters so the harness 139 stays task-agnostic. 140 141 Example — wrap a plain numpy depth callable:: 142 143 import numpy as np 144 import rpx_benchmark as rpx 145 146 def my_depth(rgb: np.ndarray) -> np.ndarray: 147 return ... # H x W float32, metres 148 149 bm = rpx.make_numpy_depth_model(my_depth, name="my_model") 150 # Hand `bm` to a task runner or the BenchmarkRunner directly. 151 152 Example — wrap a plain numpy callable:: 153 154 def my_depth(rgb: np.ndarray) -> np.ndarray: 155 return some_depth_in_metres 156 157 bm = make_numpy_depth_model(my_depth) 158 """ 159 160 task: TaskType 161 input_adapter: InputAdapter 162 model: Any 163 output_adapter: OutputAdapter 164 invoker: ModelInvoker = field(default=default_invoker) 165 name: str = "benchmarkable_model" 166 setup_hook: Optional[Callable[[], None]] = None 167 168 def setup(self) -> None: 169 if hasattr(self.input_adapter, "setup"): 170 self.input_adapter.setup() 171 if hasattr(self.output_adapter, "setup"): 172 self.output_adapter.setup() 173 if self.setup_hook is not None: 174 self.setup_hook() 175 176 def predict(self, batch: Sequence[Sample]) -> Sequence[Any]: 177 preds = [] 178 for sample in batch: 179 prepared = self.input_adapter.prepare(sample) 180 model_out = self.invoker(self.model, prepared.payload) 181 pred = self.output_adapter.finalize(model_out, prepared.context, sample) 182 preds.append(pred) 183 return preds
Compose an input adapter, a model, and an output adapter.
This is the canonical way to plug a model into the RPX benchmark
harness. The BenchmarkRunner only ever sees the
BenchmarkModel contract (setup, predict); all the
model-family-specific logic lives in the adapters so the harness
stays task-agnostic.
Example — wrap a plain numpy depth callable::
import numpy as np
import rpx_benchmark as rpx
def my_depth(rgb: np.ndarray) -> np.ndarray:
return ... # H x W float32, metres
bm = rpx.make_numpy_depth_model(my_depth, name="my_model")
# Hand `bm` to a task runner or the BenchmarkRunner directly.
Example — wrap a plain numpy callable::
def my_depth(rgb: np.ndarray) -> np.ndarray:
return some_depth_in_metres
bm = make_numpy_depth_model(my_depth)
107def default_invoker(model: Any, payload: Any) -> Any: 108 """Call ``model`` with ``payload``. Uses torch no_grad if available.""" 109 try: 110 import torch # noqa: F401 111 except ImportError: 112 return _dispatch(model, payload) 113 114 import torch 115 116 with torch.no_grad(): 117 return _dispatch(model, payload)
Call model with payload. Uses torch no_grad if available.
168 def setup(self) -> None: 169 if hasattr(self.input_adapter, "setup"): 170 self.input_adapter.setup() 171 if hasattr(self.output_adapter, "setup"): 172 self.output_adapter.setup() 173 if self.setup_hook is not None: 174 self.setup_hook()
Load checkpoints, warm CUDA, and do any other one-time init.
The runner calls this exactly once before iterating the
dataset, unless BenchmarkRunner(call_setup=False) was
passed — in which case the caller is responsible.
176 def predict(self, batch: Sequence[Sample]) -> Sequence[Any]: 177 preds = [] 178 for sample in batch: 179 prepared = self.input_adapter.prepare(sample) 180 model_out = self.invoker(self.model, prepared.payload) 181 pred = self.output_adapter.finalize(model_out, prepared.context, sample) 182 preds.append(pred) 183 return preds
Run inference on a batch of samples.
Parameters
batch : sequence of Sample
One or more samples. Length equals dataset.batch_size
except possibly for the final tail batch.
Returns
sequence
One task-specific Prediction dataclass per input sample,
in the same order. The prediction dataclass must match
what MetricSuite expects for this task.
Raises
ModelError (By convention) when a sample cannot be processed. The runner surfaces it as a clean error rather than a stack trace.
Inherited Members
72@runtime_checkable 73class InputAdapter(Protocol): 74 """Sample → model-ready payload.""" 75 76 def setup(self) -> None: # pragma: no cover - optional hook 77 """Optional one-time setup (e.g., build a processor on first use).""" 78 79 def prepare(self, sample: Sample) -> PreparedInput: ...
Sample → model-ready payload.
1771def _no_init_or_replace_init(self, *args, **kwargs): 1772 cls = type(self) 1773 1774 if cls._is_protocol: 1775 raise TypeError('Protocols cannot be instantiated') 1776 1777 # Already using a custom `__init__`. No need to calculate correct 1778 # `__init__` to call. This can lead to RecursionError. See bpo-45121. 1779 if cls.__init__ is not _no_init_or_replace_init: 1780 return 1781 1782 # Initially, `__init__` of a protocol subclass is set to `_no_init_or_replace_init`. 1783 # The first instantiation of the subclass will call `_no_init_or_replace_init` which 1784 # searches for a proper new `__init__` in the MRO. The new `__init__` 1785 # replaces the subclass' old `__init__` (ie `_no_init_or_replace_init`). Subsequent 1786 # instantiation of the protocol subclass will thus use the new 1787 # `__init__` and no longer call `_no_init_or_replace_init`. 1788 for base in cls.__mro__: 1789 init = base.__dict__.get('__init__', _no_init_or_replace_init) 1790 if init is not _no_init_or_replace_init: 1791 cls.__init__ = init 1792 break 1793 else: 1794 # should not happen 1795 cls.__init__ = object.__init__ 1796 1797 cls.__init__(self, *args, **kwargs)
76 def setup(self) -> None: # pragma: no cover - optional hook 77 """Optional one-time setup (e.g., build a processor on first use)."""
Optional one-time setup (e.g., build a processor on first use).
79 def prepare(self, sample: Sample) -> PreparedInput: ...
82@runtime_checkable 83class OutputAdapter(Protocol): 84 """Model output → RPX prediction object.""" 85 86 def setup(self) -> None: # pragma: no cover - optional hook 87 """Optional one-time setup.""" 88 89 def finalize( 90 self, 91 model_output: Any, 92 context: Dict[str, Any], 93 sample: Sample, 94 ) -> Any: ...
Model output → RPX prediction object.
1771def _no_init_or_replace_init(self, *args, **kwargs): 1772 cls = type(self) 1773 1774 if cls._is_protocol: 1775 raise TypeError('Protocols cannot be instantiated') 1776 1777 # Already using a custom `__init__`. No need to calculate correct 1778 # `__init__` to call. This can lead to RecursionError. See bpo-45121. 1779 if cls.__init__ is not _no_init_or_replace_init: 1780 return 1781 1782 # Initially, `__init__` of a protocol subclass is set to `_no_init_or_replace_init`. 1783 # The first instantiation of the subclass will call `_no_init_or_replace_init` which 1784 # searches for a proper new `__init__` in the MRO. The new `__init__` 1785 # replaces the subclass' old `__init__` (ie `_no_init_or_replace_init`). Subsequent 1786 # instantiation of the protocol subclass will thus use the new 1787 # `__init__` and no longer call `_no_init_or_replace_init`. 1788 for base in cls.__mro__: 1789 init = base.__dict__.get('__init__', _no_init_or_replace_init) 1790 if init is not _no_init_or_replace_init: 1791 cls.__init__ = init 1792 break 1793 else: 1794 # should not happen 1795 cls.__init__ = object.__init__ 1796 1797 cls.__init__(self, *args, **kwargs)
50@dataclass 51class PreparedInput: 52 """Everything a model needs for one sample, plus context for post-processing. 53 54 ``payload`` is whatever the model's forward call accepts. If it is a 55 ``dict``, the default invoker calls ``model(**payload)``; otherwise 56 ``model(payload)``. 57 58 ``context`` is a free-form dict the output adapter receives back. Use 59 it to stash things like target image size, original intrinsics, or 60 any preprocessing metadata the postprocessing step needs. 61 """ 62 63 payload: Any 64 context: Dict[str, Any] = field(default_factory=dict)
Everything a model needs for one sample, plus context for post-processing.
payload is whatever the model's forward call accepts. If it is a
dict, the default invoker calls model(**payload); otherwise
model(payload).
context is a free-form dict the output adapter receives back. Use
it to stash things like target image size, original intrinsics, or
any preprocessing metadata the postprocessing step needs.
107def default_invoker(model: Any, payload: Any) -> Any: 108 """Call ``model`` with ``payload``. Uses torch no_grad if available.""" 109 try: 110 import torch # noqa: F401 111 except ImportError: 112 return _dispatch(model, payload) 113 114 import torch 115 116 with torch.no_grad(): 117 return _dispatch(model, payload)
Call model with payload. Uses torch no_grad if available.
349def make_numpy_depth_model( 350 fn: Callable[[np.ndarray], np.ndarray], 351 *, 352 name: str = "numpy_depth_model", 353 depth_output_kind: str = "metric", 354 native_alignment: str | None = None, 355) -> BenchmarkableModel: 356 """Wrap a plain numpy depth callable as a :class:`BenchmarkableModel`. 357 358 The callable must accept a ``(H, W, 3) uint8`` RGB image and return a 359 ``(H', W') float`` depth map. If ``depth_output_kind="metric"`` the 360 values are metres. Relative outputs remain raw until the runner fits 361 one pooled alignment per ``(scene, phase)`` cell. If ``(H', W') != 362 (H, W)``, the output is bilinearly resized to match the ground truth. 363 364 Parameters 365 ---------- 366 fn : callable 367 The depth function. Signature: ``fn(rgb_uint8) -> depth_float``. 368 name : str 369 Display name used in logs and reports. 370 depth_output_kind : str 371 ``"metric"`` (default) or ``"relative"``. 372 native_alignment : str, optional 373 Alignment used for relative output. Defaults to ``"ls_affine"`` 374 for relative models and ``"none"`` for metric models. 375 376 Examples 377 -------- 378 >>> import numpy as np 379 >>> import rpx_benchmark as rpx 380 >>> def my_depth(rgb): 381 ... return np.full(rgb.shape[:2], 2.0, dtype=np.float32) 382 >>> bm = rpx.make_numpy_depth_model(my_depth, name="mine") 383 >>> bm.task is rpx.TaskType.MONOCULAR_DEPTH 384 True 385 """ 386 if depth_output_kind not in {"metric", "relative"}: 387 raise AdapterError( 388 f"depth_output_kind must be 'metric' or 'relative', got {depth_output_kind!r}", 389 ) 390 model = BenchmarkableModel( 391 task=TaskType.MONOCULAR_DEPTH, 392 input_adapter=_NumpyDepthInput(), 393 model=fn, 394 output_adapter=_NumpyDepthOutput(), 395 invoker=lambda model, payload: model(payload), 396 name=name, 397 ) 398 model.depth_output_kind = depth_output_kind 399 model.native_alignment = native_alignment or ( 400 "ls_affine" if depth_output_kind == "relative" else "none" 401 ) 402 return model
Wrap a plain numpy depth callable as a BenchmarkableModel.
The callable must accept a (H, W, 3) uint8 RGB image and return a
(H', W') float depth map. If depth_output_kind="metric" the
values are metres. Relative outputs remain raw until the runner fits
one pooled alignment per (scene, phase) cell. If (H', W') !=
(H, W), the output is bilinearly resized to match the ground truth.
Parameters
fn : callable
The depth function. Signature: fn(rgb_uint8) -> depth_float.
name : str
Display name used in logs and reports.
depth_output_kind : str
"metric" (default) or "relative".
native_alignment : str, optional
Alignment used for relative output. Defaults to "ls_affine"
for relative models and "none" for metric models.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_depth(rgb):
... return np.full(rgb.shape[:2], 2.0, dtype=np.float32)
>>> bm = rpx.make_numpy_depth_model(my_depth, name="mine")
>>> bm.task is rpx.TaskType.MONOCULAR_DEPTH
True
574def make_numpy_detection_model( 575 fn: Callable[[np.ndarray], Any], 576 *, 577 name: str = "numpy_detection_model", 578 task: TaskType = TaskType.OBJECT_DETECTION, 579) -> BenchmarkableModel: 580 """Wrap a plain numpy detection callable as a :class:`BenchmarkableModel`. 581 582 The callable must accept a ``(H, W, 3) uint8`` RGB image and 583 return either a dict with keys ``"boxes"`` / ``"scores"`` / 584 ``"labels"`` or a ``(boxes, scores, labels)`` tuple in that 585 order. Boxes are pixel coordinates in ``(x1, y1, x2, y2)`` 586 format, scores are floats in ``[0, 1]``, labels are strings. 587 588 Parameters 589 ---------- 590 fn : callable 591 ``fn(rgb_uint8) -> dict | tuple``. 592 name : str 593 Display name for reports. 594 task : TaskType 595 Use :attr:`TaskType.OBJECT_DETECTION` for closed-vocabulary 596 detection or :attr:`TaskType.OPEN_VOCAB_DETECTION` for 597 open-vocab (the Prediction contract is the same). 598 599 Examples 600 -------- 601 >>> import numpy as np 602 >>> import rpx_benchmark as rpx 603 >>> def my_det(rgb): 604 ... return { 605 ... "boxes": np.array([[10, 10, 30, 30]], dtype=np.float32), 606 ... "scores": np.array([0.9], dtype=np.float32), 607 ... "labels": ["cup"], 608 ... } 609 >>> bm = rpx.make_numpy_detection_model(my_det) 610 >>> bm.task is rpx.TaskType.OBJECT_DETECTION 611 True 612 """ 613 return BenchmarkableModel( 614 task=task, 615 input_adapter=_NumpyRgbInput(), 616 model=fn, 617 output_adapter=_NumpyDetectionOutput(), 618 invoker=_passthrough_invoker, 619 name=name, 620 )
Wrap a plain numpy detection callable as a BenchmarkableModel.
The callable must accept a (H, W, 3) uint8 RGB image and
return either a dict with keys "boxes" / "scores" /
"labels" or a (boxes, scores, labels) tuple in that
order. Boxes are pixel coordinates in (x1, y1, x2, y2)
format, scores are floats in [0, 1], labels are strings.
Parameters
fn : callable
fn(rgb_uint8) -> dict | tuple.
name : str
Display name for reports.
task : TaskType
Use TaskType.OBJECT_DETECTION for closed-vocabulary
detection or TaskType.OPEN_VOCAB_DETECTION for
open-vocab (the Prediction contract is the same).
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_det(rgb):
... return {
... "boxes": np.array([[10, 10, 30, 30]], dtype=np.float32),
... "scores": np.array([0.9], dtype=np.float32),
... "labels": ["cup"],
... }
>>> bm = rpx.make_numpy_detection_model(my_det)
>>> bm.task is rpx.TaskType.OBJECT_DETECTION
True
667def make_numpy_grounding_model( 668 fn: Callable[[np.ndarray, str], Any], 669 *, 670 name: str = "numpy_grounding_model", 671) -> BenchmarkableModel: 672 """Wrap a visual-grounding callable as a :class:`BenchmarkableModel`. 673 674 The callable takes ``(rgb_uint8, text)`` and returns either a 675 dict with keys ``"boxes"`` / ``"scores"`` or a tuple 676 ``(boxes, scores)``. Boxes are ``(x1, y1, x2, y2)`` pixel 677 coordinates; scores are floats. The referring expression 678 ``text`` is plucked from ``sample.ground_truth.text`` by the 679 adapter so the callable never sees the GT boxes. 680 681 Examples 682 -------- 683 >>> import numpy as np 684 >>> import rpx_benchmark as rpx 685 >>> def my_ground(rgb, text): 686 ... return ( 687 ... np.array([[10, 10, 30, 30]], dtype=np.float32), 688 ... np.array([0.8], dtype=np.float32), 689 ... ) 690 >>> bm = rpx.make_numpy_grounding_model(my_ground) 691 """ 692 return BenchmarkableModel( 693 task=TaskType.VISUAL_GROUNDING, 694 input_adapter=_NumpyGroundingInput(), 695 model=fn, 696 output_adapter=_NumpyGroundingOutput(), 697 invoker=_grounding_invoker, 698 name=name, 699 )
Wrap a visual-grounding callable as a BenchmarkableModel.
The callable takes (rgb_uint8, text) and returns either a
dict with keys "boxes" / "scores" or a tuple
(boxes, scores). Boxes are (x1, y1, x2, y2) pixel
coordinates; scores are floats. The referring expression
text is plucked from sample.ground_truth.text by the
adapter so the callable never sees the GT boxes.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_ground(rgb, text):
... return (
... np.array([[10, 10, 30, 30]], dtype=np.float32),
... np.array([0.8], dtype=np.float32),
... )
>>> bm = rpx.make_numpy_grounding_model(my_ground)
926def make_numpy_keypoint_model( 927 fn: Callable[[np.ndarray, np.ndarray], Any], 928 *, 929 name: str = "numpy_keypoint_model", 930) -> BenchmarkableModel: 931 """Wrap a keypoint-matching callable as a :class:`BenchmarkableModel`. 932 933 The callable takes ``(rgb_a, rgb_b)`` and returns either a dict 934 with keys ``"points0"``, ``"points1"`` and optional ``"scores"`` 935 or a 2/3-tuple in the same order. Points are ``(N, 2)`` pixel 936 coordinates. 937 938 Examples 939 -------- 940 >>> import numpy as np 941 >>> import rpx_benchmark as rpx 942 >>> def my_matcher(rgb_a, rgb_b): 943 ... pts = np.array([[10, 10], [20, 20]], dtype=np.float32) 944 ... return pts, pts 945 >>> bm = rpx.make_numpy_keypoint_model(my_matcher) 946 """ 947 return BenchmarkableModel( 948 task=TaskType.KEYPOINT_MATCHING, 949 input_adapter=_NumpyKeypointInput(), 950 model=fn, 951 output_adapter=_NumpyKeypointOutput(), 952 invoker=_keypoint_invoker, 953 name=name, 954 )
Wrap a keypoint-matching callable as a BenchmarkableModel.
The callable takes (rgb_a, rgb_b) and returns either a dict
with keys "points0", "points1" and optional "scores"
or a 2/3-tuple in the same order. Points are (N, 2) pixel
coordinates.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_matcher(rgb_a, rgb_b):
... pts = np.array([[10, 10], [20, 20]], dtype=np.float32)
... return pts, pts
>>> bm = rpx.make_numpy_keypoint_model(my_matcher)
450def make_numpy_mask_model( 451 fn: Callable[[np.ndarray], np.ndarray], 452 *, 453 name: str = "numpy_mask_model", 454) -> BenchmarkableModel: 455 """Wrap a plain numpy instance-mask callable as a :class:`BenchmarkableModel`. 456 457 The callable must accept a ``(H, W, 3) uint8`` RGB image and return 458 a ``(H', W') int`` instance mask where pixel values are instance 459 IDs (``0`` is conventionally background). If ``(H', W') != (H, W)`` 460 the output is nearest-neighbour resized to match the GT mask so 461 integer IDs are preserved. 462 463 Parameters 464 ---------- 465 fn : callable 466 ``fn(rgb_uint8) -> mask_int``. 467 name : str 468 469 Examples 470 -------- 471 >>> import numpy as np 472 >>> import rpx_benchmark as rpx 473 >>> def my_seg(rgb): 474 ... mask = np.zeros(rgb.shape[:2], dtype=np.int32) 475 ... mask[rgb.sum(-1) > 384] = 1 # trivial brightness threshold 476 ... return mask 477 >>> bm = rpx.make_numpy_mask_model(my_seg, name="mine") 478 >>> bm.task is rpx.TaskType.OBJECT_SEGMENTATION 479 True 480 """ 481 return BenchmarkableModel( 482 task=TaskType.OBJECT_SEGMENTATION, 483 input_adapter=_NumpyMaskInput(), 484 model=fn, 485 output_adapter=_NumpyMaskOutput(), 486 invoker=lambda model, payload: model(payload), 487 name=name, 488 )
Wrap a plain numpy instance-mask callable as a BenchmarkableModel.
The callable must accept a (H, W, 3) uint8 RGB image and return
a (H', W') int instance mask where pixel values are instance
IDs (0 is conventionally background). If (H', W') != (H, W)
the output is nearest-neighbour resized to match the GT mask so
integer IDs are preserved.
Parameters
fn : callable
fn(rgb_uint8) -> mask_int.
name : str
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_seg(rgb):
... mask = np.zeros(rgb.shape[:2], dtype=np.int32)
... mask[rgb.sum(-1) > 384] = 1 # trivial brightness threshold
... return mask
>>> bm = rpx.make_numpy_mask_model(my_seg, name="mine")
>>> bm.task is rpx.TaskType.OBJECT_SEGMENTATION
True
750def make_numpy_pose_model( 751 fn: Callable[[np.ndarray, np.ndarray], Any], 752 *, 753 name: str = "numpy_pose_model", 754) -> BenchmarkableModel: 755 """Wrap a relative-camera-pose callable as a :class:`BenchmarkableModel`. 756 757 The callable takes ``(rgb_a, rgb_b)`` and returns either a dict 758 with keys ``"rotation"`` (3×3 rotation matrix or 4-element 759 quaternion) and ``"translation"`` (3-vector, metres) or a 760 ``(rotation, translation)`` tuple. 761 762 Examples 763 -------- 764 >>> import numpy as np 765 >>> import rpx_benchmark as rpx 766 >>> def my_pose(rgb_a, rgb_b): 767 ... return {"rotation": np.eye(3), "translation": np.zeros(3)} 768 >>> bm = rpx.make_numpy_pose_model(my_pose) 769 """ 770 return BenchmarkableModel( 771 task=TaskType.RELATIVE_CAMERA_POSE, 772 input_adapter=_NumpyPoseInput(), 773 model=fn, 774 output_adapter=_NumpyPoseOutput(), 775 invoker=_pose_invoker, 776 name=name, 777 )
Wrap a relative-camera-pose callable as a BenchmarkableModel.
The callable takes (rgb_a, rgb_b) and returns either a dict
with keys "rotation" (3×3 rotation matrix or 4-element
quaternion) and "translation" (3-vector, metres) or a
(rotation, translation) tuple.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_pose(rgb_a, rgb_b):
... return {"rotation": np.eye(3), "translation": np.zeros(3)}
>>> bm = rpx.make_numpy_pose_model(my_pose)
825def make_numpy_sparse_depth_model( 826 fn: Callable[[np.ndarray, np.ndarray], np.ndarray], 827 *, 828 name: str = "numpy_sparse_depth_model", 829) -> BenchmarkableModel: 830 """Wrap a sparse-depth callable as a :class:`BenchmarkableModel`. 831 832 The callable takes ``(rgb_uint8, coords)`` where ``coords`` is a 833 ``(N, 2)`` float32 array of pixel coordinates and returns an 834 ``(N,)`` float32 array of depths in metres at those exact 835 coordinates. 836 837 Examples 838 -------- 839 >>> import numpy as np 840 >>> import rpx_benchmark as rpx 841 >>> def my_sd(rgb, coords): 842 ... return np.full(len(coords), 2.0, dtype=np.float32) 843 >>> bm = rpx.make_numpy_sparse_depth_model(my_sd) 844 """ 845 return BenchmarkableModel( 846 task=TaskType.SPARSE_DEPTH, 847 input_adapter=_NumpySparseDepthInput(), 848 model=fn, 849 output_adapter=_NumpySparseDepthOutput(), 850 invoker=_sparse_depth_invoker, 851 name=name, 852 )
Wrap a sparse-depth callable as a BenchmarkableModel.
The callable takes (rgb_uint8, coords) where coords is a
(N, 2) float32 array of pixel coordinates and returns an
(N,) float32 array of depths in metres at those exact
coordinates.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_sd(rgb, coords):
... return np.full(len(coords), 2.0, dtype=np.float32)
>>> bm = rpx.make_numpy_sparse_depth_model(my_sd)
1044def make_numpy_tracking_model( 1045 fn: Callable[[np.ndarray], Any], 1046 *, 1047 name: str = "numpy_tracking_model", 1048) -> BenchmarkableModel: 1049 """Wrap a per-frame tracking callable as a :class:`BenchmarkableModel`. 1050 1051 The callable takes a single ``rgb`` ``H×W×3 uint8`` frame and 1052 returns the active tracks. Accepted return shapes (see 1053 :func:`_coerce_tracking_output`): 1054 1055 - ``TrackletPrediction`` 1056 - list of :class:`Tracklet` 1057 - list of ``{"track_id": str, "boxes": (T, 4), "scores"?: (T,)}`` 1058 - ``{"tracks": [as above]}`` 1059 1060 Examples 1061 -------- 1062 >>> import numpy as np 1063 >>> import rpx_benchmark as rpx 1064 >>> def my_tracker(rgb): 1065 ... return [{ 1066 ... "track_id": "obj_0", 1067 ... "boxes": np.array([[10, 10, 50, 50]], dtype=np.float32), 1068 ... }] 1069 >>> bm = rpx.make_numpy_tracking_model(my_tracker) 1070 """ 1071 return BenchmarkableModel( 1072 task=TaskType.OBJECT_TRACKING, 1073 input_adapter=_NumpyTrackingInput(), 1074 model=fn, 1075 output_adapter=_NumpyTrackingOutput(), 1076 invoker=_tracking_invoker, 1077 name=name, 1078 )
Wrap a per-frame tracking callable as a BenchmarkableModel.
The callable takes a single rgb H×W×3 uint8 frame and
returns the active tracks. Accepted return shapes (see
_coerce_tracking_output()):
TrackletPrediction- list of
Tracklet - list of
{"track_id": str, "boxes": (T, 4), "scores"?: (T,)} {"tracks": [as above]}
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_tracker(rgb):
... return [{
... "track_id": "obj_0",
... "boxes": np.array([[10, 10, 50, 50]], dtype=np.float32),
... }]
>>> bm = rpx.make_numpy_tracking_model(my_tracker)
298def make_numpy_video_depth_model( 299 fn: Callable[[np.ndarray], np.ndarray], 300 *, 301 name: str = "numpy_video_depth_model", 302 depth_output_kind: str = "metric", 303) -> BenchmarkableModel: 304 """Wrap a plain numpy video-depth callable as a 305 :class:`BenchmarkableModel`. 306 307 The callable must accept a ``(T, H, W, 3) uint8`` RGB clip and 308 return a ``(T, H', W') float`` per-frame depth sequence (metres 309 if metric; up-to-scale if relative — the runner applies per-clip 310 ``(s, t)`` alignment for relative models). If ``(H', W') != 311 (H, W)``, the output is bilinearly resized per frame to match. 312 ``T`` must match the input clip length. 313 314 Parameters 315 ---------- 316 fn : callable 317 Signature: ``fn(rgb_seq_uint8) -> depth_seq_float``. 318 name : str 319 Display name used in logs and reports. 320 depth_output_kind : str 321 ``"metric"`` (default) — model emits metres; runner skips 322 alignment. ``"relative"`` — runner applies per-clip 323 ``(s, t)`` scale-and-shift alignment before metrics. 324 325 Examples 326 -------- 327 >>> import numpy as np 328 >>> import rpx_benchmark as rpx 329 >>> def my_clip_depth(rgb_seq): 330 ... T, H, W, _ = rgb_seq.shape 331 ... return np.full((T, H, W), 2.0, dtype=np.float32) 332 >>> bm = rpx.make_numpy_video_depth_model(my_clip_depth, name="mine") 333 >>> bm.task is rpx.TaskType.VIDEO_DEPTH 334 True 335 """ 336 model = BenchmarkableModel( 337 task=TaskType.VIDEO_DEPTH, 338 input_adapter=_NumpyVideoDepthInput(), 339 model=fn, 340 output_adapter=_NumpyVideoDepthOutput(), 341 invoker=lambda model, payload: model(payload), 342 name=name, 343 ) 344 # Propagate depth_output_kind to the runner's alignment dispatch. 345 model.depth_output_kind = depth_output_kind 346 return model
Wrap a plain numpy video-depth callable as a
BenchmarkableModel.
The callable must accept a (T, H, W, 3) uint8 RGB clip and
return a (T, H', W') float per-frame depth sequence (metres
if metric; up-to-scale if relative — the runner applies per-clip
(s, t) alignment for relative models). If (H', W') !=
(H, W), the output is bilinearly resized per frame to match.
T must match the input clip length.
Parameters
fn : callable
Signature: fn(rgb_seq_uint8) -> depth_seq_float.
name : str
Display name used in logs and reports.
depth_output_kind : str
"metric" (default) — model emits metres; runner skips
alignment. "relative" — runner applies per-clip
(s, t) scale-and-shift alignment before metrics.
Examples
>>> import numpy as np
>>> import rpx_benchmark as rpx
>>> def my_clip_depth(rgb_seq):
... T, H, W, _ = rgb_seq.shape
... return np.full((T, H, W), 2.0, dtype=np.float32)
>>> bm = rpx.make_numpy_video_depth_model(my_clip_depth, name="mine")
>>> bm.task is rpx.TaskType.VIDEO_DEPTH
True