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)])