rpx_benchmark.examples.dummy_models

 1"""Dummy models for the six RPX paper tasks: copy one, replace its body with your model.
 2
 3Each function shows exactly what RPX passes in and expects back for one task.
 4They run as-is, so you can try any task before writing code::
 5
 6    python -m rpx_benchmark.examples.benchmark_tasks --task T2 --split easy --max-samples 2 \\
 7        --model rpx_benchmark.examples.dummy_models:video_depth
 8
 9The outputs are placeholders, not predictions; their scores mean nothing.
10"""
11
12from __future__ import annotations
13
14import json
15from pathlib import Path
16
17import numpy as np
18
19
20def image_depth(rgb: np.ndarray) -> np.ndarray:
21    """T1 image depth: RGB image ``(H, W, 3)`` uint8 -> depth ``(H, W)`` in metres.
22
23    For relative (up-to-scale) depth, also pass ``--depth-output-kind relative``.
24    """
25    return np.full(rgb.shape[:2], 2.0, np.float32)
26
27
28def video_depth(clip: np.ndarray) -> np.ndarray:
29    """T2 video depth: RGB clip ``(T, H, W, 3)`` uint8 -> depth ``(T, H, W)`` in metres.
30
31    The evaluator rejects a constant map as degenerate, so this placeholder is a
32    left-to-right ramp from 1.5 m to 2.5 m.
33    """
34    t, h, w = clip.shape[:3]
35    ramp = np.linspace(1.5, 2.5, w, dtype=np.float32)
36    return np.broadcast_to(ramp, (t, h, w)).copy()
37
38
39def tracking(rgb: np.ndarray) -> list[dict]:
40    """T3 object tracking: one RGB frame ``(H, W, 3)`` -> the objects tracked so far.
41
42    Called once per frame, in order, so a real tracker keeps its state between
43    calls. Each track has a stable ``track_id`` and its ``boxes`` as
44    ``[x1, y1, x2, y2]`` pixel rows.
45    """
46    h, w = rgb.shape[:2]
47    box = np.array([[w * 0.25, h * 0.25, w * 0.75, h * 0.75]], np.float32)
48    return [{"track_id": "0", "boxes": box}]
49
50
51def camera_pose(rgb_a: np.ndarray, rgb_b: np.ndarray) -> dict:
52    """T4 relative camera pose: two RGB images -> pose of camera B relative to camera A.
53
54    ``rotation`` is a 3x3 matrix (or a quaternion) and ``translation`` a 3-vector in metres.
55    """
56    return {"rotation": np.eye(3), "translation": np.zeros(3)}
57
58
59def vqa(image_paths: tuple[Path, ...], prompt: str, max_new_tokens: int, output_kind: str) -> str:
60    """T5 VQA: one image and a question -> your model's raw text answer.
61
62    ``image_paths`` holds the target image. Answer with one JSON object: the
63    object's ``label`` and its ``bbox`` as ``[x1, y1, x2, y2]`` normalized to 0-1000.
64    Works with local models and hosted APIs alike.
65    """
66    return json.dumps({"label": "object", "bbox": [250, 250, 750, 750]})
67
68
69def incontext_vqa(
70    image_paths: tuple[Path, ...], prompt: str, max_new_tokens: int, output_kind: str
71) -> str:
72    """T6 in-context VQA: a reference crop and a target image -> the same JSON answer as T5.
73
74    ``image_paths`` is ``(reference_crop, target)``; the box refers to the target.
75    """
76    return vqa(image_paths, prompt, max_new_tokens, output_kind)