187 lines
6.2 KiB
Python
187 lines
6.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import struct
|
|
from base64 import b64decode
|
|
from typing import Any, BinaryIO, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from .utils import NpEncoder, pad_to_4_bytes, warn
|
|
|
|
|
|
GLTF_MAGIC = 0x46546C67
|
|
JSON_CHUNK_TYPE = 0x4E4F534A
|
|
BINARY_CHUNK_TYPE = 0x004E4942
|
|
DATA_URI_PREFIX = "data:application/octet-stream;base64,"
|
|
|
|
|
|
def read_header(stream: BinaryIO) -> Optional[Tuple[int, int]]:
|
|
header = stream.read(12)
|
|
if len(header) != 12:
|
|
return None
|
|
magic, version, length = struct.unpack("<III", header)
|
|
if magic != GLTF_MAGIC:
|
|
return None
|
|
return version, length
|
|
|
|
|
|
def read_chunks(stream: BinaryIO, length: int) -> Optional[Tuple[Dict[str, Any], np.ndarray]]:
|
|
json_data: Optional[Dict[str, Any]] = None
|
|
binary_data: Optional[np.ndarray] = None
|
|
|
|
while stream.tell() < length:
|
|
chunk_header = stream.read(8)
|
|
if len(chunk_header) != 8:
|
|
warn(f"Invalid chunk header size: {len(chunk_header)} bytes (expected 8)")
|
|
break
|
|
chunk_length, chunk_type = struct.unpack("<II", chunk_header)
|
|
chunk_data = stream.read(chunk_length)
|
|
if len(chunk_data) != chunk_length:
|
|
warn(
|
|
f"Incomplete chunk data: {len(chunk_data)} bytes (expected {chunk_length})"
|
|
)
|
|
return None
|
|
|
|
if chunk_type == JSON_CHUNK_TYPE:
|
|
json_data = json.loads(chunk_data)
|
|
elif chunk_type == BINARY_CHUNK_TYPE:
|
|
binary_data = np.frombuffer(chunk_data, dtype=np.uint8)
|
|
else:
|
|
warn("Unsupported chunk type")
|
|
return None
|
|
|
|
if json_data is None:
|
|
raise ValueError("Missing json header")
|
|
|
|
if binary_data is None:
|
|
binary_data = np.array([], dtype=np.uint8)
|
|
|
|
return json_data, binary_data
|
|
|
|
|
|
class GltfDocument:
|
|
"""Minimal GLTF/GLB container with helpers over core collections."""
|
|
|
|
def __init__(self, json_data: Dict[str, Any], binary_buffers: Optional[Dict[int, np.ndarray]] = None):
|
|
self._json = json_data
|
|
self._bin: Dict[int, np.ndarray] = binary_buffers or {0: np.array([], dtype=np.uint8)}
|
|
|
|
# Accessors to top-level collections
|
|
def json(self) -> Dict[str, Any]:
|
|
return self._json
|
|
|
|
def buffers(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("buffers", [])
|
|
|
|
def bufferViews(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("bufferViews", [])
|
|
|
|
def accessors(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("accessors", [])
|
|
|
|
def nodes(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("nodes", [])
|
|
|
|
def meshes(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("meshes", [])
|
|
|
|
def scenes(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("scenes", [])
|
|
|
|
def materials(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("materials", [])
|
|
|
|
def images(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("images", [])
|
|
|
|
def textures(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("textures", [])
|
|
|
|
def samplers(self) -> List[Dict[str, Any]]:
|
|
return self._json.setdefault("samplers", [])
|
|
|
|
def binary(self, index: int = 0) -> np.ndarray:
|
|
return self._bin.setdefault(index, np.array([], dtype=np.uint8))
|
|
|
|
def set_binary(self, index: int, data: np.ndarray) -> None:
|
|
self._bin[index] = data
|
|
|
|
def update_buffer_length(self, buffer_index: int = 0) -> None:
|
|
if not self.buffers():
|
|
self._json["buffers"] = [{"byteLength": int(self.binary(buffer_index).nbytes)}]
|
|
else:
|
|
self.buffers()[buffer_index]["byteLength"] = int(self.binary(buffer_index).nbytes)
|
|
|
|
|
|
def get_binary_data(doc: GltfDocument, buffer_index: int) -> np.ndarray:
|
|
"""Return binary data for a buffer index.
|
|
|
|
Supports embedded base64 buffers in GLTF as well as GLB's main buffer.
|
|
"""
|
|
buffers = doc.buffers()
|
|
if buffer_index >= len(buffers):
|
|
raise IndexError("Buffer index out of range")
|
|
binary = doc._bin.get(buffer_index)
|
|
if binary is not None and binary.size > 0:
|
|
return binary
|
|
uri = buffers[buffer_index].get("uri")
|
|
if uri and uri.startswith(DATA_URI_PREFIX):
|
|
binary_data = b64decode(uri[len(DATA_URI_PREFIX) :])
|
|
arr = np.frombuffer(binary_data, dtype=np.uint8)
|
|
doc.set_binary(buffer_index, arr)
|
|
return arr
|
|
# if no uri, return existing buffer (may be empty for GLB until provided)
|
|
return doc.binary(buffer_index)
|
|
|
|
|
|
def load_gltf_or_glb(path: str) -> GltfDocument:
|
|
with open(path, "rb") as f:
|
|
stream = f
|
|
version_and_length = read_header(stream)
|
|
if version_and_length is None: # GLTF JSON
|
|
stream.seek(0)
|
|
json_data = json.load(stream)
|
|
# embedded or external buffers will be resolved lazily
|
|
return GltfDocument(json_data=json_data)
|
|
version, length = version_and_length
|
|
if version != 2:
|
|
raise ValueError(f"Only GLB version 2 is supported, found version {version}")
|
|
json_data, binary_data = read_chunks(stream, length)
|
|
return GltfDocument(json_data=json_data, binary_buffers={0: binary_data})
|
|
|
|
|
|
def save_glb(doc: GltfDocument, stream: BinaryIO) -> None:
|
|
"""Write GLB v2 with proper 4-byte alignment for JSON and BIN chunks."""
|
|
# ensure buffers length is correct for buffer 0
|
|
doc.update_buffer_length(0)
|
|
|
|
json_bytes = json.dumps(doc.json(), cls=NpEncoder, separators=(",", ":")).encode("utf-8")
|
|
json_padded = pad_to_4_bytes(json_bytes, b"\x20")
|
|
|
|
bin_bytes = bytes(doc.binary(0))
|
|
bin_padded = pad_to_4_bytes(bin_bytes, b"\x00")
|
|
|
|
version = 2
|
|
file_header_size = 12
|
|
chunk_header_size = 8
|
|
file_length = (
|
|
file_header_size
|
|
+ chunk_header_size
|
|
+ len(json_padded)
|
|
+ chunk_header_size
|
|
+ len(bin_padded)
|
|
)
|
|
|
|
# write GLB header
|
|
stream.write(struct.pack("<III", GLTF_MAGIC, version, file_length))
|
|
|
|
# JSON chunk
|
|
stream.write(struct.pack("<II", len(json_padded), JSON_CHUNK_TYPE))
|
|
stream.write(json_padded)
|
|
|
|
# BIN chunk
|
|
stream.write(struct.pack("<II", len(bin_padded), BINARY_CHUNK_TYPE))
|
|
stream.write(bin_padded)
|
|
|