rpx_benchmark.examples.benchmark_tasks

  1"""Run six paper task examples from an installed wheel.
  2
  3``--smoke`` creates tiny synthetic fixtures and runs the real public evaluators.
  4It needs no dataset, weights, GPU, or network.
  5
  6For a real run, name the task: its data is downloaded first (only the files it
  7needs), then scored. ``--model module:function`` plugs in your model; without it
  8the matching function in ``rpx_benchmark.examples.dummy_models`` is scored, so
  9every paper task runs before you write code. ``--manifest`` scores an existing
 10manifest instead of downloading. T5/T6 use the canonical VQA JSONL contract,
 11including the verified reference crop for T6; without ``--manifest`` they use
 12the packaged, revision-pinned questions in ``rpx_benchmark/vqa/quickstart_v1``.
 13"""
 14
 15from __future__ import annotations
 16
 17import argparse
 18import hashlib
 19import io
 20import json
 21from pathlib import Path
 22from typing import Any, Callable
 23
 24import numpy as np
 25from PIL import Image
 26
 27import rpx_benchmark as rpx
 28
 29from ..exceptions import ConfigError
 30from ..quickstart import download
 31from ..tasks._pipeline import PipelineResult, TaskRunConfig
 32from ..vqa.contract import VQASample, image_locator, load_manifest, write_manifest
 33from ..vqa.evaluate import evaluate_vqa
 34from ..vqa.hub_rgb import image_cache_name, reference_cache_name
 35from . import dummy_models
 36from .profile_callable import load_callable
 37
 38TASKS: dict[
 39    str,
 40    tuple[
 41        Callable[..., rpx.BenchmarkableModel], type[TaskRunConfig], Callable[[Any], PipelineResult]
 42    ],
 43] = {
 44    "T1": (rpx.make_numpy_depth_model, rpx.MonocularDepthRunConfig, rpx.run_monocular_depth),
 45    "T2": (rpx.make_numpy_video_depth_model, rpx.VideoDepthRunConfig, rpx.run_video_depth),
 46    "T3": (rpx.make_numpy_tracking_model, rpx.ObjectTrackingRunConfig, rpx.run_object_tracking),
 47    "T4": (rpx.make_numpy_pose_model, rpx.RelativePoseRunConfig, rpx.run_relative_pose),
 48}
 49
 50
 51#: Hub manifest folder for each task's ``download_split`` call.
 52#: Packaged, revision-pinned VQA questions for T5/T6 runs without --manifest.
 53VQA_QUICKSTART = Path(__file__).resolve().parents[1] / "vqa" / "quickstart_v1"
 54
 55#: Dummy model scored for each task when no --model is given.
 56DUMMY_MODELS: dict[str, Callable[..., Any]] = {
 57    "T1": dummy_models.image_depth,
 58    "T2": dummy_models.video_depth,
 59    "T3": dummy_models.tracking,
 60    "T4": dummy_models.camera_pose,
 61    "T5": dummy_models.vqa,
 62    "T6": dummy_models.incontext_vqa,
 63}
 64
 65HUB_MANIFESTS = {
 66    "T1": ("monocular_depth", None),
 67    "T2": ("video_depth", None),
 68    "T3": ("object_tracking", None),
 69    "T4": ("relative_camera_pose", "relative_pose"),
 70}
 71
 72
 73def demo_depth(rgb):
 74    """Constant integration fixture; not pretrained inference."""
 75    return np.full(rgb.shape[:-1], 2.1, np.float32)
 76
 77
 78def demo_video_depth(rgb):
 79    """Non-degenerate depth ramp for the video validator's synthetic case."""
 80    return np.broadcast_to(np.linspace(1.9, 2.1, rgb.shape[2]), rgb.shape[:3]).astype(np.float32)
 81
 82
 83def demo_tracking(rgb):
 84    """One fixed track for synthetic data; real trackers need persistent state."""
 85    return [{"track_id": "1", "boxes": np.array([[3, 2, 12, 10]], np.float32)}]
 86
 87
 88def demo_pose(rgb_a, rgb_b):
 89    """Fixed 10-cm translation fixture; not a camera pose estimator."""
 90    return {"rotation": np.eye(3), "translation": np.array([0.1, 0.0, 0.0])}
 91
 92
 93def demo_vqa(image_paths, prompt, max_new_tokens, output_kind):
 94    """Fixed synthetic localization; never use this as a benchmark model."""
 95    return '{"label":"cup","bbox":[0,0,1000,1000]}'
 96
 97
 98def create_smoke(root: Path, task: str) -> tuple[Path, Callable[..., Any]]:
 99    """Create clearly synthetic local inputs with no released annotation edits."""
