tripocli/tripo/glb.py

"""Read metadata and node names out of a GLB/GLTF file, using only the standard library.

The studio builds ``part_name_list`` from the mesh's *node names*, which live
inside the model file — the project detail API never returns them.  Unsegmented
models carry a single ``tripo_node_<operator_id>`` node; a segmented model
carries one ``tripo_part_<n>`` node per part.
"""

import json
import os
import struct

GLB_MAGIC = b"glTF"
JSON_CHUNK = 0x4E4F534A


def gltf_document(data):
    """Return the glTF JSON document from a ``.glb`` or ``.gltf`` byte string."""
    if data[:4] == GLB_MAGIC:
        offset = 12
        while offset < len(data):
            length, kind = struct.unpack_from("<II", data, offset)
            offset += 8
            chunk = data[offset:offset + length]
            offset += length
            if kind == JSON_CHUNK:
                return json.loads(chunk.decode("utf-8"))
        raise ValueError("GLB 缺少 JSON chunk")
    if data[:18] == b"Kaydara FBX Binary":
        raise ValueError("该文件实际是 FBX 二进制格式（而非 GLB），可用 `tripocli export -p <pid> -f glb --wait` 转为 GLB")
    return json.loads(data.decode("utf-8"))


def read_gltf_file(path):
    """Read just the JSON header chunk from a local ``.glb`` or ``.gltf`` file."""
    with open(path, "rb") as handle:
        head = handle.read(20)
        if head[:4] == GLB_MAGIC and len(head) >= 20:
            length, kind = struct.unpack_from("<II", head, 12)
            if kind == JSON_CHUNK:
                chunk = handle.read(length)
                return json.loads(chunk.decode("utf-8"))
        handle.seek(0)
        return gltf_document(handle.read())


def node_names(data):
    """Named nodes of the model, in file order."""
    document = gltf_document(data)
    return [node.get("name") for node in (document.get("nodes") or []) if node.get("name")]


def inspect_file(path):
    """Extract game-relevant stats (triangles, bbox, textures, joints, animations) from a GLB/GLTF."""
    abspath = os.path.abspath(path)
    file_size = os.path.getsize(abspath)
    doc = read_gltf_file(abspath)

    accessors = doc.get("accessors") or []
    meshes = doc.get("meshes") or []
    nodes = doc.get("nodes") or []
    materials = doc.get("materials") or []
    textures = doc.get("textures") or []
    images = doc.get("images") or []
    skins = doc.get("skins") or []
    animations = doc.get("animations") or []
    extensions = doc.get("extensionsUsed") or []

    vertices = 0
    triangles = 0
    primitives_count = 0
    min_pt = [float("inf"), float("inf"), float("inf")]
    max_pt = [float("-inf"), float("-inf"), float("-inf")]
    has_bbox = False

    for mesh in meshes:
        for prim in mesh.get("primitives") or []:
            primitives_count += 1
            attrs = prim.get("attributes") or {}
            pos_idx = attrs.get("POSITION")
            pos_count = 0
            if pos_idx is not None and 0 <= pos_idx < len(accessors):
                pos_acc = accessors[pos_idx]
                pos_count = pos_acc.get("count") or 0
                vertices += pos_count
                amin = pos_acc.get("min")
                amax = pos_acc.get("max")
                if isinstance(amin, list) and isinstance(amax, list) and len(amin) >= 3 and len(amax) >= 3:
                    has_bbox = True
                    for axis in range(3):
                        min_pt[axis] = min(min_pt[axis], float(amin[axis]))
                        max_pt[axis] = max(max_pt[axis], float(amax[axis]))

            mode = prim.get("mode", 4)  # 4 = TRIANGLES
            idx_acc_id = prim.get("indices")
            if idx_acc_id is not None and 0 <= idx_acc_id < len(accessors):
                idx_count = accessors[idx_acc_id].get("count") or 0
                if mode == 4:
                    triangles += idx_count // 3
                elif mode in (5, 6):
                    triangles += max(0, idx_count - 2)
            elif mode == 4 and pos_count:
                triangles += pos_count // 3

    bbox = None
    if has_bbox:
        size = [round(max_pt[i] - min_pt[i], 4) for i in range(3)]
        bbox = {
            "min": [round(v, 4) for v in min_pt],
            "max": [round(v, 4) for v in max_pt],
            "size": size,
            "ground_offset_y": round(-min_pt[1], 4),
        }

    has_pbr = False
    for mat in materials:
        pbr = mat.get("pbrMetallicRoughness") or {}
        if pbr.get("metallicRoughnessTexture") or mat.get("normalTexture"):
            has_pbr = True
            break

    joint_indices = set()
    for skin in skins:
        joint_indices.update(skin.get("joints") or [])

    mesh_nodes = [
        n.get("name") for i, n in enumerate(nodes)
        if n.get("name") and ("mesh" in n or i not in joint_indices)
    ]
    anim_list = [
        {"name": a.get("name") or f"clip_{i}", "channels": len(a.get("channels") or [])}
        for i, a in enumerate(animations)
    ]

    warnings = []
    if not meshes and anim_list:
        warnings.append("纯骨骼动画轨道文件（无网格），在引擎中需挂载到 rig 导出的角色本体网格上播放")
    if triangles > 100000:
        warnings.append(
            f"面数较高 ({triangles:,} 三角面)，放入实时游戏场景建议先运行: "
            f"tripocli optimize {os.path.basename(path)}"
        )
    if meshes and not textures:
        warnings.append("无贴图白模（materials 未包含纹理贴图）")

    return {
        "file": abspath,
        "name": os.path.basename(abspath),
        "file_size": file_size,
        "file_size_mb": round(file_size / (1024 * 1024), 2),
        "meshes": len(meshes),
        "primitives": primitives_count,
        "vertices": vertices,
        "triangles": triangles,
        "bbox": bbox,
        "materials": len(materials),
        "textures": len(textures),
        "images": len(images),
        "has_pbr_maps": has_pbr,
        "is_rigged": bool(skins),
        "joints": len(joint_indices),
        "animations": anim_list,
        "draco_compressed": "KHR_draco_mesh_compression" in extensions,
        "extensions": extensions,
        "mesh_nodes": mesh_nodes[:20],
        "warnings": warnings,
    }