ego-dex-viewer / app.py
pablovela5620's picture
final working state
a5b259e
"""
Demonstrates integrating Rerun visualization with Gradio.
Provides example implementations of data streaming, keypoint annotation, and dynamic
visualization across multiple Gradio tabs using Rerun's recording and visualization capabilities.
"""
import math
import os
import tempfile
import time
import uuid
import cv2
import gradio as gr
import rerun as rr
import rerun.blueprint as rrb
from gradio_rerun import Rerun
from gradio_rerun.events import (
SelectionChange,
TimelineChange,
TimeUpdate,
)
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from huggingface_hub import HfApi, HfFileSystem
from pathlib import Path
API = HfApi() # thin REST wrapper
FS = HfFileSystem() # fsspec-style filesystem :contentReference[oaicite:0]{index=0}
REPO = "pablovela5620/ego-dex-rrd"
SEQUENCE_NAMES = [
"add_remove_lid",
"arrange_topple_dominoes",
"assemble_disassemble_furniture_bench_chair",
"assemble_disassemble_furniture_bench_desk",
"assemble_disassemble_furniture_bench_drawer",
"assemble_disassemble_furniture_bench_lamp",
"assemble_disassemble_furniture_bench_square_table",
"assemble_disassemble_furniture_bench_stool",
"assemble_disassemble_legos",
"assemble_disassemble_soft_legos",
"assemble_disassemble_structures",
"assemble_disassemble_tiles",
"assemble_jenga",
"basic_fold",
"basic_pick_place",
"boil_serve_egg",
"braid_unbraid",
"build_unstack_lego",
"charge_uncharge_airpods",
"charge_uncharge_device",
"clean_cups",
"clean_surface",
"clean_tableware",
"clip_unclip_papers",
"color",
"crumple_flatten_paper",
"deal_gather_cards",
"declutter_desk",
"dry_hands",
"fidget_magnetic_spinner_rings",
"flip_coin",
"flip_pages",
"fold_stack_unstack_unfold_cloths",
"fold_unfold_paper_basic",
"fold_unfold_paper_origami",
"fry_bread",
"fry_egg",
"gather_roll_dice",
"insert_dump_blocks",
"insert_remove_airpods",
"insert_remove_bagging",
"insert_remove_bookshelf",
"insert_remove_cups_from_rack",
"insert_remove_drawer",
"insert_remove_furniture_bench_cabinet",
"insert_remove_furniture_bench_round_table",
"insert_remove_plug_socket",
"insert_remove_shirt_in_tube",
"insert_remove_tennis_ball",
"insert_remove_usb",
"insert_remove_utensils",
"knead_slime",
"load_dispense_ice",
"lock_unlock_key",
"make_sandwich",
"measure_objects",
"open_close_insert_remove_box",
"open_close_insert_remove_case",
"open_close_insert_remove_tupperware",
"paint_clean_brush",
"peel_place_sticker",
"pick_place_food",
"pick_up_and_put_down_case_or_bag",
"play_mancala",
"play_piano",
"play_reset_connect_four",
"point_and_click_remote",
"pour",
"push_pop_toy",
"put_away_set_up_board_game",
"put_in_take_out_glasses",
"put_toothpaste_on_toothbrush",
"rake_smooth_zen_garden",
"roll_ball",
"scoop_dump_ice",
"screw_unscrew_allen_fixture",
"screw_unscrew_bottle_cap",
"screw_unscrew_fingers_fixture",
"set_up_clean_up_chessboard",
"setup_cleanup_table",
"sleeve_unsleeve_cards",
"slot_batteries",
"sort_beads",
"stack",
"stack_remove_jenga",
"stack_unstack_bowls",
"stack_unstack_cups",
"stack_unstack_plates",
"stack_unstack_tupperware",
"staple_paper",
"stock_unstock_fridge",
"sweep_dustpan",
"thread_unthread_bead_necklace",
"throw_and_catch_ball",
"throw_collect_objects",
"tie_and_untie_shoelace",
"tie_untie_rubberband",
"type_keyboard",
"use_chopsticks",
"use_rubiks_cube",
"vertical_pick_place",
"wash_fruit",
"wash_kitchen_dishes",
"wash_put_away_dishes",
"wipe_kitchen_surfaces",
"wipe_screen",
"wrap",
"wrap_unwrap_food",
"write",
"zip_unzip_bag",
"zip_unzip_case",
]
def show_dataset(task_name: str, episode_index: str):
episode_index = f"{int(episode_index):05d}"
url_str = f"https://huggingface.co/datasets/pablovela5620/ego-dex-rrd/resolve/main/{task_name}/{episode_index}.rrd"
return url_str
def list_episodes(task:str) -> list[str]:
"""
Return ["00000", "00001", ...] for the chosen task folder.
"""
# fastest: one HTTP hit that returns the whole tree once
files = API.list_repo_files(REPO, repo_type="dataset")
return sorted(
{Path(f).stem for f in files
if f.startswith(f"{task}/") and f.endswith(".rrd")}
)
default_task = SEQUENCE_NAMES[0] # "add_remove_lid"
initial_eps = list_episodes(default_task) # ["00000", "00001", …]
with gr.Blocks() as demo:
with gr.Tab("Hosted RRD"):
with gr.Row():
with gr.Column(scale=1):
task_name = gr.Dropdown(
label="Task Name",
choices=SEQUENCE_NAMES,
value=default_task,
)
episode_index = gr.Dropdown(
label="Episode Index",
choices=initial_eps,
value=initial_eps[0] if initial_eps else None,
)
def _update_eps(t): # Gradio wants a fn
eps = list_episodes(t)
return gr.update(choices=eps, value=eps[0] if eps else None)
task_name.change(_update_eps, inputs=task_name, outputs=episode_index)
button = gr.Button("Show Dataset")
with gr.Column(scale=4):
viewer = Rerun(
streaming=True,
panel_states={
"time": "collapsed",
"blueprint": "hidden",
"selection": "hidden",
},
)
# choose_rrd.change(lambda x: x, inputs=[choose_rrd], outputs=[viewer])
button.click(
fn=show_dataset,
inputs=[task_name, episode_index],
outputs=[viewer]
)
if __name__ == "__main__":
demo.launch(ssr_mode=False)