100    root.mkdir(parents=True, exist_ok=True)
101    rgb = np.full((16, 20, 3), 128, np.uint8)
102    Image.fromarray(rgb).save(root / "rgb.png")
103    if task in {"T5", "T6"}:
104        image = image_locator({"scene_id": "scene012", "kind": "mos", "phase": 1, "frame": "00113"})
105        row = dict(
106            sample_id=f"synthetic-{task}",
107            scene_id="scene012",
108            kind="mos",
109            phase=1,
110            frame="00113",
111            img_w=20,
112            img_h=16,
113            image=image,
114            question_type="attr_single_color",
115            question="Locate the gray cup.",
116            answer="cup",
117            answer_bbox=[0, 0, 19, 15],
118        )
119        if task == "T6":
120            buffer = io.BytesIO()
121            Image.new("RGB", (4, 4), "gray").save(buffer, format="PNG")
122            crop = buffer.getvalue()
123            row.update(
124                question_type="inctx_attr_single_color",
125                reference_image={**image, "shard": "synthetic/reference/rgb.tar"},
126                reference_crop_bbox=[0, 0, 3, 3],
127                reference_crop_sha256=hashlib.sha256(crop).hexdigest(),
128                source_sample_id="synthetic-T5",
129            )
130        vqa_sample = VQASample.from_dict(row)
131        cache = root / "images"
132        cache.mkdir()
133        # Populate the normal cache with SYNTHETIC pixels, only under this
134        # isolated fixture root. Real runs fetch the manifest's actual images.
135        Image.fromarray(rgb).save(cache / image_cache_name(vqa_sample))
136        if task == "T6":
137            (cache / reference_cache_name(vqa_sample)).write_bytes(crop)
138        manifest = root / "manifest.jsonl"
139        write_manifest([vqa_sample], manifest)
140        return manifest, demo_vqa
141    Image.fromarray(np.full((16, 20), 2000, np.uint16)).save(root / "depth.png")
142    mask = np.zeros((16, 20), np.uint8)
143    mask[2:10, 3:12] = 1
144    Image.fromarray(mask).save(root / "mask.png")
145    np.savez(root / "pose.npz", position=np.zeros(3), orientation=np.array([0.0, 0.0, 0.0, 1.0]))
146    sample: dict[str, Any] = dict(
147        id=f"synthetic-{task}",
148        scene="scene012",
149        scene_id="scene012",
150        phase="interaction",
151        difficulty="easy",
152        rgb="rgb.png",
153    )
154    task_names = {
155        "T1": "monocular_depth",
156        "T2": "video_depth",
157        "T3": "object_tracking",
158        "T4": "relative_camera_pose",
159    }
160    fn: Callable[..., Any]
161    if task == "T1":
162        sample["depth"] = "depth.png"
163        fn = demo_depth
164    elif task == "T2":
165        sample.update(phase=1, frame_filenames=["rgb.png"] * 3, depth_filenames=["depth.png"] * 3)
166        fn = demo_video_depth
167    elif task == "T3":
168        sample["mask"] = "mask.png"
169        fn = demo_tracking
170    else:
171        np.savez(
172            root / "pose_b.npz",
173            position=np.array([0.1, 0.0, 0.0]),
174            orientation=np.array([0.0, 0.0, 0.0, 1.0]),
175        )
176        sample.update(rgb_b="rgb.png", pose_a="pose.npz", pose_b="pose_b.npz")
177        fn = demo_pose
178    manifest = root / "manifest.json"
179    manifest.write_text(
180        json.dumps(
181            {"task": task_names[task], "root": str(root.resolve()), "samples": [sample]}, indent=2
182        )
183        + "\n"
184    )
185    return manifest, fn
186
187
188def run_task(
189    task: str,
190    manifest: Path,
191    fn: Callable[..., Any],
192    output: Path,
193    *,
194    device: str = "cpu",
195    split: str = "hard",
196    max_samples: int | None = None,
197    image_cache: Path | None = None,
198    model_name: str = "custom",
199    model_revision: str = "unspecified",
200    depth_output_kind: str = "metric",
201) -> dict[str, Any]:
202    """Run one public task protocol with a supplied callable and manifest."""
203    if output.exists() and any(output.iterdir()):
204        raise ConfigError(f"Use a new, empty output directory: {output}")
205    if task in {"T5", "T6"}:
206        samples = load_manifest(manifest)
207        if any(sample.is_in_context != (task == "T6") for sample in samples):
208            raise ConfigError(f"{task} manifest uses the wrong one-image/two-image protocol")
209        if max_samples is not None:
210            raise ConfigError("For T5/T6 select a subset manifest first; all its rows are scored")
211        return evaluate_vqa(
212            manifest,
213            fn,
214            model_name=model_name,
215            model_revision=model_revision,
216            output_dir=output,
217            image_cache=image_cache or output / "images",
218        )
219    wrap, config, runner = TASKS[task]
220    kwargs = {"depth_output_kind": depth_output_kind} if task in {"T1", "T2"} else {}
221    cfg = config(
222        model=wrap(fn, name=model_name, **kwargs),
223        manifest_path=str(manifest),
224        output_dir=str(output),
225        device=device,
226        split=split,
227        require_cuda=device.startswith("cuda"),
228        skip_flops=True,
229        max_samples=max_samples,
230    )
231    result, _, paths = runner(cfg)
232    config_path = output / "run_config.json"
233    config_path.write_text(
234        json.dumps(
235            {
236                "task": task,
237                "model": model_name,
238                "model_revision": model_revision,
239                "manifest_sha256": hashlib.sha256(manifest.read_bytes()).hexdigest(),
240                "device": device,
241                "split": split,
242                "max_samples": max_samples,
243                "depth_output_kind": depth_output_kind if task in {"T1", "T2"} else None,
244                "protocol": "public_callable_task_runner",
245                "skip_flops": True,
246            },
247            indent=2,
248        )
249        + "\n"
250    )
251    paths["run_config"] = config_path
252    return {
253        "task": task,
254        "num_samples": result.num_samples,
255        "aggregated": result.aggregated,
256        "paths": {key: str(path) for key, path in paths.items()},
257    }
258
259
260TASK_LABELS = {
261    "T1": "Image depth",
262    "T2": "Video depth",
263    "T3": "Object tracking",
264    "T4": "Relative camera pose",
265    "T5": "VQA",
266    "T6": "In-context VQA",
267}
268# Headline metrics shown in the console summary (all metrics are in result files).
269HEADLINE = {
270    "T1": (("absrel", "AbsRel"), ("delta1", "δ1")),
271    "T2": (("absrel", "AbsRel"), ("tgm", "TGM")),
272    "T3": (("mota", "MOTA"), ("idsw", "IDSW")),
273    "T4": (("rotation_error_deg", "Rot°"), ("translation_angular_deg", "Transl°")),
274    "T5": (("bbox_accuracy_at_0_5", "Acc@0.5"), ("bbox_mean_giou", "GIoU")),
275    "T6": (("bbox_accuracy_at_0_5", "Acc@0.5"), ("bbox_mean_giou", "GIoU")),
276}
277
278
279def summary_line(task: str, report: dict[str, Any], out_dir: Path) -> str:
280    """One readable console line: task, sample count, headline metrics, output folder."""
281    metrics = report.get("aggregated") or report.get("metrics") or {}
282    shown = []
283    for key, label in HEADLINE[task]:
284        value = metrics.get(key)
285        if isinstance(value, (int, float)):
286            shown.append(f"{label} {value:.3f}")
287    samples = report.get("num_samples", "?")
288    return f"  {task}  {TASK_LABELS[task]:<21}{samples:>5} samples   {' · '.join(shown):<34} → {out_dir}"
289
290
291def default_output(task: str, model: str | None) -> Path:
292    """``rpx_results/<task>/<model>``, numbered so repeated runs never collide."""
293    label = model.replace(":", "-").replace("/", "-") if model else "dummy"
294    base = Path("rpx_results") / task / label
295    candidate, n = base, 1
296    while candidate.exists() and any(candidate.iterdir()):
297        n += 1
298        candidate = base.with_name(f"{label}-{n}")
299    return candidate
300
301
302def main() -> None:
303    parser = argparse.ArgumentParser(description=__doc__)
304    parser.add_argument("--task", choices=[*TASKS, "T5", "T6", "all"], required=True)
305    parser.add_argument("--smoke", action="store_true")
306    parser.add_argument("--manifest", type=Path)
307    parser.add_argument("--model", help="module:function; the model is loaded by your module")
308    parser.add_argument("--model-name", default="custom")
309    parser.add_argument("--model-revision", default="unspecified")
310    parser.add_argument("--device", default="cpu")
311    parser.add_argument("--split", choices=["easy", "medium", "hard"], default="hard")
312    parser.add_argument("--max-samples", type=int)
313    parser.add_argument("--depth-output-kind", choices=["metric", "relative"], default="metric")
314    parser.add_argument("--image-cache", type=Path)
315    parser.add_argument(
316        "--full",
317        action="store_true",
318        help="T5/T6: score the paper's full question list instead of the packaged quick start",
319    )
320    parser.add_argument(
321        "--output",
322        type=Path,
323        help="results folder (default: rpx_results/<task>/<model>, numbered if taken)",
324    )
325    parser.add_argument(
326        "--json",
327        action="store_true",
328        help="print the full reports as JSON lines instead of a summary",
329    )
330    args = parser.parse_args()
331    if args.smoke:
332        args.output = args.output or Path("rpx-check")
333        if args.model or args.manifest or args.image_cache or args.device != "cpu":
334            parser.error(
335                "--smoke uses isolated synthetic fixtures and CPU; omit model/manifest/cache"
336            )
337        if args.output.exists() and any(args.output.iterdir()):
338            parser.error("Use a new, empty smoke output directory")
339        task_list = list(TASKS) + ["T5", "T6"] if args.task == "all" else [args.task]
340        if not args.json:
341            print("RPX smoke check: synthetic inputs, CPU, real evaluators")
342        for task in task_list:
343            root = args.output / task
344            manifest, fn = create_smoke(root / "synthetic_inputs", task)
345            report = run_task(
346                task,
347                manifest,
348                fn,
349                root / "results",
350                split="easy",
351                image_cache=root / "synthetic_inputs/images",
352                model_name="synthetic-integration-fixture",
353                model_revision="demo-v1",
354            )
355            if args.json:
356                print(json.dumps({"task": task, "synthetic": True, "report": report}, default=str))
357            else:
358                print(summary_line(task, report, root / "results"))
359        if not args.json:
360            print(
361                f"Done: {len(task_list)} task(s) ran end to end; full results under {args.output}/.\n"
362                "Next, score your own model: https://irvlutd.github.io/RPX/toolkit-docs/models/"
363            )
364    else:
365        if args.task == "all":
366            parser.error("Real runs score one task at a time: pass --task T1 ... T6")
367        if args.full and (args.task not in {"T5", "T6"} or args.manifest):
368            parser.error("--full selects the full T5/T6 question list; drop it, or drop --manifest")
369        output = args.output or default_output(args.task, args.model)
370        manifest = args.manifest
371        max_samples = args.max_samples
372        image_cache = args.image_cache
373        if manifest is None:
374            manifest = download(args.task, args.split, args.max_samples, full=args.full)
375            max_samples = None  # the downloaded manifest already holds exactly these rows
376            if args.task in {"T5", "T6"}:
377                image_cache = image_cache or manifest.parent / "vqa-images"
378        if args.model:
379            fn, model_name = load_callable(args.model), args.model_name
380        else:
381            fn, model_name = DUMMY_MODELS[args.task], "dummy"
382            print(
383                "No --model given: scoring the dummy model in rpx_benchmark.examples.dummy_models. "
384                "Pass --model my_model:predict to score yours."
385            )
386        print(f"Scoring {TASK_LABELS[args.task]} ...")
387        report = run_task(
388            args.task,
389            manifest,
390            fn,
391            output,
392            device=args.device,
393            split=args.split,
394            max_samples=max_samples,
395            image_cache=image_cache,
396            model_name=model_name,
397            model_revision=args.model_revision,
398            depth_output_kind=args.depth_output_kind,
399        )
400        if args.json:
401            print(json.dumps(report, default=str))
402        else:
403            print(summary_line(args.task, report, output))
404
405
406if __name__ == "__main__":
407    main()