update

app.py

@@ -23,13 +23,17 @@ gradio_tmp_dir = os.path.join(
 os.makedirs(gradio_tmp_dir, exist_ok=True)
 os.environ["GRADIO_TEMP_DIR"] = gradio_tmp_dir
 
+import colorsys
 import shutil
 import uuid
 import xml.etree.ElementTree as ET
+import zipfile
 from pathlib import Path
 
 import gradio as gr
+import numpy as np
 import pandas as pd
+import trimesh
 from app_style import custom_theme, lighting_css
 from embodied_gen.utils.tags import VERSION
 
@@ -49,9 +53,16 @@ except Exception as e:
 # --- Configuration & Data Loading ---
 RUNNING_MODE = "hf_remote"  # local or hf_remote
 CSV_FILE = "dataset_index.csv"
+HF_REPO_ID = "HorizonRobotics/EmbodiedGenData"
+HF_LOCAL_DIR = "EmbodiedGenData"
+# Camera distance as a multiple of each mesh's bounding radius. Meshes are in
+# real-world meters (sizes vary widely), so the camera radius must scale per
+# asset. Larger -> camera further away -> object appears smaller on load.
+CAMERA_ZOOM = 3.2
 
 # Compatible with huggingface space zero GPU
 import spaces
+from huggingface_hub import snapshot_download
 
 
 @spaces.GPU
@@ -62,23 +73,47 @@ def fake_gpu_init():
 fake_gpu_init()
 
 if RUNNING_MODE == "local":
-    DATA_ROOT = "/horizon-bucket/robot_lab/datasets/embodiedgen/assets"
+    DATA_ROOT = "/horizon-bucket/robot_lab/datasets/embodiedgen/assets_v2"
 elif RUNNING_MODE == "hf_remote":
-    from huggingface_hub import snapshot_download
-
+    # Only fetch the index and preview videos up front; per-asset mesh/urdf/
+    # usd/mjcf files are pulled lazily on demand (see `ensure_hf_files`).
     snapshot_download(
-        repo_id="HorizonRobotics/EmbodiedGenData",
+        repo_id=HF_REPO_ID,
         repo_type="dataset",
-        allow_patterns=f"dataset/**",
-        local_dir="EmbodiedGenData",
+        allow_patterns=[
+            f"dataset/{CSV_FILE}",
+            "dataset/**/*.mp4",
+        ],
+        local_dir=HF_LOCAL_DIR,
         local_dir_use_symlinks=False,
     )
-    DATA_ROOT = "EmbodiedGenData/dataset"
+    DATA_ROOT = os.path.join(HF_LOCAL_DIR, "dataset")
 else:
     raise ValueError(
         f"Unknown RUNNING_MODE: {RUNNING_MODE}, must be 'local' or 'hf_remote'."
     )
 
+
+def ensure_hf_files(rel_patterns: list[str]) -> None:
+    """Lazily download files for one asset in hf_remote mode.
+
+    Args:
+        rel_patterns: Glob patterns relative to ``DATA_ROOT`` (the repo's
+            ``dataset/`` folder), e.g. ``"cat/sub/uid/mesh/**"``.
+
+    No-op when running locally, where all files already exist on disk.
+    """
+    if RUNNING_MODE != "hf_remote":
+        return
+    snapshot_download(
+        repo_id=HF_REPO_ID,
+        repo_type="dataset",
+        allow_patterns=[f"dataset/{p}" for p in rel_patterns],
+        local_dir=HF_LOCAL_DIR,
+        local_dir_use_symlinks=False,
+    )
+
+
 csv_path = os.path.join(DATA_ROOT, CSV_FILE)
 df = pd.read_csv(csv_path)
 TMP_DIR = os.path.join(
@@ -131,6 +166,82 @@ def _unique_path(
     return str(dst)
 
 
+def _bounding_radius(mesh: trimesh.Trimesh | trimesh.Scene) -> float:
+    """Radius of the mesh/scene bounding sphere (half the diagonal)."""
+    lo, hi = mesh.bounds
+    return float(np.linalg.norm(np.asarray(hi) - np.asarray(lo)) / 2)
+
+
+def _camera_position(radius: float | None) -> tuple:
+    """Initial camera `(alpha, beta, radius)`; only push the distance."""
+    if not radius or radius <= 0:
+        return (None, None, None)
+    return (None, None, CAMERA_ZOOM * radius)
+
+
+def _visual_radius(visual_path: str | None) -> float | None:
+    if not visual_path:
+        return None
+    return _bounding_radius(trimesh.load(visual_path, process=False))
+
+
+def _pastel_color(i: int) -> list[int]:
+    """Distinct pastel RGBA for the i-th convex piece (golden-ratio hues)."""
+    h = (i * 0.6180339887498949) % 1.0
+    r, g, b = colorsys.hsv_to_rgb(h, 0.45, 0.98)
+    return [int(r * 255), int(g * 255), int(b * 255), 255]
+
+
+def _colored_collision_path(
+    collision_path: str | None, session_hash: str
+) -> tuple[str | None, float | None]:
+    """Export the collision mesh as a GLB with one color per convex piece.
+
+    The convex-decomposition pieces are stored as separate objects inside the
+    single collision ``.obj``; they are recovered via connected-component split
+    and each is tinted a distinct pastel color for visualization.
+
+    Returns ``(glb_path, bounding_radius)`` for camera framing.
+    """
+    if not collision_path:
+        return None, None
+
+    tmp_root = (
+        Path(os.environ.get("GRADIO_TEMP_DIR", "/tmp"))
+        / "model3d-cache"
+        / session_hash
+    )
+    tmp_root.mkdir(parents=True, exist_ok=True)
+
+    # rolling cleanup for same kind
+    prev = _prev_temp.get(session_hash, {})
+    old = prev.get("collision")
+    if old and old.exists():
+        old.unlink()
+
+    loaded = trimesh.load(collision_path, process=False)
+    if isinstance(loaded, trimesh.Scene):
+        parts = list(loaded.geometry.values())
+    else:
+        parts = loaded.split(only_watertight=False)
+        if len(parts) == 0:
+            parts = [loaded]
+
+    scene = trimesh.Scene()
+    for i, part in enumerate(parts):
+        part.visual = trimesh.visual.ColorVisuals(
+            mesh=part, vertex_colors=_pastel_color(i)
+        )
+        scene.add_geometry(part)
+
+    dst = tmp_root / f"collision-{uuid.uuid4().hex}.glb"
+    scene.export(str(dst))
+
+    prev["collision"] = dst
+    _prev_temp[session_hash] = prev
+    return str(dst), _bounding_radius(scene)
+
+
 # --- Helper Functions (data filtering) ---
 def get_primary_categories():
     return sorted(df["primary_category"].dropna().unique())
@@ -356,6 +467,18 @@ def show_asset_from_gallery(
         )
 
     row = subset.iloc[index]
+
+    # In hf_remote mode, pull only what the viewer needs for this asset.
+    rel_dir = row["asset_dir"]
+    urdf_name = os.path.basename(row["urdf_path"])
+    ensure_hf_files(
+        [
+            f"{rel_dir}/{urdf_name}",
+            f"{rel_dir}/mesh/**",
+            f"{rel_dir}/textures/**",
+        ]
+    )
+
     visual_path, collision_path, desc = _extract_mesh_paths(row)
 
     urdf_path = os.path.join(DATA_ROOT, row["urdf_path"])
@@ -427,41 +550,146 @@ def render_meshes(
         if visual_path
         else None
     )
-    collision_unique = (
-        _unique_path(collision_path, session_hash, "collision")
-        if collision_path
-        else None
-    )
+    visual_cam = _camera_position(_visual_radius(visual_path))
+
+    if collision_path:
+        collision_unique, collision_r = _colored_collision_path(
+            collision_path, session_hash
+        )
+    else:
+        collision_unique, collision_r = None, None
+    collision_cam = _camera_position(collision_r)
 
     if switch_viewer:
         yield (
-            gr.update(value=visual_unique),
+            gr.update(value=visual_unique, camera_position=visual_cam),
             gr.update(value=None, visible=False),
-            gr.update(value=collision_unique, visible=True),
+            gr.update(
+                value=collision_unique,
+                visible=True,
+                camera_position=collision_cam,
+            ),
             False,
         )
     else:
         yield (
-            gr.update(value=visual_unique),
-            gr.update(value=collision_unique, visible=True),
+            gr.update(value=visual_unique, camera_position=visual_cam),
+            gr.update(
+                value=collision_unique,
+                visible=True,
+                camera_position=collision_cam,
+            ),
             gr.update(value=None, visible=False),
             True,
         )
 
 
-def create_asset_zip(asset_dir: str, req: gr.Request):
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    os.makedirs(user_dir, exist_ok=True)
+def _rel_dir(asset_dir: str) -> str:
+    """Repo-relative asset dir (path under DATA_ROOT), for HF glob patterns."""
+    return os.path.relpath(asset_dir, DATA_ROOT)
+
+
+def _find_urdf_stem(asset_dir: str) -> str | None:
+    for f in os.listdir(asset_dir):
+        if f.lower().endswith(".urdf"):
+            return os.path.splitext(f)[0]
+    return None
+
+
+def _zip_items(zip_path: str, items: list[tuple[str, str]]) -> str:
+    """Write files/dirs into a zip.
+
+    Args:
+        zip_path: Output archive path.
+        items: ``(src_abspath, arcname)`` pairs. Directories are walked and
+            stored under ``arcname``.
+    """
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for src, arcname in items:
+            if os.path.isdir(src):
+                for root_d, _, files in os.walk(src):
+                    for fn in files:
+                        fp = os.path.join(root_d, fn)
+                        rel = os.path.relpath(fp, src)
+                        zf.write(fp, os.path.join(arcname, rel))
+            elif os.path.isfile(src):
+                zf.write(src, arcname)
+    return zip_path
+
+
+def _fmt_zip_path(
+    asset_dir: str, fmt: str, req: gr.Request
+) -> tuple[str, str]:
+    """Build the per-format output zip path named ``<uid>.zip``.
+
+    Returns ``(zip_path, uid)``. Each format gets its own subdir so the
+    downloaded filename stays ``<uid>.zip`` without colliding across formats.
+    """
+    uid = os.path.basename(os.path.normpath(asset_dir))
+    fmt_dir = os.path.join(TMP_DIR, str(req.session_hash), fmt)
+    os.makedirs(fmt_dir, exist_ok=True)
+    return os.path.join(fmt_dir, f"{uid}.zip"), uid
+
+
+def download_urdf(asset_dir: str, req: gr.Request) -> str | None:
+    """Package the full asset folder (URDF + all referenced files) as-is."""
+    if not asset_dir or not os.path.isdir(asset_dir):
+        gr.Warning("Please select an asset first.")
+        return None
+    ensure_hf_files([f"{_rel_dir(asset_dir)}/**"])
+    zip_path, uid = _fmt_zip_path(asset_dir, "urdf", req)
+    _zip_items(zip_path, [(asset_dir, "")])
+    gr.Info(f"✅ {uid}.zip (URDF) is ready.")
+    return zip_path
 
-    asset_folder_name = os.path.basename(os.path.normpath(asset_dir))
-    zip_path_base = os.path.join(user_dir, asset_folder_name)
 
-    archive_path = shutil.make_archive(
-        base_name=zip_path_base, format='zip', root_dir=asset_dir
-    )
-    gr.Info(f"✅ {asset_folder_name}.zip is ready and can be downloaded.")
+def download_usd(asset_dir: str, req: gr.Request) -> str | None:
+    """Package the ``.usd`` file together with its ``textures/`` folder."""
+    if not asset_dir or not os.path.isdir(asset_dir):
+        gr.Warning("Please select an asset first.")
+        return None
+
+    rel = _rel_dir(asset_dir)
+    stem = _find_urdf_stem(asset_dir)
+    if stem is not None:
+        ensure_hf_files([f"{rel}/{stem}.usd", f"{rel}/textures/**"])
+    else:
+        ensure_hf_files([f"{rel}/*.usd", f"{rel}/textures/**"])
+
+    items: list[tuple[str, str]] = []
+    for f in os.listdir(asset_dir):
+        if f.lower().endswith(".usd"):
+            items.append((os.path.join(asset_dir, f), f))
+    tex_dir = os.path.join(asset_dir, "textures")
+    if os.path.isdir(tex_dir):
+        items.append((tex_dir, "textures"))
+
+    if not any(s.lower().endswith(".usd") for s, _ in items):
+        gr.Warning("No USD file available for this asset.")
+        return None
 
-    return archive_path
+    zip_path, uid = _fmt_zip_path(asset_dir, "usd", req)
+    _zip_items(zip_path, items)
+    gr.Info(f"✅ {uid}.zip (USD) is ready.")
+    return zip_path
+
+
+def download_mjcf(asset_dir: str, req: gr.Request) -> str | None:
+    """Package the ``mjcf/`` folder."""
+    if not asset_dir or not os.path.isdir(asset_dir):
+        gr.Warning("Please select an asset first.")
+        return None
+
+    ensure_hf_files([f"{_rel_dir(asset_dir)}/mjcf/**"])
+    mjcf_dir = os.path.join(asset_dir, "mjcf")
+    if not os.path.isdir(mjcf_dir):
+        gr.Warning("No MJCF folder available for this asset.")
+        return None
+
+    zip_path, uid = _fmt_zip_path(asset_dir, "mjcf", req)
+    _zip_items(zip_path, [(mjcf_dir, "mjcf")])
+    gr.Info(f"✅ {uid}.zip (MJCF) is ready.")
+    return zip_path
 
 
 def start_session(req: gr.Request) -> None:
@@ -481,7 +709,7 @@ with gr.Blocks(
     css=css,
     title="3D Asset Library",
 ) as demo:
-    # gr.HTML(lighting_css, visible=False)
+    gr.HTML(lighting_css, visible=False)
     gr.Markdown(
         """
         ## 🏛️ ***EmbodiedGen***: 3D Asset Gallery Explorer
@@ -582,21 +810,21 @@ with gr.Blocks(
                     with gr.TabItem("Visual Mesh") as t1:
                         viewer = gr.Model3D(
                             label="🧊 3D Model Viewer",
-                            height=500,
+                            height=380,
                             clear_color=[0.95, 0.95, 0.95],
                             elem_id="visual_mesh",
                         )
                     with gr.TabItem("Collision Mesh") as t2:
                         collision_viewer_a = gr.Model3D(
                             label="🧊 Collision Mesh",
-                            height=500,
+                            height=380,
                             clear_color=[0.95, 0.95, 0.95],
                             elem_id="collision_mesh_a",
                             visible=True,
                         )
                         collision_viewer_b = gr.Model3D(
                             label="🧊 Collision Mesh",
-                            height=500,
+                            height=380,
                             clear_color=[0.95, 0.95, 0.95],
                             elem_id="collision_mesh_b",
                             visible=False,
@@ -633,13 +861,18 @@ with gr.Blocks(
                         label="URDF File Path", interactive=False, lines=2
                     )
                 with gr.Row():
-                    extract_btn = gr.Button(
-                        "📥 Extract Asset",
+                    urdf_dl_btn = gr.DownloadButton(
+                        label="⬇️ Download URDF",
+                        variant="primary",
+                        interactive=False,
+                    )
+                    usd_dl_btn = gr.DownloadButton(
+                        label="⬇️ Download USD",
                         variant="primary",
                         interactive=False,
                     )
-                    download_btn = gr.DownloadButton(
-                        label="⬇️ Download Asset",
+                    mjcf_dl_btn = gr.DownloadButton(
+                        label="⬇️ Download MJCF",
                         variant="primary",
                         interactive=False,
                     )
@@ -720,13 +953,23 @@ with gr.Blocks(
             switch_viewer_state,
         ],
     ).success(
-        lambda: (gr.Button(interactive=True), gr.Button(interactive=False)),
-        outputs=[extract_btn, download_btn],
+        lambda: (
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+        ),
+        outputs=[urdf_dl_btn, usd_dl_btn, mjcf_dl_btn],
     )
 
-    extract_btn.click(
-        fn=create_asset_zip, inputs=[asset_folder], outputs=[download_btn]
-    ).success(fn=lambda: gr.update(interactive=True), outputs=download_btn)
+    urdf_dl_btn.click(
+        fn=download_urdf, inputs=[asset_folder], outputs=[urdf_dl_btn]
+    )
+    usd_dl_btn.click(
+        fn=download_usd, inputs=[asset_folder], outputs=[usd_dl_btn]
+    )
+    mjcf_dl_btn.click(
+        fn=download_mjcf, inputs=[asset_folder], outputs=[mjcf_dl_btn]
+    )
 
     demo.load(start_session)
     demo.unload(end_session)