rpx_benchmark.tasks

15@dataclass
16class MonocularDepthRunConfig(TaskRunConfig):
17    compute_temporal_stability: bool = True

15@dataclass
16class SegmentationRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class ObjectDetectionRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class VisualGroundingRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class RelativePoseRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class KeypointMatchingRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class SparseDepthRunConfig(TaskRunConfig):
17    pass

15@dataclass
16class ObjectTrackingRunConfig(TaskRunConfig):
17    pass

20def run_monocular_depth(cfg: MonocularDepthRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.MONOCULAR_DEPTH,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25        compute_ts=cfg.compute_temporal_stability,
26    )

20def run_segmentation(cfg: SegmentationRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.OBJECT_SEGMENTATION,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25        compute_ts=True,
26    )

20def run_object_detection(cfg: ObjectDetectionRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.OBJECT_DETECTION,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25    )

28def run_open_vocab_detection(cfg: ObjectDetectionRunConfig) -> PipelineResult:
29    return run_pipeline(
30        task=TaskType.OPEN_VOCAB_DETECTION,
31        primary_metric=PRIMARY_METRIC,
32        cfg=cfg,
33    )

20def run_visual_grounding(cfg: VisualGroundingRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.VISUAL_GROUNDING,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25    )

20def run_relative_pose(cfg: RelativePoseRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.RELATIVE_CAMERA_POSE,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25        compute_ts=False,
26    )

20def run_keypoint_matching(cfg: KeypointMatchingRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.KEYPOINT_MATCHING,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25        compute_ts=False,
26    )

20def run_sparse_depth(cfg: SparseDepthRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.SPARSE_DEPTH,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25    )

20def run_object_tracking(cfg: ObjectTrackingRunConfig) -> PipelineResult:
21    return run_pipeline(
22        task=TaskType.OBJECT_TRACKING,
23        primary_metric=PRIMARY_METRIC,
24        cfg=cfg,
25        compute_ts=False,
26    )

 38@dataclass
 39class TaskRunConfig:
 40    """Base config for every task runner.
 41
 42    Subclass with ``@dataclass`` when a task wants extra fields. The
 43    user supplies ``model`` — the toolkit doesn't ship any.
 44    """
 45
 46    model: Optional[BenchmarkableModel] = None
 47
 48    split: Difficulty | str = Difficulty.HARD
 49    repo_id: Optional[str] = None
 50    cache_dir: Optional[str] = None
 51    revision: Optional[str] = DEFAULT_REVISION
 52    batch_size: int = 1
 53    #: Refuse CPU fallback when a CUDA run was requested. Production smoke
 54    #: and benchmark launchers set this True; lightweight library users keep
 55    #: the historical best-effort fallback by default.
 56    require_cuda: bool = False
 57    #: Skip the first-sample torch FLOP counter. Useful for smoke gates and
 58    #: memory-constrained GPUs where instrumentation can OOM independently of
 59    #: the real forward pass.
 60    skip_flops: bool = False
 61    #: Keep only the first N manifest entries. This is a sample count for
 62    #: frame tasks and a clip count for Video Depth.
 63    max_samples: Optional[int] = None
 64
 65    #: When set, the pipeline skips :func:`download_split` and loads
 66    #: the manifest from disk. Use for runs against a locally-staged
 67    #: lossless dataset (e.g. ``dataset_hub.cli manifest`` output)
 68    #: before the v2-webp tree has been uploaded to HF. Path must point
 69    #: at the JSON file (typically ``<staging>/manifests/<task>/<split>.json``);
 70    #: the ``root`` inside that JSON resolves modality file paths.
 71    manifest_path: Optional[str] = None
 72
 73    device: str = "cuda"
 74    output_dir: Optional[str] = None
 75    progress: Optional[ProgressCallback] = None
 76    #: Optional protocol-specific metric suite. The default remains the
 77    #: registered public suite for the task.
 78    metric_suite: Optional[MetricSuite] = None
 79
 80    #: When True, mirror the per-run output directory to UTD Box at
 81    #: ``<box_folder_id>/<task>/<model>/<split>/`` after the run finishes.
 82    #: Requires ``BOX_DEVELOPER_TOKEN`` in the environment. Idempotent
 83    #: re-uploads (size-matched skip).
 84    upload_to_box: bool = False
 85    #: Box folder id to root the upload under. Defaults to the team's
 86    #: RPX-Outputs folder. Unused unless ``upload_to_box=True``.
 87    box_folder_id: Optional[str] = None
 88
 89    def __post_init__(self) -> None:
 90        if self.model is None:
 91            raise ConfigError(
 92                f"{type(self).__name__}.model is required. Wrap your "
 93                "model via rpx.make_numpy_<task>_model(fn) or build a "
 94                "BenchmarkableModel with your own input/output adapters.",
 95            )
 96        if isinstance(self.split, str):
 97            try:
 98                self.split = Difficulty(self.split)
 99            except ValueError as e:
