| import json |
| import gradio as gr |
| from huggingface_hub import hf_hub_download |
|
|
| DATASET_REPO = "atharva-manav/physicsbench-sampled" |
|
|
| TASKS = [ |
| "pb-hard-balance-hard-v1", |
| "pb-hard-balance-medium-v1", |
| "pb-hetbi-bowl-table-v1", |
| "pb-hetbi-cracker-box-jaw-v1", |
| "pb-hetbi-cracker-box-suction-v1", |
| "pb-hetbi-cup-inversion-v1", |
| "pb-hetbi-plate-shelf-v1", |
| "pb-hobi-keyboard-typing-v1", |
| "pb-hobi-plate-shelf-v1", |
| "pb-pr-cup-extract-v1", |
| "pb-pr-domino-select-v1", |
| "pb-pr-domino-single-v1", |
| "pb-pr-edge-slide-v1", |
| "pb-pr-galileo-ramp-drop-v1", |
| "pb-pr-mass-sort-v1", |
| "pb-pr-mirror-pick-place-v1", |
| "pb-pr-occluder-push-v1", |
| "pb-pr-packing-v1", |
| "pb-pr-recover-peg-insert-v1", |
| "pb-pr-seesaw-balance-v1", |
| "pb-pr-sheltered-grasp-v1", |
| "pb-pr-slide-catch-v1", |
| "pb-pr-slip-recovery-v1", |
| "pb-pr-tool-retrieve-v1", |
| "pb-puzzle-ball-cage-escape-v1", |
| "pb-puzzle-gap-funnel-v1", |
| "pb-puzzle-indirect-push-v1", |
| "pb-puzzle-pick-place-ycb-v1", |
| "pb-puzzle-ramp-sort-v1", |
| "pb-puzzle-shape-stack-v1", |
| "pb-puzzle-stack-collapse-recovery-v1", |
| ] |
|
|
| _info_cache = {} |
|
|
|
|
| def get_info(task): |
| if task not in _info_cache: |
| try: |
| path = hf_hub_download( |
| repo_id=DATASET_REPO, |
| filename=f"{task}/meta/info.json", |
| repo_type="dataset", |
| ) |
| with open(path) as f: |
| _info_cache[task] = json.load(f) |
| except Exception as e: |
| return None |
| return _info_cache[task] |
|
|
|
|
| def get_cameras(task): |
| info = get_info(task) |
| if not info: |
| return [] |
| return [k for k, v in info["features"].items() if v.get("dtype") == "video"] |
|
|
|
|
| def on_task_change(task): |
| cameras = get_cameras(task) |
| info = get_info(task) |
| n_ep = info.get("total_episodes", 5) if info else 5 |
| task_desc = info.get("task_id", task) if info else task |
| return ( |
| gr.Dropdown(choices=cameras, value=cameras[0] if cameras else None, interactive=True), |
| gr.Slider(minimum=0, maximum=n_ep - 1, step=1, value=0), |
| task_desc, |
| ) |
|
|
|
|
| def load_video(task, episode_idx, camera): |
| if not task or camera is None: |
| return None, "Pick a task and camera then click Load." |
| try: |
| episode_idx = int(episode_idx) |
| path = hf_hub_download( |
| repo_id=DATASET_REPO, |
| filename=f"{task}/videos/{camera}/chunk-000/file-{episode_idx:03d}.mp4", |
| repo_type="dataset", |
| ) |
| return path, "" |
| except Exception as e: |
| return None, f"Error loading video: {e}" |
|
|
|
|
| with gr.Blocks(title="PhysicsBench Visualizer", theme=gr.themes.Soft()) as demo: |
| gr.Markdown("# PhysicsBench Sampled Dataset Visualizer") |
| gr.Markdown( |
| "31 physics-manipulation tasks · 5 episodes each · multiple cameras \n" |
| f"Data: [{DATASET_REPO}](https://huggingface.co/datasets/{DATASET_REPO})" |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(scale=1): |
| task_dd = gr.Dropdown(choices=TASKS, value=TASKS[0], label="Task", interactive=True) |
| camera_dd = gr.Dropdown(choices=[], label="Camera", interactive=True) |
| episode_slider = gr.Slider(minimum=0, maximum=4, step=1, value=0, label="Episode", interactive=True) |
| task_desc = gr.Textbox(label="Task ID", interactive=False) |
| load_btn = gr.Button("Load Video", variant="primary") |
| err_box = gr.Textbox(label="", visible=True, interactive=False, show_label=False) |
|
|
| with gr.Column(scale=2): |
| video_out = gr.Video(label="Episode Video", autoplay=True) |
|
|
| task_dd.change(on_task_change, inputs=[task_dd], outputs=[camera_dd, episode_slider, task_desc]) |
| load_btn.click(load_video, inputs=[task_dd, episode_slider, camera_dd], outputs=[video_out, err_box]) |
| demo.load(on_task_change, inputs=[task_dd], outputs=[camera_dd, episode_slider, task_desc]) |
|
|
| demo.launch() |
|
|
