Spaces:
Running
Running
| """Build the object catalog from NVIDIA SimReady Central assets on the Hugging Face Hub. | |
| For every asset in the dataset that a robot arm could handle (mass <= MAX_KG), writes | |
| objects/<id>.glb the decimated visual geometry, metres, z up | |
| and objects/catalog.js, the list app.js searches: name, class, description, mass, | |
| bounding-box size and a link to the asset's thumbnail and page on the Hub. | |
| Mass, semantic class and description come from each asset's payloads/base.usda; the size | |
| and the geometry come from composing the asset's root USD with usd-core. | |
| uv pip install usd-core huggingface_hub trimesh fast-simplification numpy | |
| python tools/build_objects.py | |
| for f in objects/*.glb; do npx @gltf-transform/cli optimize $f $f --compress meshopt --simplify false \\ | |
| --texture-compress false; done | |
| """ | |
| import json | |
| import re | |
| import sys | |
| from collections import Counter, defaultdict | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| import fast_simplification | |
| import numpy as np | |
| import trimesh | |
| from huggingface_hub import HfApi, hf_hub_download | |
| from pxr import Usd, UsdGeom | |
| REPO = "nvidia/simready-assets" | |
| MAX_KG = 35.0 # heaviest object the heaviest catalog robot can lift | |
| MAX_GEOMETRY = 40e6 # skip huge scans | |
| FACES = 2500 | |
| OUT = Path(__file__).resolve().parent.parent / "objects" | |
| HUB = f"https://huggingface.co/datasets/{REPO}" | |
| def assets(): | |
| by = defaultdict(dict) | |
| for f in HfApi().list_repo_tree(REPO, repo_type="dataset", recursive=True): | |
| if hasattr(f, "size"): | |
| by[f.path.split("/")[0]][f.path] = f.size | |
| out = [] | |
| for name, files in sorted(by.items()): | |
| roots = [p for p in files if p.count("/") == 2 and p.split("/")[1] == "simready_usd" and p.endswith(".usd")] | |
| base = [p for p in files if p.endswith("payloads/base.usda")] | |
| if not roots or not base: | |
| continue | |
| root = roots[0] | |
| thumb = root.replace("simready_usd/", "simready_usd/.thumbs/256x256/") + "_transparent.png" | |
| geo = sum(s for p, s in files.items() if p.endswith(".usd") and "/payloads/" in p) | |
| out.append({"id": name, "root": root, "base": base[0], "thumb": thumb if thumb in files else None, "geo": geo, | |
| "files": [p for p in files if "/simready_usd/" in p and "/Textures/" not in p and "/source_assets/" not in p | |
| and "/.thumbs/" not in p and not p.endswith((".png", ".jpg", ".mdl"))]}) | |
| return out | |
| def metadata(a): | |
| text = Path(hf_hub_download(REPO, a["base"], repo_type="dataset")).read_text(errors="ignore") | |
| mass = sum(float(m) for m in re.findall(r"pxr:usd:physics_mass = ([0-9.eE+-]+)", text)) | |
| if not mass: | |
| m = re.search(r"physics:mass = ([0-9.eE+-]+)", text) | |
| mass = float(m.group(1)) if m else None | |
| cls = re.search(r'semantics:labels:wikidata_class = \["([^"]+)"', text) | |
| doc = re.search(r'omni:simready:documentation = "([^"]*)"', text) | |
| kind = re.search(r'string type = "([^"]+)"', text) | |
| return {"mass": mass, "cls": cls.group(1) if cls else None, "doc": doc.group(1) if doc else "", "type": kind.group(1) if kind else None} | |
| def pretty(name): | |
| words = re.sub(r"_[a-z]\d\d$", "", name).replace("_", " ") | |
| words = re.sub(r"([a-z])([A-Z])", r"\1 \2", words).lower() | |
| return words[:1].upper() + words[1:] | |
| def geometry(a): | |
| for p in a["files"]: | |
| hf_hub_download(REPO, p, repo_type="dataset") | |
| root = hf_hub_download(REPO, a["root"], repo_type="dataset") | |
| stage = Usd.Stage.Open(root) | |
| up = UsdGeom.GetStageUpAxis(stage) | |
| scale = UsdGeom.GetStageMetersPerUnit(stage) or 1.0 | |
| cache = UsdGeom.XformCache() | |
| parts = [] | |
| for prim in Usd.PrimRange.Stage(stage, Usd.TraverseInstanceProxies()): | |
| if not prim.IsA(UsdGeom.Mesh) or UsdGeom.Imageable(prim).ComputePurpose() not in (UsdGeom.Tokens.default_, UsdGeom.Tokens.render): | |
| continue | |
| mesh = UsdGeom.Mesh(prim) | |
| pts, counts, idx = mesh.GetPointsAttr().Get(), mesh.GetFaceVertexCountsAttr().Get(), mesh.GetFaceVertexIndicesAttr().Get() | |
| if not pts or not counts: | |
| continue | |
| pts = np.asarray(pts, float) | |
| M = np.asarray(cache.GetLocalToWorldTransform(prim), float) # row-vector convention | |
| pts = (np.c_[pts, np.ones(len(pts))] @ M)[:, :3] | |
| faces, k = [], 0 | |
| idx = np.asarray(idx) | |
| for c in counts: | |
| faces += [[idx[k], idx[k + i], idx[k + i + 1]] for i in range(1, c - 1)] | |
| k += c | |
| color = mesh.GetDisplayColorAttr().Get() | |
| parts.append((pts, np.asarray(faces, int), tuple(color[0]) if color else None)) | |
| if not parts: | |
| return None | |
| meshes = [] | |
| total = sum(len(f) for _, f, _ in parts) | |
| for pts, faces, color in parts: | |
| target = max(60, int(FACES * len(faces) / total)) | |
| if len(faces) > target: | |
| pts, faces = fast_simplification.simplify(pts.astype(np.float32), faces.astype(np.int32), 1 - target / len(faces)) | |
| m = trimesh.Trimesh(pts * scale, faces, process=True) | |
| rgba = (np.r_[color, 1] if color else np.array([0.78, 0.74, 0.68, 1])) * 255 | |
| m.visual = trimesh.visual.TextureVisuals(material=trimesh.visual.material.PBRMaterial(baseColorFactor=rgba.astype(np.uint8), roughnessFactor=0.7)) | |
| meshes.append(m) | |
| scene = trimesh.Scene(meshes) | |
| if up == "Y": | |
| scene.apply_transform(trimesh.transformations.rotation_matrix(np.pi / 2, [1, 0, 0])) | |
| lo, hi = scene.bounds | |
| scene.apply_translation([-(lo[0] + hi[0]) / 2, -(lo[1] + hi[1]) / 2, -lo[2]]) # centred, resting on z = 0 | |
| return scene | |
| def build(a): | |
| try: | |
| meta = metadata(a) | |
| except Exception as e: | |
| return a["id"], None, f"metadata: {e!r}" | |
| if meta["mass"] is None or meta["mass"] > MAX_KG or a["geo"] > MAX_GEOMETRY: | |
| return a["id"], None, "skipped" | |
| try: | |
| scene = geometry(a) | |
| except Exception as e: | |
| return a["id"], None, f"geometry: {e!r}" | |
| if scene is None: | |
| return a["id"], None, "no mesh" | |
| dims = scene.extents | |
| if max(dims) > 3: # mis-scaled asset | |
| return a["id"], None, f"size {dims}" | |
| (OUT / f"{a['id']}.glb").write_bytes(trimesh.exchange.gltf.export_glb(scene)) | |
| variant = re.search(r"_([a-z]\d\d)$", a["id"]) | |
| name = (meta["cls"][:1].upper() + meta["cls"][1:]) if meta["cls"] else pretty(a["id"]) | |
| item = {"id": a["id"], "name": name + (f" ({variant.group(1)})" if variant else ""), "cls": meta["cls"], "doc": meta["doc"][:160], "kg": round(meta["mass"], 4), | |
| "size": [round(float(x), 4) for x in dims], "thumb": a["thumb"]} | |
| return a["id"], item, "ok" | |
| def dedupe(items): | |
| """Several folders share a class (three "Bottle (a01)"): name those after their folder instead.""" | |
| count = Counter(i["name"] for i in items) | |
| for i in items: | |
| if count[i["name"]] > 1: | |
| variant = re.search(r"_([a-z]\d\d)$", i["id"]) | |
| i["name"] = pretty(i["id"]) + (f" ({variant.group(1)})" if variant else "") | |
| def main(): | |
| OUT.mkdir(exist_ok=True) | |
| todo = assets() | |
| if len(sys.argv) > 1: | |
| todo = [a for a in todo if a["id"] in sys.argv[1:]] | |
| print(f"{len(todo)} assets", flush=True) | |
| items, stats = [], defaultdict(int) | |
| with ThreadPoolExecutor(12) as pool: | |
| for i, (aid, item, status) in enumerate(pool.map(build, todo)): | |
| stats[status.split(":")[0]] += 1 | |
| if item: | |
| items.append(item) | |
| elif status != "skipped": | |
| print(aid, status, flush=True) | |
| if i % 100 == 0: | |
| print(i, dict(stats), flush=True) | |
| dedupe(items) | |
| items.sort(key=lambda x: x["name"]) | |
| (OUT / "catalog.js").write_text(f"// Generated by tools/build_objects.py from {HUB} (CC-BY-4.0).\n" | |
| f"const OBJECTS_SOURCE = {json.dumps(HUB)};\nconst OBJECTS = " | |
| + json.dumps(items, separators=(",", ":")) + ";\n") | |
| print(len(items), "objects", dict(stats)) | |
| if __name__ == "__main__": | |
| main() | |