100                raise ConfigError(
101                    f"Unknown split {self.split!r}. Use one of: {[d.value for d in Difficulty]}",
102                ) from e
103        if self.batch_size < 1:
104            raise ConfigError(f"batch_size must be >= 1, got {self.batch_size}")
105        if self.max_samples is not None and self.max_samples < 1:
106            raise ConfigError(f"max_samples must be >= 1, got {self.max_samples}")

145def run_pipeline(
146    *,
147    task: TaskType,
148    primary_metric: str,
149    cfg: TaskRunConfig,
150    compute_ts: bool = True,
151    compute_sgc: bool = False,
152) -> PipelineResult:
153    """Run a task end-to-end: download, load, predict, score, report.
154
155    Parameters
156    ----------
157    task : TaskType
158    primary_metric : str
159        Metric key that drives ESD-weighted phase scoring.
160    cfg : TaskRunConfig
161        Must already have ``model`` set.
162    compute_ts : bool
163        Compute Temporal Stability. Disable for tasks where adjacent
164        samples are unrelated (e.g. keypoint matching).
165    compute_sgc : bool
166        Compute Stack Geometric Coherence. Only meaningful for
167        segmentation with depth available.
168    """
169    cfg.device = resolve_device(cfg.device, require_cuda=cfg.require_cuda)
170    split_name = cfg.split.value if isinstance(cfg.split, Difficulty) else str(cfg.split)
171    repo_id = cfg.repo_id or DEFAULT_REPO_ID
172
173    log.info("pipeline: task=%s split=%s device=%s", task.value, split_name, cfg.device)
174
175    if cfg.manifest_path:
176        manifest_path = Path(cfg.manifest_path)
177        if not manifest_path.exists():
178            raise ConfigError(
179                f"manifest_path does not exist: {manifest_path}",
180                hint="Generate per-task manifests with "
181                "`python -m rpx_benchmark.dataset_hub.cli manifest "
182                "--src <raw> --staging <out> --splits "
183                "benchmark/data/splits/scene_splits.json` and pass the "
184                "resulting <out>/manifests/<task>/<split>.json path.",
185            )
186        log.info("using local manifest: %s", manifest_path)
187    else:
188        manifest_path = download_split(
189            task=task,
190            split=cfg.split,
191            repo_id=repo_id,
192            cache_dir=cfg.cache_dir,
193            revision=cfg.revision,
194            max_samples=cfg.max_samples,
195        )
196    dataset = RPXDataset.from_manifest(
197        manifest_path,
198        batch_size=cfg.batch_size,
199        max_samples=cfg.max_samples,
200    )
201    log.info("loaded %d samples from %s", len(dataset), manifest_path)
202
203    model = cfg.model
204    name = getattr(model, "name", "model")
205
206    model.setup()
207    efficiency = _count_params(model)
208    if efficiency.params_m is not None:
209        log.info("model %s: %.2f M params", name, efficiency.params_m)
210    # SystemCard auto-detects GPU / precision / batch_size and is
211    # stamped onto every cell row so cross-host aggregation can
212    # disambiguate cells produced on different hardware.
213    efficiency.system_card = SystemCard.auto_detect(
214        precision=getattr(model, "native_precision", "fp32"),
215        batch_size=cfg.batch_size,
216    )
217
218    runner = BenchmarkRunner(
219        model=model,
220        dataset=dataset,
221        metric_suite=cfg.metric_suite or MetricSuite.for_task(task),
222        call_setup=False,
223    )
224    result, dr_report = runner.run_with_report(
225        primary_metric=primary_metric,
226        model_name=name,
227        efficiency=efficiency,
228        compute_ts=compute_ts,
229        compute_sgc_flag=compute_sgc,
230        skip_flops=cfg.skip_flops,
231        progress=cfg.progress,
232    )
233
234    safe_name = name.replace("/", "__")
235    out_dir = Path(cfg.output_dir or f"./rpx_results/{safe_name}/{split_name}")
236    out_dir.mkdir(parents=True, exist_ok=True)
237    json_path = out_dir / "result.json"
238    md_path = out_dir / "summary.md"
239    cells_path = out_dir / "cells.parquet"
240    per_sample_path = out_dir / "per_sample_metrics.parquet"
241
242    try:
243        import pyarrow as pa
244        import pyarrow.parquet as pq
245
246        pq.write_table(pa.Table.from_pylist(result.per_sample), per_sample_path)
247    except ImportError as exc:
248        raise ConfigError(
249            "Writing canonical per-sample metrics requires pyarrow.",
250            hint="Install 'rpx-benchmark[hub]' or use the RPX Docker image.",
251        ) from exc
252
253    # cells.parquet is the canonical artefact downstream Φ / 𝒥 / paper-
254    # table fills read from. Bucket per_sample (one row per frame) into
255    # per-(scene, phase) cells with the mean of every numeric metric.
256    # Mirrors the per-clip pipeline (_video_pipeline.py) so both tasks
257    # feed the same downstream path.
258    cells = cells_from_per_sample(
259        result.per_sample,
260        model_name=name,
261        task=task.value,
262        system_card=efficiency.system_card,
263    )
264    if cells:
265        write_cells(cells, cells_path)
266        log.info("wrote %d cell rows to %s", len(cells), cells_path)
267    else:
268        log.warning(
269            "no cells emitted for %s/%s — per_sample lacks scene/phase fields. "
270            "This usually means the dataset manifest's samples don't carry "
271            "scene_id + phase metadata, or the runner's metric_suite stripped "
272            "them. Downstream Φ/J aggregation will be unable to bucket.",
273            name, split_name,
274        )
275
276    write_json(
277        json_path,
278        task=task.value,
279        model_name=name,
280        split=split_name,
281        repo_id=repo_id,
282        result=result,
283        dr_report=dr_report,
284    )
285    md_path.write_text(
286        format_markdown_summary(
287            task=task.value,
288            model_name=name,
289            split=split_name,
290            repo_id=repo_id,
291            result=result,
292            dr_report=dr_report,
293        ),
294        encoding="utf-8",
295    )
296    log.info("wrote %s and %s", json_path, md_path)
297
298    paths: Dict[str, Path] = {
299        "json": json_path,
300        "markdown": md_path,
301        "out_dir": out_dir,
302        "per_sample_metrics": per_sample_path,
303    }
304    if cells:
305        paths["cells"] = cells_path
306
307    if cfg.upload_to_box:
308        # Lazy import keeps Box deps (requests) out of the import path
309        # for users who never upload.
310        from ..box_upload import DEFAULT_BOX_FOLDER_ID, upload_run_dir
311
312        try:
313            summary = upload_run_dir(
314                out_dir,
315                task=task.value,
316                model_name=name,
317                split=split_name,
318                root_folder_id=cfg.box_folder_id or DEFAULT_BOX_FOLDER_ID,
319            )
320            log.info(
321                "box: %d uploaded, %d skipped → box:%s",
322                summary["uploaded"],
323                summary["skipped"],
324                summary["remote_path"],
325            )
326            paths["box_remote"] = summary["remote_path"]
327        except Exception as e:  # noqa: BLE001
328            # Don't lose the result.json on the floor just because the
329            # token expired — log loudly and continue.
330            log.error(
331                "Box upload failed (%s: %s) — local artefacts at %s are intact; "
332                "use `scripts/sync_results_to_box.py` to recover with a fresh token.",
333                type(e).__name__,
334                e,
335                out_dir,
336            )
337
338    return result, dr_report, paths

