Spaces:
Running
Running
| 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 """ | |
| <div class="viewer-empty"> | |
| <div>No GLB file is available for this selection.</div> | |
| </div> | |
| """ | |
| title = f"{embodiment} / {object_name} / episode {episode}" | |
| srcdoc = f"""<!doctype html> | |
| <html> | |
| <head> | |
| <meta charset="utf-8" /> | |
| <meta name="viewport" content="width=device-width, initial-scale=1" /> | |
| <style> | |
| html, body {{ | |
| width: 100%; | |
| height: 100%; | |
| margin: 0; | |
| overflow: hidden; | |
| background: #f7f7f5; | |
| color: #232323; | |
| font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; | |
| }} | |
| #viewer {{ | |
| position: fixed; | |
| inset: 0; | |
| touch-action: none; | |
| user-select: none; | |
| }} | |
| canvas {{ | |
| display: block; | |
| width: 100%; | |
| height: 100%; | |
| cursor: grab; | |
| outline: none; | |
| touch-action: none; | |
| }} | |
| #label {{ | |
| position: fixed; | |
| left: 14px; | |
| top: 12px; | |
| max-width: min(460px, calc(100vw - 28px)); | |
| padding: 8px 10px; | |
| border: 1px solid rgba(0, 0, 0, 0.12); | |
| background: rgba(255, 255, 255, 0.86); | |
| font-size: 13px; | |
| line-height: 1.3; | |
| word-break: break-word; | |
| pointer-events: none; | |
| }} | |
| #status {{ | |
| position: fixed; | |
| left: 14px; | |
| bottom: 12px; | |
| max-width: min(520px, calc(100vw - 28px)); | |
| padding: 8px 10px; | |
| border: 1px solid rgba(0, 0, 0, 0.12); | |
| background: rgba(255, 255, 255, 0.86); | |
| font-size: 12px; | |
| line-height: 1.3; | |
| color: #494949; | |
| word-break: break-word; | |
| pointer-events: none; | |
| }} | |
| </style> | |
| </head> | |
| <body> | |
| <div id="viewer"></div> | |
| <div id="label"></div> | |
| <div id="status">Loading GLB...</div> | |
| <script> | |
| window.addEventListener("error", (event) => {{ | |
| const status = document.getElementById("status"); | |
| if (status) {{ | |
| status.textContent = "Viewer error: " + (event.message || "unknown script error"); | |
| }} | |
| }}); | |
| window.addEventListener("unhandledrejection", (event) => {{ | |
| const status = document.getElementById("status"); | |
| if (status) {{ | |
| status.textContent = "Viewer error: " + String(event.reason || "unhandled promise rejection"); | |
| }} | |
| }}); | |
| </script> | |
| <script type="importmap"> | |
| {{ | |
| "imports": {{ | |
| "three": "https://unpkg.com/three@0.160.0/build/three.module.js", | |
| "three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/" | |
| }} | |
| }} | |
| </script> | |
| <script type="module"> | |
| import * as THREE from "three"; | |
| import {{ OrbitControls }} from "three/addons/controls/OrbitControls.js"; | |
| import {{ GLTFLoader }} from "three/addons/loaders/GLTFLoader.js"; | |
| const modelUrls = {json.dumps(urls)}; | |
| const title = {json.dumps(title)}; | |
| const container = document.getElementById("viewer"); | |
| const status = document.getElementById("status"); | |
| window.__hrdexInteractionCount = 0; | |
| window.__hrdexLastCameraPosition = null; | |
| window.__hrdexAnimationClipCount = 0; | |
| window.__hrdexAnimationTime = 0; | |
| document.getElementById("label").textContent = title; | |
| const scene = new THREE.Scene(); | |
| scene.background = new THREE.Color(0xf7f7f5); | |
| const clock = new THREE.Clock(); | |
| let mixer = null; | |
| const glbToViewerRotation = new THREE.Euler(-Math.PI / 2, 0, 0); | |
| const renderer = new THREE.WebGLRenderer({{ antialias: true, alpha: false }}); | |
| renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2)); | |
| renderer.setSize(window.innerWidth, window.innerHeight); | |
| renderer.outputColorSpace = THREE.SRGBColorSpace; | |
| renderer.domElement.tabIndex = 0; | |
| renderer.domElement.addEventListener("pointerdown", () => {{ | |
| renderer.domElement.focus(); | |
| renderer.domElement.style.cursor = "grabbing"; | |
| }}); | |
| window.addEventListener("pointerup", () => {{ | |
| renderer.domElement.style.cursor = "grab"; | |
| }}); | |
| renderer.domElement.addEventListener("contextmenu", (event) => event.preventDefault()); | |
| container.appendChild(renderer.domElement); | |
| const camera = new THREE.PerspectiveCamera(45, window.innerWidth / window.innerHeight, 0.001, 1000); | |
| camera.position.set(1.4, 0.9, 1.4); | |
| const controls = new OrbitControls(camera, renderer.domElement); | |
| controls.enableDamping = true; | |
| controls.dampingFactor = 0.08; | |
| controls.enableRotate = true; | |
| controls.enableZoom = true; | |
| controls.enablePan = true; | |
| controls.screenSpacePanning = true; | |
| controls.target.set(0, 0, 0); | |
| controls.mouseButtons = {{ | |
| LEFT: THREE.MOUSE.ROTATE, | |
| MIDDLE: THREE.MOUSE.DOLLY, | |
| RIGHT: THREE.MOUSE.PAN | |
| }}; | |
| controls.touches = {{ | |
| ONE: THREE.TOUCH.ROTATE, | |
| TWO: THREE.TOUCH.DOLLY_PAN | |
| }}; | |
| controls.addEventListener("start", () => {{ | |
| window.__hrdexInteractionCount += 1; | |
| status.textContent = "Interacting: drag to rotate, scroll to zoom, right-drag to pan."; | |
| }}); | |
| controls.addEventListener("change", () => {{ | |
| window.__hrdexLastCameraPosition = camera.position.toArray(); | |
| }}); | |
| controls.addEventListener("end", () => {{ | |
| status.textContent = "Drag to rotate, scroll to zoom, right-drag to pan."; | |
| }}); | |
| scene.add(new THREE.HemisphereLight(0xffffff, 0x8a8a8a, 2.4)); | |
| const keyLight = new THREE.DirectionalLight(0xffffff, 2.2); | |
| keyLight.position.set(3, 5, 4); | |
| scene.add(keyLight); | |
| function frameObject(object) {{ | |
| const box = new THREE.Box3().setFromObject(object); | |
| const size = box.getSize(new THREE.Vector3()); | |
| const center = box.getCenter(new THREE.Vector3()); | |
| const maxDim = Math.max(size.x, size.y, size.z) || 1; | |
| const distance = maxDim / (2 * Math.tan(THREE.MathUtils.degToRad(camera.fov) / 2)); | |
| object.position.sub(center); | |
| controls.target.set(0, 0, 0); | |
| camera.near = Math.max(maxDim / 1000, 0.001); | |
| camera.far = Math.max(maxDim * 100, 100); | |
| camera.position.set(distance * 0.85, distance * 0.55, distance * 0.85); | |
| camera.updateProjectionMatrix(); | |
| controls.update(); | |
| status.textContent = "Drag to rotate, scroll to zoom, right-drag to pan."; | |
| }} | |
| function playAnimations(gltf, model) {{ | |
| const clips = gltf.animations || []; | |
| window.__hrdexAnimationClipCount = clips.length; | |
| if (clips.length === 0) {{ | |
| status.textContent = "No animation clips in this GLB. Drag to rotate, scroll to zoom, right-drag to pan."; | |
| return; | |
| }} | |
| mixer = new THREE.AnimationMixer(model); | |
| for (const clip of clips) {{ | |
| const action = mixer.clipAction(clip); | |
| action.reset(); | |
| action.setLoop(THREE.LoopRepeat, Infinity); | |
| action.play(); | |
| }} | |
| status.textContent = `Playing ${{clips.length}} animation clip${{clips.length === 1 ? "" : "s"}}. Drag to rotate, scroll to zoom, right-drag to pan.`; | |
| }} | |
| function onModelLoaded(gltf) {{ | |
| const model = gltf.scene; | |
| model.rotation.copy(glbToViewerRotation); | |
| model.updateMatrixWorld(true); | |
| scene.add(model); | |
| frameObject(model); | |
| playAnimations(gltf, model); | |
| }} | |
| function loadModel(urlIndex = 0) {{ | |
| const currentUrl = modelUrls[urlIndex]; | |
| new GLTFLoader().load( | |
| currentUrl, | |
| onModelLoaded, | |
| undefined, | |
| (error) => {{ | |
| if (urlIndex + 1 < modelUrls.length) {{ | |
| console.warn(`Failed to load ${{currentUrl}}; trying alternate route.`, error); | |
| status.textContent = "Retrying alternate GLB route..."; | |
| loadModel(urlIndex + 1); | |
| return; | |
| }} | |
| console.error(error); | |
| status.textContent = "Failed to load GLB: " + modelUrls.join(" or "); | |
| }} | |
| ); | |
| }} | |
| loadModel(); | |
| function onResize() {{ | |
| const width = window.innerWidth; | |
| const height = window.innerHeight; | |
| camera.aspect = width / height; | |
| camera.updateProjectionMatrix(); | |
| renderer.setSize(width, height); | |
| }} | |
| window.addEventListener("resize", onResize); | |
| function animate() {{ | |
| requestAnimationFrame(animate); | |
| const delta = clock.getDelta(); | |
| if (mixer) {{ | |
| mixer.update(delta); | |
| window.__hrdexAnimationTime += delta; | |
| }} | |
| controls.update(); | |
| renderer.render(scene, camera); | |
| }} | |
| animate(); | |
| </script> | |
| </body> | |
| </html>""" | |
| return f""" | |
| <iframe | |
| title="HRDexDB 3D GLB viewer" | |
| class="hrdex-viewer-frame" | |
| srcdoc="{html.escape(srcdoc, quote=True)}" | |
| sandbox="allow-scripts allow-same-origin allow-pointer-lock" | |
| ></iframe> | |
| """ | |
| 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)]) | |