Spaces:
Running on Zero
Running on Zero
kwabs22
Claude Opus 5.5
Light unlit scenes with daylight, colour assets from thumbnails, hint at the next tier
98ca16a | """Target-language lexicon entries: load a ledger item's chosen asset as a mesh. | |
| Every mesh comes back normalized the same way, whatever its tier: | |
| Y-up (glTF convention), metres, centred on x/z, resting on y = 0. | |
| """ | |
| import glob | |
| import gzip | |
| import json | |
| import os | |
| from functools import lru_cache | |
| from urllib.parse import quote | |
| import numpy as np | |
| import trimesh | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| from .grammar import DEFAULT_HEIGHT_M, NOMINAL_HEIGHT_M, candidate_label | |
| from .lexicon import LEXICON | |
| SCENE_FACE_BUDGET = 20000 # per object, for the configured scene | |
| class UnusableAsset(Exception): | |
| """License or content rules out this candidate; the next one is tried.""" | |
| def license_ok(license_str): | |
| s = (license_str or "").lower() | |
| return bool(s) and "nc" not in s.replace("-", " ").split() and "nd" not in s.replace("-", " ").split() | |
| def _normalize(mesh, meters_per_unit=1.0, z_up=False): | |
| mesh = mesh.copy() | |
| mesh.apply_scale(meters_per_unit) | |
| if z_up: | |
| mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])) | |
| lo, hi = mesh.bounds | |
| mesh.apply_translation([-(lo[0] + hi[0]) / 2, -lo[1], -(lo[2] + hi[2]) / 2]) | |
| return mesh | |
| def _flatten(loaded): | |
| if isinstance(loaded, trimesh.Scene): | |
| to_geometry = getattr(loaded, "to_geometry", None) | |
| loaded = to_geometry() if to_geometry else loaded.dump(concatenate=True) | |
| if hasattr(loaded.visual, "to_color"): | |
| loaded.visual = loaded.visual.to_color() | |
| return loaded | |
| # ---------------- simready (OpenUSD, metric) ---------------- | |
| def _simready_files(cand): | |
| """Geometry + thumbnail only: textures are ~98% of a package (~1 GB).""" | |
| name = cand["name"] | |
| root = snapshot_download( | |
| repo_id=cand["repo"], repo_type="dataset", revision=cand["revision"], | |
| allow_patterns=[f"{name}/com.nvidia.simready.packaging.json", | |
| f"{name}/simready_usd/*.usd", f"{name}/simready_usd/*.usda", | |
| f"{name}/simready_usd/*.usdc", f"{name}/simready_usd/.thumbs/256x256/*.usd.png"]) | |
| return os.path.join(root, name) | |
| def _usd_mesh(usd_path): | |
| from pxr import Usd, UsdGeom | |
| stage = Usd.Stage.Open(usd_path) | |
| mpu = UsdGeom.GetStageMetersPerUnit(stage) or 0.01 | |
| z_up = UsdGeom.GetStageUpAxis(stage) == UsdGeom.Tokens.z | |
| cache = UsdGeom.XformCache() | |
| parts = [] | |
| for prim in Usd.PrimRange.Stage(stage, Usd.TraverseInstanceProxies()): | |
| if not prim.IsA(UsdGeom.Mesh): | |
| continue | |
| path = str(prim.GetPath()).lower() | |
| if "collision" in path or "collider" in path: | |
| continue | |
| img = UsdGeom.Imageable(prim) | |
| if img.ComputeVisibility() == UsdGeom.Tokens.invisible or img.ComputePurpose() in ( | |
| UsdGeom.Tokens.guide, UsdGeom.Tokens.proxy): | |
| continue | |
| geom = UsdGeom.Mesh(prim) | |
| points = geom.GetPointsAttr().Get() | |
| counts = geom.GetFaceVertexCountsAttr().Get() | |
| indices = geom.GetFaceVertexIndicesAttr().Get() | |
| if points is None or counts is None or len(points) == 0 or len(counts) == 0: | |
| continue | |
| pts = np.asarray(points, dtype=np.float64) | |
| xf = np.array(cache.GetLocalToWorldTransform(prim), dtype=np.float64) # row-vector convention | |
| pts = (np.c_[pts, np.ones(len(pts))] @ xf)[:, :3] | |
| idx = np.asarray(indices, dtype=np.int64) | |
| counts = np.asarray(counts, dtype=np.int64) | |
| if np.all(counts == 3): | |
| faces = idx.reshape(-1, 3) | |
| else: | |
| faces, at = [], 0 | |
| for c in counts: # fan-triangulate polygons | |
| faces.extend((idx[at], idx[at + k], idx[at + k + 1]) for k in range(1, c - 1)) | |
| at += c | |
| faces = np.asarray(faces, dtype=np.int64) | |
| parts.append(trimesh.Trimesh(pts, faces, process=False)) | |
| if not parts: | |
| raise UnusableAsset(f"no visible mesh geometry in {os.path.basename(usd_path)}") | |
| return _normalize(trimesh.util.concatenate(parts), mpu, z_up) | |
| def load_simready(cand): | |
| folder = _simready_files(cand) | |
| with open(os.path.join(folder, "com.nvidia.simready.packaging.json"), encoding="utf-8") as f: | |
| lic = json.load(f).get("license") | |
| if not license_ok(lic): | |
| raise UnusableAsset(f"license {lic!r}") | |
| roots = sorted(glob.glob(os.path.join(folder, "simready_usd", "*.usd*"))) | |
| if not roots: | |
| raise UnusableAsset("no root USD file") | |
| thumbs = sorted(glob.glob(os.path.join(folder, "simready_usd", ".thumbs", "256x256", "*.usd.png"))) | |
| mesh, url, color = _usd_mesh(roots[0]), None, None | |
| if thumbs: | |
| rel = "/".join([cand["name"], "simready_usd", ".thumbs", "256x256", os.path.basename(thumbs[0])]) | |
| url = f"https://huggingface.co/datasets/{cand['repo']}/resolve/{cand['revision']}/{quote(rel)}" | |
| color = thumbnail_color(thumbs[0]) | |
| if color is not None: # geometry comes without textures; the thumbnail still knows the colour | |
| mesh.visual = trimesh.visual.ColorVisuals(mesh, face_colors=np.tile(color, (len(mesh.faces), 1))) | |
| return mesh, {"license": lic, "metric": True, "thumbnail": url, | |
| "color": None if color is None else "#%02x%02x%02x" % tuple(color[:3])} | |
| def thumbnail_color(path): | |
| """Mean RGBA of the object in a thumbnail: transparent pixels, or pixels close to the | |
| border colour (the backdrop), are left out. None when nothing is left.""" | |
| from PIL import Image | |
| img = np.asarray(Image.open(path).convert("RGBA").resize((64, 64)), dtype=np.float64) | |
| rgb, alpha = img[..., :3], img[..., 3] | |
| if (alpha < 128).any(): | |
| mask = alpha >= 128 | |
| else: | |
| border = np.concatenate([rgb[0], rgb[-1], rgb[:, 0], rgb[:, -1]]) | |
| mask = np.linalg.norm(rgb - np.median(border, axis=0), axis=-1) > 30 | |
| if mask.sum() < 20: | |
| return None | |
| return np.r_[rgb[mask].mean(axis=0), 255].astype(np.uint8) | |
| # ---------------- objaverse (GLB, no real-world scale) ---------------- | |
| def _objaverse_license(uid, shard): | |
| path = hf_hub_download("allenai/objaverse", f"metadata/{shard}.json.gz", repo_type="dataset") | |
| with gzip.open(path, "rt", encoding="utf-8") as f: | |
| return json.load(f).get(uid, {}).get("license") | |
| def load_objaverse(cand): | |
| rel = LEXICON.object_path(cand["uid"]) # glbs/000-023/<uid>.glb | |
| lic = _objaverse_license(cand["uid"], rel.split("/")[1]) | |
| if not license_ok(lic): | |
| raise UnusableAsset(f"license {lic!r}") | |
| path = hf_hub_download("allenai/objaverse", rel, repo_type="dataset") | |
| mesh = _flatten(trimesh.load(path, force="scene")) | |
| if len(mesh.faces) == 0: | |
| raise UnusableAsset("empty mesh") | |
| return _normalize(mesh), {"license": lic, "metric": False, "thumbnail": None, "color": None} | |
| # ---------------- construction tiers ---------------- | |
| def load_primitive(cand, word): | |
| h = NOMINAL_HEIGHT_M.get(word, DEFAULT_HEIGHT_M) | |
| shape = cand.get("shape", "box") | |
| if shape == "sphere": | |
| mesh = trimesh.creation.icosphere(subdivisions=3, radius=h / 2) | |
| elif shape in ("cylinder", "cone"): | |
| make = trimesh.creation.cylinder if shape == "cylinder" else trimesh.creation.cone | |
| mesh = make(radius=h / 2, height=h, sections=32) | |
| mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])) # axis z -> y | |
| elif shape == "plane": | |
| mesh = trimesh.creation.box([4 * h, 0.01, 4 * h]) | |
| else: | |
| mesh = trimesh.creation.box([h, h, h]) | |
| return _normalize(mesh), {"license": None, "metric": False, "thumbnail": None, "color": None} | |
| def load_generation(cand, word): | |
| raise UnusableAsset("generation tier is not wired up yet") | |
| # ---------------- entry points ---------------- | |
| def _load(key): | |
| cand = json.loads(key) | |
| word = cand.pop("_word") | |
| source = cand["source"] | |
| if source == "simready": | |
| return load_simready(cand) | |
| if source == "objaverse": | |
| return load_objaverse(cand) | |
| if source == "primitive": | |
| return load_primitive(cand, word) | |
| return load_generation(cand, word) | |
| def load_candidate(cand, word): | |
| mesh, info = _load(json.dumps({**cand, "_word": word}, sort_keys=True)) | |
| return mesh.copy(), info | |
| def settle(item): | |
| """Make the item's chosen candidate loadable, skipping ones that can't be used. | |
| Returns the full-detail mesh. Skips are recorded on the ledger so the | |
| discussion can see what was passed over and why. | |
| """ | |
| dec = item["decision"] | |
| while True: | |
| try: | |
| mesh, info = load_candidate(dec["chosen"], item["word"]) | |
| break | |
| except UnusableAsset as e: | |
| dec["skipped"].append({"candidate": candidate_label(dec["chosen"]), "reason": str(e)}) | |
| if not dec["alternatives"]: | |
| raise UnusableAsset(f"no usable {dec['tier']} asset for '{item['word']}' ({e})") | |
| dec["chosen"] = dec["alternatives"].pop(0) | |
| item["license"] = info["license"] | |
| item["thumbnail"] = info["thumbnail"] # a public URL, so the chat can show and compare it | |
| item["color"] = info["color"] | |
| height = item["height_m"] | |
| native = float(mesh.extents[1]) or 1.0 | |
| if height["source"] == "user": | |
| pass | |
| elif info["metric"]: | |
| height.update(value=round(native, 3), source="asset") | |
| else: | |
| height.update(value=NOMINAL_HEIGHT_M.get(item["word"], DEFAULT_HEIGHT_M), source="default") | |
| mesh.apply_scale(height["value"] / native) | |
| return mesh | |
| def scene_mesh(mesh): | |
| """A lighter copy for the configured scene; the ladder keeps full detail.""" | |
| from .simplify import decimate | |
| return decimate(mesh, SCENE_FACE_BUDGET) if len(mesh.faces) > SCENE_FACE_BUDGET else mesh | |