Blame for vli/assets.py

"""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 ----------------

@lru_cache(maxsize=64)
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