File size: 7,505 Bytes
9f85448
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
"""Minimal glTF binary reader."""

from __future__ import annotations

import json
import struct
from dataclasses import dataclass, field
from pathlib import Path

import cv2
import numpy as np

_JSON_CHUNK = 0x4E4F534A
_BIN_CHUNK = 0x004E4942

_DTYPES = {
    5120: np.int8, 5121: np.uint8, 5122: np.int16,
    5123: np.uint16, 5125: np.uint32, 5126: np.float32,
}
_COUNTS = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT2": 4, "MAT3": 9, "MAT4": 16}


@dataclass
class Primitive:
    """One drawable mesh piece."""

    positions: np.ndarray
    normals: np.ndarray
    joints: np.ndarray
    weights: np.ndarray
    triangles: np.ndarray
    color: tuple[int, int, int] = (185, 180, 175)
    uv: np.ndarray | None = None
    texture: np.ndarray | None = None


@dataclass
class Node:
    """One scene node."""

    name: str
    matrix: np.ndarray
    children: list[int] = field(default_factory=list)
    parent: int | None = None


@dataclass
class Gltf:
    """Loaded model."""

    nodes: list[Node]
    primitives: list[Primitive]
    joints: list[int]
    inverse_bind: np.ndarray
    globals: np.ndarray


def _node_matrix(node: dict) -> np.ndarray:
    """Local transform of one node."""
    if "matrix" in node:
        return np.array(node["matrix"], dtype=np.float64).reshape(4, 4).T
    m = np.eye(4)
    if "rotation" in node:
        x, y, z, w = node["rotation"]
        m[:3, :3] = np.array([
            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
        ])
    if "scale" in node:
        m[:3, :3] = m[:3, :3] @ np.diag(node["scale"])
    if "translation" in node:
        m[:3, 3] = node["translation"]
    return m


def _decode_image(js: dict, blob: bytes, image_index: int) -> np.ndarray | None:
    """Decode one embedded glTF image to a BGR array."""
    images = js.get("images") or []
    if not 0 <= image_index < len(images):
        return None
    image = images[image_index]
    if "bufferView" not in image:
        return None
    view = js["bufferViews"][image["bufferView"]]
    start = view.get("byteOffset", 0)
    raw = np.frombuffer(blob, dtype=np.uint8, count=view["byteLength"], offset=start)
    decoded = cv2.imdecode(raw, cv2.IMREAD_COLOR)
    return decoded if decoded is not None and decoded.size else None


def _material_texture(js: dict, blob: bytes, material_index: int) -> np.ndarray | None:
    """Base-color texture of one material, if it has an embedded one."""
    pbr = js["materials"][material_index].get("pbrMetallicRoughness", {})
    ref = pbr.get("baseColorTexture")
    if ref is None:
        return None
    textures = js.get("textures") or []
    if not 0 <= ref.get("index", -1) < len(textures):
        return None
    source = textures[ref["index"]].get("source")
    return None if source is None else _decode_image(js, blob, source)


class _Reader:
    """Accessor decoder."""

    def __init__(self, js: dict, blob: bytes) -> None:
        self.js = js
        self.blob = blob

    def __call__(self, index: int) -> np.ndarray:
        """Read one accessor."""
        acc = self.js["accessors"][index]
        dtype = _DTYPES[acc["componentType"]]
        ncomp = _COUNTS[acc["type"]]
        count = acc["count"]
        if "bufferView" not in acc:
            return np.zeros((count, ncomp), dtype=np.float64)
        view = self.js["bufferViews"][acc["bufferView"]]
        start = view.get("byteOffset", 0) + acc.get("byteOffset", 0)
        stride = view.get("byteStride") or np.dtype(dtype).itemsize * ncomp
        item = np.dtype(dtype).itemsize * ncomp
        if stride == item:
            raw = np.frombuffer(self.blob, dtype=dtype, count=count * ncomp, offset=start)
            out = raw.reshape(count, ncomp)
        else:
            rows = [np.frombuffer(self.blob, dtype=dtype, count=ncomp, offset=start + i * stride)
                    for i in range(count)]
            out = np.stack(rows)
        return out.astype(np.float64) if dtype == np.float32 else out


def load(path: str | Path) -> Gltf:
    """Read a glb file."""
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Character model not found: {path}")
    data = path.read_bytes()
    magic, _version, total = struct.unpack_from("<III", data, 0)
    if magic != 0x46546C67:
        raise ValueError(f"Not a binary glTF: {path}")

    js: dict | None = None
    blob = b""
    off = 12
    while off < total:
        length, kind = struct.unpack_from("<II", data, off)
        chunk = data[off + 8:off + 8 + length]
        if kind == _JSON_CHUNK:
            js = json.loads(chunk)
        elif kind == _BIN_CHUNK:
            blob = bytes(chunk)
        off += 8 + length + ((4 - length % 4) % 4 if length % 4 else 0)
    if js is None:
        raise ValueError(f"No json chunk in {path}")

    read = _Reader(js, blob)
    nodes = [Node(name=n.get("name", f"node{i}"), matrix=_node_matrix(n),
                  children=list(n.get("children", []))) for i, n in enumerate(js["nodes"])]
    for i, node in enumerate(nodes):
        for c in node.children:
            nodes[c].parent = i

    globals_ = np.tile(np.eye(4), (len(nodes), 1, 1))
    order = [i for i in range(len(nodes)) if nodes[i].parent is None]
    while order:
        i = order.pop(0)
        parent = nodes[i].parent
        base = globals_[parent] if parent is not None else np.eye(4)
        globals_[i] = base @ nodes[i].matrix
        order.extend(nodes[i].children)

    skin = js["skins"][0]
    joints = list(skin["joints"])
    inverse_bind = read(skin["inverseBindMatrices"]).reshape(-1, 4, 4).transpose(0, 2, 1)

    primitives: list[Primitive] = []
    for mesh in js.get("meshes", []):
        for prim in mesh["primitives"]:
            attrs = prim["attributes"]
            if "JOINTS_0" not in attrs or "WEIGHTS_0" not in attrs:
                continue
            pos = read(attrs["POSITION"])
            normals = read(attrs["NORMAL"]) if "NORMAL" in attrs else np.zeros_like(pos)
            jnt = read(attrs["JOINTS_0"]).astype(np.int32)
            wts = read(attrs["WEIGHTS_0"]).astype(np.float64)
            wsum = wts.sum(axis=1, keepdims=True)
            wts = wts / np.where(wsum > 1e-9, wsum, 1.0)
            if "indices" in prim:
                idx = read(prim["indices"]).astype(np.int64).reshape(-1)
            else:
                idx = np.arange(len(pos), dtype=np.int64)
            color = (185, 180, 175)
            texture = None
            mat = prim.get("material")
            if mat is not None:
                pbr = js["materials"][mat].get("pbrMetallicRoughness", {})
                base = pbr.get("baseColorFactor")
                if base:
                    color = tuple(int(np.clip(c * 255, 0, 255)) for c in base[2::-1])
                texture = _material_texture(js, blob, mat)
            uv = read(attrs["TEXCOORD_0"]) if "TEXCOORD_0" in attrs else None
            if texture is None:
                uv = None
            primitives.append(Primitive(positions=pos, normals=normals, joints=jnt,
                                        weights=wts, triangles=idx.reshape(-1, 3),
                                        color=color, uv=uv, texture=texture))
    return Gltf(nodes=nodes, primitives=primitives, joints=joints,
                inverse_bind=inverse_bind, globals=globals_)