35@dataclass
36class TaskSpec:
37    """Per-task plugin record.
38
39    Attributes
40    ----------
41    task, display_name, description
42        Identification.
43    primary_metric
44        Metric key driving the ESD-weighted phase score.
45    required_modalities
46        Hint for hub download scope (``rgb``, ``depth``, ``mask`` …).
47    higher_is_better
48        Direction of the primary metric.
49    run
50        End-to-end callable ``(cfg) -> (BenchmarkResult, report, paths)``.
51    temporal_stability_fn, geometric_coherence_fn
52        Optional deployment-readiness hooks. The runner calls them if
53        set; tasks without a hook skip that metric silently.
54    """
55
56    task: TaskType
57    display_name: str
58    primary_metric: str
59    required_modalities: List[str]
60    run: Callable[[Any], TaskRunResult]
61    higher_is_better: bool = False
62    description: str = ""
63    temporal_stability_fn: Optional[TemporalStabilityFn] = None
64    geometric_coherence_fn: Optional[GeometricCoherenceFn] = None

70def register_task(spec: TaskSpec) -> TaskSpec:
71    if spec.task in _TASK_REGISTRY:
72        raise ConfigError(
73            f"Task {spec.task.value!r} is already registered.",
74            hint="Call unregister_task() first to replace it.",
75        )
76    _TASK_REGISTRY[spec.task] = spec
77    log.debug("registered task %s (primary=%s)", spec.task.value, spec.primary_metric)
78    return spec

81def unregister_task(task: TaskType) -> bool:
82    return _TASK_REGISTRY.pop(task, None) is not None

85def get_task_spec(task: TaskType) -> TaskSpec:
86    if task not in _TASK_REGISTRY:
87        raise ConfigError(
88            f"No runner registered for task {task.value!r}.",
89            hint="Registered: " + ", ".join(sorted(t.value for t in _TASK_REGISTRY)),
90        )
91    return _TASK_REGISTRY[task]

94def available_tasks() -> List[TaskType]:
95    return sorted(_TASK_REGISTRY.keys(), key=lambda t: t.value)

98def iter_task_specs() -> List[TaskSpec]:
99    return [_TASK_REGISTRY[t] for t in available_tasks()]