HorizonRobotics/EmbodiedGen-Gallery-Explorer
1
1# Project EmbodiedGen2#3# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or14# implied. See the License for the specific language governing15# permissions and limitations under the License.16 17 18import os19 20gradio_tmp_dir = os.path.join(21 os.path.dirname(os.path.abspath(__file__)), "gradio_cache"22)23os.makedirs(gradio_tmp_dir, exist_ok=True)24os.environ["GRADIO_TEMP_DIR"] = gradio_tmp_dir25 26import colorsys27import shutil28import uuid29import xml.etree.ElementTree as ET30import zipfile31from pathlib import Path32 33import gradio as gr34import numpy as np35import pandas as pd36import trimesh37from app_style import custom_theme38from embodied_gen.utils.tags import VERSION39 40try:41 from embodied_gen.utils.gpt_clients import GPT_CLIENT as gpt_client42 43 gpt_client.check_connection()44 GPT_AVAILABLE = True45except Exception as e:46 gpt_client = None47 GPT_AVAILABLE = False48 print(49 f"Warning: GPT client could not be initialized. Search will be disabled. Error: {e}"50 )51 52 53# --- Configuration & Data Loading ---54RUNNING_MODE = "hf_remote" # local or hf_remote55CSV_FILE = "dataset_index.csv"56HF_REPO_ID = "HorizonRobotics/EmbodiedGenData"57HF_LOCAL_DIR = "EmbodiedGenData"58CAMERA_ZOOM = 3.259# Gradio's default per-event concurrency limit is 1, which would make every60# user's HF snapshot_download (potentially several seconds) serialize behind61# whichever user clicked first. Give the HF-fetching events their own shared62# pool so several users can download/preview assets in parallel, bounded to63# avoid hammering the HF API with too many simultaneous requests.64HF_IO_CONCURRENCY_ID = "hf_asset_io"65HF_IO_CONCURRENCY_LIMIT = 666 67# Compatible with huggingface space zero GPU68import spaces69from huggingface_hub import snapshot_download70 71 72@spaces.GPU73def fake_gpu_init():74 pass75 76 77fake_gpu_init()78 79if RUNNING_MODE == "local":80 DATA_ROOT = "/horizon-bucket/robot_lab/datasets/embodiedgen/assets_v2"81elif RUNNING_MODE == "hf_remote":82 # Only fetch the index and preview videos up front; per-asset mesh/urdf/83 # usd/mjcf files are pulled lazily on demand (see `ensure_hf_files`).84 snapshot_download(85 repo_id=HF_REPO_ID,86 repo_type="dataset",87 allow_patterns=[88 f"dataset/{CSV_FILE}",89 "dataset/**/*.mp4",90 ],91 local_dir=HF_LOCAL_DIR,92 )93 DATA_ROOT = os.path.join(HF_LOCAL_DIR, "dataset")94else:95 raise ValueError(96 f"Unknown RUNNING_MODE: {RUNNING_MODE}, must be 'local' or 'hf_remote'."97 )98 99 100def ensure_hf_files(rel_patterns: list[str]) -> None:101 """Lazily download files for one asset in hf_remote mode.102 103 Args:104 rel_patterns: Glob patterns relative to ``DATA_ROOT`` (the repo's105 ``dataset/`` folder), e.g. ``"cat/sub/uid/mesh/**"``.106 107 No-op when running locally, where all files already exist on disk.108 Progress toasts are emitted by the calling action handlers.109 """110 if RUNNING_MODE != "hf_remote":111 return112 snapshot_download(113 repo_id=HF_REPO_ID,114 repo_type="dataset",115 allow_patterns=[f"dataset/{p}" for p in rel_patterns],116 local_dir=HF_LOCAL_DIR,117 )118 119 120csv_path = os.path.join(DATA_ROOT, CSV_FILE)121df = pd.read_csv(csv_path)122TMP_DIR = os.path.join(123 os.path.dirname(os.path.abspath(__file__)), "sessions/asset_viewer"124)125os.makedirs(TMP_DIR, exist_ok=True)126 127 128# --- Custom CSS for Styling ---129css = """130.gradio-container .gradio-group { box-shadow: 0 2px 4px rgba(0,0,0,0.05) !important; }131#asset-gallery { border: 1px solid #E5E7EB; border-radius: 8px; padding: 8px; background-color: #F9FAFB; }132#download-row button { white-space: nowrap; }133"""134 135lighting_css = """136<style>137#visual_mesh canvas { filter: brightness(2.2) !important; }138#collision_mesh_a canvas, #collision_mesh_b canvas { filter: brightness(1.0) !important; }139</style>140"""141 142_prev_temp = {}143 144 145def _unique_path(146 src_path: str | None, session_hash: str, kind: str147) -> str | None:148 """Link/copy src to GRADIO_TEMP_DIR/session_hash with random filename. Always return a fresh URL."""149 if not src_path:150 return None151 tmp_root = (152 Path(os.environ.get("GRADIO_TEMP_DIR", "/tmp"))153 / "model3d-cache"154 / session_hash155 )156 tmp_root.mkdir(parents=True, exist_ok=True)157 158 # rolling cleanup for same kind159 prev = _prev_temp.get(session_hash, {})160 old = prev.get(kind)161 if old and old.exists():162 old.unlink()163 164 ext = Path(src_path).suffix or ".glb"165 dst = tmp_root / f"{kind}-{uuid.uuid4().hex}{ext}"166 shutil.copy2(src_path, dst)167 168 prev[kind] = dst169 _prev_temp[session_hash] = prev170 return str(dst)171 172 173def _bounding_radius(mesh: trimesh.Trimesh | trimesh.Scene) -> float:174 """Radius of the mesh/scene bounding sphere (half the diagonal)."""175 lo, hi = mesh.bounds176 return float(np.linalg.norm(np.asarray(hi) - np.asarray(lo)) / 2)177 178 179def _camera_position(radius: float | None) -> tuple:180 """Initial camera `(alpha, beta, radius)`; only push the distance."""181 if not radius or radius <= 0:182 return (None, None, None)183 return (None, None, CAMERA_ZOOM * radius)184 185 186def _visual_radius(visual_path: str | None) -> float | None:187 if not visual_path:188 return None189 return _bounding_radius(trimesh.load(visual_path, process=False))190 191 192def _pastel_color(i: int) -> list[int]:193 """Distinct pastel RGBA for the i-th convex piece (golden-ratio hues)."""194 h = (i * 0.6180339887498949) % 1.0195 r, g, b = colorsys.hsv_to_rgb(h, 0.45, 0.98)196 return [int(r * 255), int(g * 255), int(b * 255), 255]197 198 199def _colored_collision_path(200 collision_path: str | None, session_hash: str201) -> tuple[str | None, float | None]:202 """Export the collision mesh as a GLB with one color per convex piece.203 204 The convex-decomposition pieces are stored as separate objects inside the205 single collision ``.obj``; they are recovered via connected-component split206 and each is tinted a distinct pastel color for visualization.207 208 Returns ``(glb_path, bounding_radius)`` for camera framing.209 """210 if not collision_path:211 return None, None212 213 tmp_root = (214 Path(os.environ.get("GRADIO_TEMP_DIR", "/tmp"))215 / "model3d-cache"216 / session_hash217 )218 tmp_root.mkdir(parents=True, exist_ok=True)219 220 # rolling cleanup for same kind221 prev = _prev_temp.get(session_hash, {})222 old = prev.get("collision")223 if old and old.exists():224 old.unlink()225 226 loaded = trimesh.load(collision_path, process=False)227 if isinstance(loaded, trimesh.Scene):228 parts = list(loaded.geometry.values())229 else:230 parts = loaded.split(only_watertight=False)231 if len(parts) == 0:232 parts = [loaded]233 234 scene = trimesh.Scene()235 for i, part in enumerate(parts):236 part.visual = trimesh.visual.ColorVisuals(237 mesh=part, vertex_colors=_pastel_color(i)238 )239 scene.add_geometry(part)240 241 dst = tmp_root / f"collision-{uuid.uuid4().hex}.glb"242 scene.export(str(dst))243 244 prev["collision"] = dst245 _prev_temp[session_hash] = prev246 return str(dst), _bounding_radius(scene)247 248 249# --- Helper Functions (data filtering) ---250def get_primary_categories():251 return sorted(df["primary_category"].dropna().unique())252 253 254def get_secondary_categories(primary):255 if not primary:256 return []257 return sorted(258 df[df["primary_category"] == primary]["secondary_category"]259 .dropna()260 .unique()261 )262 263 264def get_categories(primary, secondary):265 if not primary or not secondary:266 return []267 return sorted(268 df[269 (df["primary_category"] == primary)270 & (df["secondary_category"] == secondary)271 ]["category"]272 .dropna()273 .unique()274 )275 276 277def get_assets(primary, secondary, category):278 if not primary or not secondary:279 return [], gr.update(interactive=False), pd.DataFrame()280 281 subset = df[282 (df["primary_category"] == primary)283 & (df["secondary_category"] == secondary)284 ]285 if category:286 subset = subset[subset["category"] == category]287 288 items = []289 for row in subset.itertuples():290 asset_dir = os.path.join(DATA_ROOT, row.asset_dir)291 video_path = None292 if pd.notna(asset_dir) and os.path.exists(asset_dir):293 for f in os.listdir(asset_dir):294 if f.lower().endswith(".mp4"):295 video_path = os.path.join(asset_dir, f)296 break297 items.append(298 video_path299 if video_path300 else "https://dummyimage.com/512x512/cccccc/000000&text=No+Preview"301 )302 303 return items, gr.update(interactive=True), subset304 305 306def search_assets(query: str, top_k: int):307 if not GPT_AVAILABLE or not query:308 gr.Warning(309 "GPT client is not available or query is empty. Cannot perform search."310 )311 return [], gr.update(interactive=False), pd.DataFrame()312 313 gr.Info(f"Searching for assets matching: '{query}'...")314 315 keywords = query.split()316 keyword_filter = pd.Series([False] * len(df), index=df.index)317 for keyword in keywords:318 keyword_filter |= df['description'].str.contains(319 keyword, case=False, na=False320 )321 322 candidates = df[keyword_filter]323 324 if len(candidates) > 100:325 candidates = candidates.head(100)326 327 if candidates.empty:328 gr.Warning("No assets found matching the keywords.")329 return [], gr.update(interactive=True), pd.DataFrame()330 331 try:332 descriptions = [333 f"{idx}: {desc}" for idx, desc in candidates['description'].items()334 ]335 descriptions_text = "\n".join(descriptions)336 337 prompt = f"""338 A user is searching for 3D assets with the query: "{query}".339 Below is a list of available assets, each with an ID and a description.340 Please evaluate how well each asset description matches the user's query and rate them on a scale from 0 to 10, where 10 is a perfect match.341 342 Your task is to return a list of the top {top_k} asset IDs, sorted from the most relevant to the least relevant.343 The output format must be a simple comma-separated list of IDs, for example: "123,45,678". Do not add any other text.344 345 Asset Descriptions:346 {descriptions_text}347 348 User Query: "{query}"349 350 Top {top_k} sorted asset IDs:351 """352 response = gpt_client.query(prompt)353 sorted_ids_str = response.strip().split(',')354 sorted_ids = [355 int(id_str.strip())356 for id_str in sorted_ids_str357 if id_str.strip().isdigit()358 ]359 top_assets = df.loc[sorted_ids].head(top_k)360 except Exception as e:361 gr.Error(f"An error occurred while using GPT for ranking: {e}")362 top_assets = candidates.head(top_k)363 364 items = []365 for row in top_assets.itertuples():366 asset_dir = os.path.join(DATA_ROOT, row.asset_dir)367 video_path = None368 if pd.notna(row.asset_dir) and os.path.exists(asset_dir):369 for f in os.listdir(asset_dir):370 if f.lower().endswith(".mp4"):371 video_path = os.path.join(asset_dir, f)372 break373 items.append(374 video_path375 if video_path376 else "https://dummyimage.com/512x512/cccccc/000000&text=No+Preview"377 )378 379 gr.Info(f"Found {len(items)} assets.")380 return items, gr.update(interactive=True), top_assets381 382 383def _extract_mesh_paths(row) -> tuple[str | None, str | None, str]:384 desc = row["description"]385 urdf_path = os.path.join(DATA_ROOT, row["urdf_path"])386 asset_dir = os.path.join(DATA_ROOT, row["asset_dir"])387 visual_mesh_path = None388 collision_mesh_path = None389 390 if pd.notna(urdf_path) and os.path.exists(urdf_path):391 try:392 tree = ET.parse(urdf_path)393 root = tree.getroot()394 395 visual_mesh_element = root.find('.//visual/geometry/mesh')396 if visual_mesh_element is not None:397 visual_mesh_filename = visual_mesh_element.get('filename')398 if visual_mesh_filename:399 glb_filename = (400 os.path.splitext(visual_mesh_filename)[0] + ".glb"401 )402 potential_path = os.path.join(asset_dir, glb_filename)403 if os.path.exists(potential_path):404 visual_mesh_path = potential_path405 406 collision_mesh_element = root.find('.//collision/geometry/mesh')407 if collision_mesh_element is not None:408 collision_mesh_filename = collision_mesh_element.get(409 'filename'410 )411 if collision_mesh_filename:412 potential_collision_path = os.path.join(413 asset_dir, collision_mesh_filename414 )415 if os.path.exists(potential_collision_path):416 collision_mesh_path = potential_collision_path417 418 except ET.ParseError:419 desc = f"Error: Failed to parse URDF at {urdf_path}. {desc}"420 except Exception as e:421 desc = f"An error occurred while processing URDF: {str(e)}. {desc}"422 423 return visual_mesh_path, collision_mesh_path, desc424 425 426def show_asset_from_gallery(427 evt: gr.SelectData,428 primary: str,429 secondary: str,430 category: str,431 search_query: str,432 gallery_df: pd.DataFrame,433):434 """Parse the selected asset and return raw paths + metadata."""435 index = evt.index436 437 if search_query and gallery_df is not None and not gallery_df.empty:438 subset = gallery_df439 else:440 if not primary or not secondary:441 return (442 None, # visual_path443 None, # collision_path444 "Error: Primary or secondary category not selected.",445 None, # asset_dir446 None, # urdf_path447 "N/A",448 "N/A",449 "N/A",450 "N/A",451 )452 453 subset = df[454 (df["primary_category"] == primary)455 & (df["secondary_category"] == secondary)456 ]457 if category:458 subset = subset[subset["category"] == category]459 460 if subset.empty or index >= len(subset):461 return (462 None,463 None,464 "Error: Selection index is out of bounds or data is missing.",465 None,466 None,467 "N/A",468 "N/A",469 "N/A",470 "N/A",471 )472 473 row = subset.iloc[index]474 475 # In hf_remote mode, pull only what the two mesh viewers need: the visual476 # GLB + collision OBJ (mesh/), plus the tiny URDF used to locate them and477 # show metadata. USD/MJCF/USD-textures are fetched only on button click.478 gr.Info("⏳ Loading 3D model, please wait...")479 rel_dir = row["asset_dir"]480 urdf_name = os.path.basename(row["urdf_path"])481 ensure_hf_files(482 [483 f"{rel_dir}/{urdf_name}",484 f"{rel_dir}/mesh/**",485 ]486 )487 488 visual_path, collision_path, desc = _extract_mesh_paths(row)489 490 urdf_path = os.path.join(DATA_ROOT, row["urdf_path"])491 asset_dir = os.path.join(DATA_ROOT, row["asset_dir"])492 493 # read extra info494 est_type_text = "N/A"495 est_height_text = "N/A"496 est_mass_text = "N/A"497 est_mu_text = "N/A"498 499 if pd.notna(urdf_path) and os.path.exists(urdf_path):500 try:501 tree = ET.parse(urdf_path)502 root = tree.getroot()503 category_elem = root.find('.//extra_info/category')504 if category_elem is not None and category_elem.text:505 est_type_text = category_elem.text.strip()506 height_elem = root.find('.//extra_info/real_height')507 if height_elem is not None and height_elem.text:508 est_height_text = height_elem.text.strip()509 mass_elem = root.find('.//extra_info/min_mass')510 if mass_elem is not None and mass_elem.text:511 est_mass_text = mass_elem.text.strip()512 mu_elem = root.find('.//collision/gazebo/mu2')513 if mu_elem is not None and mu_elem.text:514 est_mu_text = mu_elem.text.strip()515 except Exception:516 pass517 518 return (519 visual_path,520 collision_path,521 desc,522 asset_dir,523 urdf_path,524 est_type_text,525 est_height_text,526 est_mass_text,527 est_mu_text,528 )529 530 531def render_meshes(532 visual_path: str | None,533 collision_path: str | None,534 switch_viewer: bool,535 req: gr.Request,536):537 session_hash = getattr(req, "session_hash", "default")538 539 if switch_viewer:540 yield (541 gr.update(value=None),542 gr.update(value=None, visible=False),543 gr.update(value=None, visible=True),544 True,545 )546 else:547 yield (548 gr.update(value=None),549 gr.update(value=None, visible=True),550 gr.update(value=None, visible=False),551 True,552 )553 554 visual_unique = (555 _unique_path(visual_path, session_hash, "visual")556 if visual_path557 else None558 )559 visual_cam = _camera_position(_visual_radius(visual_path))560 561 if collision_path:562 collision_unique, collision_r = _colored_collision_path(563 collision_path, session_hash564 )565 else:566 collision_unique, collision_r = None, None567 collision_cam = _camera_position(collision_r)568 569 if visual_unique or collision_unique:570 gr.Info("✅ 3D model loaded.")571 572 if switch_viewer:573 yield (574 gr.update(value=visual_unique, camera_position=visual_cam),575 gr.update(value=None, visible=False),576 gr.update(577 value=collision_unique,578 visible=True,579 camera_position=collision_cam,580 ),581 False,582 )583 else:584 yield (585 gr.update(value=visual_unique, camera_position=visual_cam),586 gr.update(587 value=collision_unique,588 visible=True,589 camera_position=collision_cam,590 ),591 gr.update(value=None, visible=False),592 True,593 )594 595 596def _rel_dir(asset_dir: str) -> str:597 """Repo-relative asset dir (path under DATA_ROOT), for HF glob patterns."""598 return os.path.relpath(asset_dir, DATA_ROOT)599 600 601def _find_urdf_stem(asset_dir: str) -> str | None:602 for f in os.listdir(asset_dir):603 if f.lower().endswith(".urdf"):604 return os.path.splitext(f)[0]605 return None606 607 608def _zip_items(609 zip_path: str,610 items: list[tuple[str, str]],611 exclude_suffixes: tuple[str, ...] = (),612) -> str:613 """Write files/dirs into a zip.614 615 Args:616 zip_path: Output archive path.617 items: ``(src_abspath, arcname)`` pairs. Directories are walked and618 stored under ``arcname``.619 exclude_suffixes: Lowercase filename suffixes to skip when walking620 directories, e.g. ``(".glb",)``.621 """622 with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:623 for src, arcname in items:624 if os.path.isdir(src):625 for root_d, _, files in os.walk(src):626 for fn in files:627 if exclude_suffixes and fn.lower().endswith(628 exclude_suffixes629 ):630 continue631 fp = os.path.join(root_d, fn)632 rel = os.path.relpath(fp, src)633 zf.write(fp, os.path.join(arcname, rel))634 elif os.path.isfile(src):635 zf.write(src, arcname)636 return zip_path637 638 639def _fmt_zip_path(640 asset_dir: str, fmt: str, req: gr.Request641) -> tuple[str, str, str]:642 """Build the per-format output zip path named ``<uid>_<fmt>.zip``.643 644 Returns ``(zip_path, uid, zip_name)``.645 """646 uid = os.path.basename(os.path.normpath(asset_dir))647 user_dir = os.path.join(TMP_DIR, str(req.session_hash))648 os.makedirs(user_dir, exist_ok=True)649 zip_name = f"{uid}_{fmt}.zip"650 return os.path.join(user_dir, zip_name), uid, zip_name651 652 653def download_urdf(asset_dir: str, req: gr.Request) -> str | None:654 """Package the ``.urdf`` file(s) and the ``mesh/`` folder (minus ``.glb``).655 656 The visual/collision meshes referenced by the URDF live in ``mesh/`` as657 ``.obj`` + material files; the ``.glb`` (a separate viewer asset) and all658 other formats (USD/MJCF/video/...) are intentionally excluded.659 """660 if not asset_dir or not os.path.isdir(asset_dir):661 gr.Warning("Please select an asset first.")662 return None663 gr.Info("⏳ Preparing URDF asset for download, please wait...")664 rel = _rel_dir(asset_dir)665 ensure_hf_files([f"{rel}/*.urdf", f"{rel}/mesh/**"])666 667 items: list[tuple[str, str]] = [668 (os.path.join(asset_dir, f), f)669 for f in os.listdir(asset_dir)670 if f.lower().endswith(".urdf")671 ]672 mesh_dir = os.path.join(asset_dir, "mesh")673 if os.path.isdir(mesh_dir):674 items.append((mesh_dir, "mesh"))675 676 zip_path, uid, zip_name = _fmt_zip_path(asset_dir, "urdf", req)677 _zip_items(zip_path, items, exclude_suffixes=(".glb",))678 gr.Info(f"✅ {zip_name} is ready.")679 return zip_path680 681 682def download_usd(asset_dir: str, req: gr.Request) -> str | None:683 """Package the ``.usd`` file together with its ``textures/`` folder."""684 if not asset_dir or not os.path.isdir(asset_dir):685 gr.Warning("Please select an asset first.")686 return None687 688 gr.Info("⏳ Preparing USD asset for download, please wait...")689 rel = _rel_dir(asset_dir)690 stem = _find_urdf_stem(asset_dir)691 if stem is not None:692 ensure_hf_files([f"{rel}/{stem}.usd", f"{rel}/textures/**"])693 else:694 ensure_hf_files([f"{rel}/*.usd", f"{rel}/textures/**"])695 696 items: list[tuple[str, str]] = []697 for f in os.listdir(asset_dir):698 if f.lower().endswith(".usd"):699 items.append((os.path.join(asset_dir, f), f))700 tex_dir = os.path.join(asset_dir, "textures")701 if os.path.isdir(tex_dir):702 items.append((tex_dir, "textures"))703 704 if not any(s.lower().endswith(".usd") for s, _ in items):705 gr.Warning("No USD file available for this asset.")706 return None707 708 zip_path, uid, zip_name = _fmt_zip_path(asset_dir, "usd", req)709 _zip_items(zip_path, items)710 gr.Info(f"✅ {zip_name} is ready.")711 return zip_path712 713 714def download_mjcf(asset_dir: str, req: gr.Request) -> str | None:715 """Package the ``mjcf/`` folder."""716 if not asset_dir or not os.path.isdir(asset_dir):717 gr.Warning("Please select an asset first.")718 return None719 720 gr.Info("⏳ Preparing MJCF asset for download, please wait...")721 ensure_hf_files([f"{_rel_dir(asset_dir)}/mjcf/**"])722 mjcf_dir = os.path.join(asset_dir, "mjcf")723 if not os.path.isdir(mjcf_dir):724 gr.Warning("No MJCF folder available for this asset.")725 return None726 727 zip_path, uid, zip_name = _fmt_zip_path(asset_dir, "mjcf", req)728 _zip_items(zip_path, [(mjcf_dir, "mjcf")])729 gr.Info(f"✅ {zip_name} is ready.")730 return zip_path731 732 733def download_affordance(asset_dir: str, req: gr.Request) -> str | None:734 """Package the ``affordance/`` folder."""735 if not asset_dir or not os.path.isdir(asset_dir):736 gr.Warning("Please select an asset first.")737 return None738 739 gr.Info("⏳ Preparing affordance asset for download, please wait...")740 ensure_hf_files([f"{_rel_dir(asset_dir)}/affordance/**"])741 affordance_dir = os.path.join(asset_dir, "affordance")742 if not os.path.isdir(affordance_dir):743 gr.Warning("No affordance folder available for this asset.")744 return None745 746 zip_path, uid, zip_name = _fmt_zip_path(asset_dir, "affordance", req)747 _zip_items(zip_path, [(affordance_dir, "affordance")])748 gr.Info(f"✅ {zip_name} is ready.")749 return zip_path750 751 752# Download buttons keyed by format; order matches the outputs list used when753# locking/unlocking them together.754_DL_KEYS = ("urdf", "usd", "mjcf", "affordance")755_DL_LABELS = {756 "urdf": "⬇️ Download URDF",757 "usd": "⬇️ Download USD",758 "mjcf": "⬇️ Download MJCF",759 "affordance": "⬇️ Affordance",760}761 762 763def _lock_downloads(active: str) -> tuple:764 """Disable all download buttons; show a persistent hint on the clicked one.765 766 Keeps the user from spamming clicks while the (possibly slow) zip is being767 prepared, since a toast alone neither blocks clicks nor stays put.768 """769 return tuple(770 gr.update(interactive=False, label=f"⏳ Preparing {k.upper()}...")771 if k == active772 else gr.update(interactive=False)773 for k in _DL_KEYS774 )775 776 777def _unlock_downloads() -> tuple:778 """Re-enable buttons and restore labels; keep the freshly built file value."""779 return tuple(780 gr.update(interactive=True, label=_DL_LABELS[k]) for k in _DL_KEYS781 )782 783 784def _reset_downloads() -> tuple:785 """Enable + restore labels + clear stale file value (used on asset switch)."""786 return tuple(787 gr.update(interactive=True, label=_DL_LABELS[k], value=None)788 for k in _DL_KEYS789 )790 791 792def start_session(req: gr.Request) -> None:793 user_dir = os.path.join(TMP_DIR, str(req.session_hash))794 os.makedirs(user_dir, exist_ok=True)795 796 797def end_session(req: gr.Request) -> None:798 user_dir = os.path.join(TMP_DIR, str(req.session_hash))799 if os.path.exists(user_dir):800 shutil.rmtree(user_dir)801 802 803# --- UI ---804with gr.Blocks(805 theme=custom_theme,806 css=css,807 title="3D Asset Library",808) as demo:809 # gr.HTML(lighting_css, visible=False)810 gr.Markdown(811 """812 ## 🏛️ ***EmbodiedGen***: 3D Asset Gallery Explorer813 814 **🔖 Version**: {VERSION}815 <p style="display: flex; gap: 10px; flex-wrap: nowrap;">816 <a href="https://horizonrobotics.github.io/EmbodiedGen">817 <img alt="📖 Documentation" src="https://img.shields.io/badge/📖-Documentation-blue">818 </a>819 <a href="https://arxiv.org/abs/2607.07459">820 <img alt="📄 arXiv" src="https://img.shields.io/badge/📄-arXiv-b31b1b">821 </a>822 <a href="https://github.com/HorizonRobotics/EmbodiedGen">823 <img alt="💻 GitHub" src="https://img.shields.io/badge/GitHub-000000?logo=github">824 </a>825 <a href="https://youtu.be/MIkJJSVM8L4">826 <img alt="🎥 Video" src="https://img.shields.io/badge/🎥-Video-red">827 </a>828 </p>829 830 Browse and visualize the EmbodiedGen 3D asset database. Select categories to filter and click on a preview to load the model.831 832 """.format(VERSION=VERSION),833 elem_classes=["header"],834 )835 836 primary_list = get_primary_categories()837 primary_val = primary_list[0] if primary_list else None838 secondary_list = get_secondary_categories(primary_val)839 secondary_val = secondary_list[0] if secondary_list else None840 category_list = get_categories(primary_val, secondary_val)841 category_val = category_list[0] if category_list else None842 asset_folder = gr.State(value=None)843 gallery_df_state = gr.State()844 845 switch_viewer_state = gr.State(value=False)846 847 with gr.Row(equal_height=False):848 with gr.Column(scale=1, min_width=350):849 with gr.Group():850 gr.Markdown("### Search Asset with Descriptions")851 search_box = gr.Textbox(852 label="🔎 Enter your search query",853 placeholder="e.g., 'a red chair with four legs'",854 interactive=GPT_AVAILABLE,855 )856 top_k_slider = gr.Slider(857 minimum=1,858 maximum=50,859 value=10,860 step=1,861 label="Number of results",862 interactive=GPT_AVAILABLE,863 )864 search_button = gr.Button(865 "Search", variant="primary", interactive=GPT_AVAILABLE866 )867 if not GPT_AVAILABLE:868 gr.Markdown(869 "<p style='color: #ff4b4b;'>⚠️ GPT client not available. Search is disabled.</p>"870 )871 872 with gr.Group():873 gr.Markdown("### Select Asset Category")874 primary = gr.Dropdown(875 choices=primary_list,876 value=primary_val,877 label="🗂️ Primary Category",878 )879 secondary = gr.Dropdown(880 choices=secondary_list,881 value=secondary_val,882 label="📂 Secondary Category",883 )884 category = gr.Dropdown(885 choices=category_list,886 value=category_val,887 label="🏷️ Asset Category",888 )889 890 with gr.Group():891 initial_assets, _, initial_df = get_assets(892 primary_val, secondary_val, category_val893 )894 gallery = gr.Gallery(895 value=initial_assets,896 label="🖼️ Asset Previews",897 columns=3,898 height="auto",899 allow_preview=True,900 elem_id="asset-gallery",901 interactive=bool(category_val),902 )903 904 with gr.Column(scale=2, min_width=500):905 with gr.Group():906 with gr.Tabs():907 with gr.TabItem("Visual Mesh") as t1:908 viewer = gr.Model3D(909 label="🧊 3D Model Viewer",910 height=380,911 clear_color=[0.95, 0.95, 0.95],912 elem_id="visual_mesh",913 )914 with gr.TabItem("Collision Mesh") as t2:915 collision_viewer_a = gr.Model3D(916 label="🧊 Collision Mesh",917 height=380,918 clear_color=[0.95, 0.95, 0.95],919 elem_id="collision_mesh_a",920 visible=True,921 )922 collision_viewer_b = gr.Model3D(923 label="🧊 Collision Mesh",924 height=380,925 clear_color=[0.95, 0.95, 0.95],926 elem_id="collision_mesh_b",927 visible=False,928 )929 930 t1.select(931 fn=lambda: None,932 js="() => { window.dispatchEvent(new Event('resize')); }",933 )934 t2.select(935 fn=lambda: None,936 js="() => { window.dispatchEvent(new Event('resize')); }",937 )938 939 with gr.Row():940 est_type_text = gr.Textbox(941 label="Asset category", interactive=False942 )943 est_height_text = gr.Textbox(944 label="Real height(.m)", interactive=False945 )946 est_mass_text = gr.Textbox(947 label="Mass(.kg)", interactive=False948 )949 est_mu_text = gr.Textbox(950 label="Friction coefficient", interactive=False951 )952 with gr.Row():953 desc_box = gr.Textbox(954 label="📝 Asset Description", interactive=False955 )956 with gr.Accordion(label="Asset Details", open=False):957 urdf_file = gr.Textbox(958 label="URDF File Path", interactive=False, lines=2959 )960 with gr.Row(elem_id="download-row"):961 # DownloadButtons that build their zip into their own value962 # on click; a chained JS handler then triggers the browser963 # download from that value (one-click). If the JS ever964 # fails, the value is still populated, so a second manual965 # click downloads it via the button's native behavior.966 urdf_dl_btn = gr.DownloadButton(967 label="⬇️ Download URDF",968 variant="primary",969 interactive=False,970 )971 usd_dl_btn = gr.DownloadButton(972 label="⬇️ Download USD",973 variant="primary",974 interactive=False,975 )976 mjcf_dl_btn = gr.DownloadButton(977 label="⬇️ Download MJCF",978 variant="primary",979 interactive=False,980 )981 affordance_dl_btn = gr.DownloadButton(982 label="⬇️ Affordance",983 variant="primary",984 interactive=False,985 )986 987 search_button.click(988 fn=search_assets,989 inputs=[search_box, top_k_slider],990 outputs=[gallery, gallery, gallery_df_state],991 )992 search_box.submit(993 fn=search_assets,994 inputs=[search_box, top_k_slider],995 outputs=[gallery, gallery, gallery_df_state],996 )997 998 def update_on_primary_change(p):999 s_choices = get_secondary_categories(p)1000 initial_assets, gallery_update, initial_df = get_assets(p, None, None)1001 return (1002 gr.update(choices=s_choices, value=None),1003 gr.update(choices=[], value=None),1004 initial_assets,1005 gallery_update,1006 initial_df,1007 )1008 1009 def update_on_secondary_change(p, s):1010 c_choices = get_categories(p, s)1011 asset_previews, gallery_update, gallery_df = get_assets(p, s, None)1012 return (1013 gr.update(choices=c_choices, value=None),1014 asset_previews,1015 gallery_update,1016 gallery_df,1017 )1018 1019 def update_assets(p, s, c):1020 asset_previews, gallery_update, gallery_df = get_assets(p, s, c)1021 return asset_previews, gallery_update, gallery_df1022 1023 primary.change(1024 fn=update_on_primary_change,1025 inputs=[primary],1026 outputs=[secondary, category, gallery, gallery, gallery_df_state],1027 )1028 secondary.change(1029 fn=update_on_secondary_change,1030 inputs=[primary, secondary],1031 outputs=[category, gallery, gallery, gallery_df_state],1032 )1033 category.change(1034 fn=update_assets,1035 inputs=[primary, secondary, category],1036 outputs=[gallery, gallery, gallery_df_state],1037 )1038 1039 gallery.select(1040 fn=show_asset_from_gallery,1041 inputs=[primary, secondary, category, search_box, gallery_df_state],1042 concurrency_id=HF_IO_CONCURRENCY_ID,1043 concurrency_limit=HF_IO_CONCURRENCY_LIMIT,1044 outputs=[1045 (visual_path_state := gr.State()),1046 (collision_path_state := gr.State()),1047 desc_box,1048 asset_folder,1049 urdf_file,1050 est_type_text,1051 est_height_text,1052 est_mass_text,1053 est_mu_text,1054 ],1055 ).then(1056 fn=render_meshes,1057 inputs=[visual_path_state, collision_path_state, switch_viewer_state],1058 outputs=[1059 viewer,1060 collision_viewer_a,1061 collision_viewer_b,1062 switch_viewer_state,1063 ],1064 ).success(1065 fn=_reset_downloads,1066 outputs=[urdf_dl_btn, usd_dl_btn, mjcf_dl_btn, affordance_dl_btn],1067 )1068 1069 # After the zip is built into the button's own value, pass that file value1070 # straight into the JS handler and trigger the browser download from its1071 # URL. Reading the value via `inputs` avoids DOM-timing/selector issues.1072 download_js = """1073 (f) => {1074 if (f && f.url) {1075 const a = document.createElement('a');1076 a.href = f.url;1077 a.download = (f.orig_name || 'asset.zip').split('/').pop();1078 document.body.appendChild(a);1079 a.click();1080 document.body.removeChild(a);1081 }1082 return [];1083 }1084 """1085 1086 dl_btns = [urdf_dl_btn, usd_dl_btn, mjcf_dl_btn, affordance_dl_btn]1087 1088 # Each flow: lock all buttons (prevent spam clicks) -> build the zip ->1089 # JS-trigger the browser download -> unlock. Selecting another asset1090 # resets the buttons independently (see `.success` above).1091 urdf_dl_btn.click(1092 fn=lambda: _lock_downloads("urdf"), outputs=dl_btns, queue=False1093 ).then(1094 fn=download_urdf,1095 inputs=[asset_folder],1096 outputs=[urdf_dl_btn],1097 concurrency_id=HF_IO_CONCURRENCY_ID,1098 concurrency_limit=HF_IO_CONCURRENCY_LIMIT,1099 ).then(fn=lambda *a: None, inputs=[urdf_dl_btn], js=download_js).then(1100 fn=_unlock_downloads, outputs=dl_btns1101 )1102 1103 usd_dl_btn.click(1104 fn=lambda: _lock_downloads("usd"), outputs=dl_btns, queue=False1105 ).then(1106 fn=download_usd,1107 inputs=[asset_folder],1108 outputs=[usd_dl_btn],1109 concurrency_id=HF_IO_CONCURRENCY_ID,1110 concurrency_limit=HF_IO_CONCURRENCY_LIMIT,1111 ).then(fn=lambda *a: None, inputs=[usd_dl_btn], js=download_js).then(1112 fn=_unlock_downloads, outputs=dl_btns1113 )1114 1115 mjcf_dl_btn.click(1116 fn=lambda: _lock_downloads("mjcf"), outputs=dl_btns, queue=False1117 ).then(1118 fn=download_mjcf,1119 inputs=[asset_folder],1120 outputs=[mjcf_dl_btn],1121 concurrency_id=HF_IO_CONCURRENCY_ID,1122 concurrency_limit=HF_IO_CONCURRENCY_LIMIT,1123 ).then(fn=lambda *a: None, inputs=[mjcf_dl_btn], js=download_js).then(1124 fn=_unlock_downloads, outputs=dl_btns1125 )1126 1127 affordance_dl_btn.click(1128 fn=lambda: _lock_downloads("affordance"),1129 outputs=dl_btns,1130 queue=False,1131 ).then(1132 fn=download_affordance,1133 inputs=[asset_folder],1134 outputs=[affordance_dl_btn],1135 concurrency_id=HF_IO_CONCURRENCY_ID,1136 concurrency_limit=HF_IO_CONCURRENCY_LIMIT,1137 ).then(1138 fn=lambda *a: None, inputs=[affordance_dl_btn], js=download_js1139 ).then(fn=_unlock_downloads, outputs=dl_btns)1140 1141 demo.load(start_session)1142 demo.unload(end_session)1143 1144 1145if __name__ == "__main__":1146 # Serve gallery videos / meshes that live under DATA_ROOT (outside cwd).1147 demo.launch()1148