from __future__ import annotations import html import json from pathlib import Path from urllib.parse import quote import gradio as gr ROOT = Path(__file__).resolve().parent GLB_ROOT = ROOT / "assets" / "glb" CATALOG_PATH = ROOT / "assets" / "glb_catalog.json" DATASET_REPO_ID = "HRDexDB/HRDexDB-glb" DATASET_REVISION = "main" SPACE_REPO_ID = "HRDexDB/HRDexDB-Visualizer" SPACE_REVISION = "main" def _episode_sort_key(path: Path) -> tuple[int, str]: try: return (int(path.stem), path.stem) except ValueError: return (10**9, path.stem) def build_catalog() -> dict[str, dict[str, list[str]]]: if CATALOG_PATH.exists(): with CATALOG_PATH.open() as handle: return json.load(handle) catalog: dict[str, dict[str, list[str]]] = {} if not GLB_ROOT.exists(): return catalog for embodiment_dir in sorted(path for path in GLB_ROOT.iterdir() if path.is_dir()): objects: dict[str, list[str]] = {} for object_dir in sorted(path for path in embodiment_dir.iterdir() if path.is_dir()): episodes = [path.stem for path in sorted(object_dir.glob("*.glb"), key=_episode_sort_key)] if episodes: objects[object_dir.name] = episodes if objects: catalog[embodiment_dir.name] = objects return catalog CATALOG = build_catalog() def first_selection() -> tuple[str | None, str | None, str | None]: if not CATALOG: return None, None, None embodiment = next(iter(CATALOG)) object_name = next(iter(CATALOG[embodiment])) episode = CATALOG[embodiment][object_name][0] return embodiment, object_name, episode def dataset_glb_path(embodiment: str | None, object_name: str | None, episode: str | None) -> Path | None: if not embodiment or not object_name or not episode: return None if episode not in CATALOG.get(embodiment, {}).get(object_name, []): return None return Path(embodiment) / object_name / f"{episode}.glb" def local_glb_path(embodiment: str | None, object_name: str | None, episode: str | None) -> Path | None: dataset_path = dataset_glb_path(embodiment, object_name, episode) if dataset_path is None: return None return Path("assets") / "glb" / embodiment / object_name / f"{episode}.glb" def glb_urls(embodiment: str | None, object_name: str | None, episode: str | None) -> list[str]: dataset_path = dataset_glb_path(embodiment, object_name, episode) local_path = local_glb_path(embodiment, object_name, episode) if dataset_path is None or local_path is None: return [] quoted_dataset_path = quote(dataset_path.as_posix()) quoted_local_path = quote(local_path.as_posix()) return [ f"https://huggingface.co/datasets/{DATASET_REPO_ID}/resolve/{DATASET_REVISION}/{quoted_dataset_path}", f"https://huggingface.co/spaces/{SPACE_REPO_ID}/resolve/{SPACE_REVISION}/{quoted_local_path}", "/gradio_api/file=" + quoted_local_path, "/file=" + quoted_local_path, ] def viewer_html(embodiment: str | None, object_name: str | None, episode: str | None) -> str: urls = glb_urls(embodiment, object_name, episode) if not urls: return """
No GLB file is available for this selection.
""" title = f"{embodiment} / {object_name} / episode {episode}" srcdoc = f"""
Loading GLB...
""" return f""" """ def selection_summary(embodiment: str | None, object_name: str | None, episode: str | None) -> str: dataset_path = dataset_glb_path(embodiment, object_name, episode) local_path = local_glb_path(embodiment, object_name, episode) if dataset_path is None or local_path is None: return "No selected GLB." file_path = ROOT / local_path try: size_mb = file_path.stat().st_size / (1024 * 1024) return f"Selected: `{dataset_path.as_posix()}` ({size_mb:.1f} MB local fallback)" except OSError: return f"Selected: `{dataset_path.as_posix()}`" def object_choices(embodiment: str | None) -> list[str]: return list(CATALOG.get(embodiment or "", {}).keys()) def episode_choices(embodiment: str | None, object_name: str | None) -> list[str]: return CATALOG.get(embodiment or "", {}).get(object_name or "", []) def on_embodiment_change(embodiment: str | None): objects = object_choices(embodiment) object_name = objects[0] if objects else None episodes = episode_choices(embodiment, object_name) episode = episodes[0] if episodes else None return ( gr.update(choices=objects, value=object_name), gr.update(choices=episodes, value=episode), viewer_html(embodiment, object_name, episode), selection_summary(embodiment, object_name, episode), ) def on_object_change(embodiment: str | None, object_name: str | None): episodes = episode_choices(embodiment, object_name) episode = episodes[0] if episodes else None return ( gr.update(choices=episodes, value=episode), viewer_html(embodiment, object_name, episode), selection_summary(embodiment, object_name, episode), ) def on_episode_change(embodiment: str | None, object_name: str | None, episode: str | None): return viewer_html(embodiment, object_name, episode), selection_summary(embodiment, object_name, episode) initial_embodiment, initial_object, initial_episode = first_selection() CSS = """ .gradio-container { max-width: none !important; } .hrdex-layout { gap: 14px; } .hrdex-controls { min-width: 260px; max-width: 340px; } .hrdex-viewer-frame { display: block; width: 100%; height: min(74vh, 760px); min-height: 520px; border: 1px solid #d6d6d0; background: #f7f7f5; } .viewer-empty { height: min(74vh, 760px); min-height: 520px; display: flex; align-items: center; justify-content: center; border: 1px solid #d6d6d0; background: #f7f7f5; color: #555; } @media (max-width: 760px) { .hrdex-controls { max-width: none; } .hrdex-viewer-frame, .viewer-empty { height: 68vh; min-height: 420px; } } """ with gr.Blocks(title="HRDexDB Visualizer", css=CSS) as demo: gr.Markdown("# HRDexDB Visualizer") with gr.Row(elem_classes=["hrdex-layout"]): with gr.Column(scale=0, min_width=280, elem_classes=["hrdex-controls"]): embodiment_dropdown = gr.Dropdown( label="Embodiment", choices=list(CATALOG.keys()), value=initial_embodiment, interactive=True, ) object_dropdown = gr.Dropdown( label="Object", choices=object_choices(initial_embodiment), value=initial_object, interactive=True, ) episode_dropdown = gr.Dropdown( label="Episode", choices=episode_choices(initial_embodiment, initial_object), value=initial_episode, interactive=True, ) summary = gr.Markdown(selection_summary(initial_embodiment, initial_object, initial_episode)) with gr.Column(scale=1, min_width=420): viewer = gr.HTML(viewer_html(initial_embodiment, initial_object, initial_episode)) embodiment_dropdown.change( on_embodiment_change, inputs=embodiment_dropdown, outputs=[object_dropdown, episode_dropdown, viewer, summary], ) object_dropdown.change( on_object_change, inputs=[embodiment_dropdown, object_dropdown], outputs=[episode_dropdown, viewer, summary], ) episode_dropdown.change( on_episode_change, inputs=[embodiment_dropdown, object_dropdown, episode_dropdown], outputs=[viewer, summary], ) if __name__ == "__main__": demo.launch(allowed_paths=[str(GLB_ROOT)])