Coverage for trimesh/exchange/gltf/extensions.py: 86%
113 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-08-31 18:21 +0000
1"""
2gltf_extensions.py
3------------------
5Extension registry for glTF import/export with scope-based handlers.
6Each scope has a TypedDict defining the context passed to handlers.
7"""
9from collections.abc import Callable, Iterable
10from typing import Any, Literal, TypeAlias, TypedDict
12import numpy as np
14from ...constants import log
15from ...iteration import IndexedDict
16from ...typed import NDArray
18# Scopes define where in the glTF load/export process handlers run:
19# material - after parsing material, can override PBR values
20# texture_source - when resolving texture image index
21# primitive - after loading primitive, can add face_attributes
22# primitive_preprocess - before accessor reads, can modify accessors in-place
23# primitive_export - during mesh export, can compress/modify primitive data
24Scope: TypeAlias = Literal[
25 "material", "texture_source", "primitive", "primitive_preprocess", "primitive_export"
26]
29# ----------------------------------------------------------------------
30# TypedDict contexts for each scope
31# ----------------------------------------------------------------------
32#
33# These TypedDicts define the MINIMUM fields passed to handlers for each scope.
34# Additional fields may be added in future versions for new functionality.
35#
36# FOR FORWARD COMPATIBILITY: Handlers should access only the fields they need
37# and ignore unknown fields. The context is passed as a plain dict at runtime,
38# so handlers can safely use dict.get() for optional access or simply not
39# reference fields they don't need.
40#
41# Example handler pattern:
42#
43# def my_handler(context: MaterialContext) -> dict | None:
44# # Access only what you need - additional fields won't break this
45# data = context["data"]
46# images = context["images"]
47# return {"baseColorFactor": [1, 0, 0, 1]}
48#
49# ----------------------------------------------------------------------
52class MaterialContext(TypedDict):
53 """Context for material scope handlers."""
55 data: dict[str, Any]
56 parse_textures: Callable[..., dict[str, Any]]
57 images: list
60class TextureSourceContext(TypedDict):
61 """Context for texture_source scope handlers."""
63 data: dict[str, Any]
66class PrimitiveContext(TypedDict):
67 """Context for primitive scope handlers (post-load)."""
69 data: dict[str, Any]
70 primitive: dict
71 mesh_kwargs: dict
72 accessors: list
75class PrimitivePreprocessContext(TypedDict):
76 """Context for primitive_preprocess scope handlers (pre-load)."""
78 data: dict[str, Any]
79 primitive: dict
80 accessors: list
81 views: list
84class PrimitiveExportContext(TypedDict):
85 """Context for primitive_export scope handlers (during export)."""
87 mesh: Any
88 name: str
89 tree: dict
90 # a `bufferView` is the position of an entry in here, so the order matters
91 buffer_items: IndexedDict
92 primitive: dict
93 # the arrays the primitive's accessors were built from, in the dtype they
94 # would have been written with, keyed by accessor index: a handler storing
95 # them itself never unpacks bytes back into numpy, and the exporter stores
96 # them after all if no handler claims them
97 arrays: dict[int, NDArray]
100# Handler type alias - handlers receive a context dict
101Handler: TypeAlias = Callable[[Any], Any]
103# callback to parse material dict and resolve texture references
104# signature: (*, data: dict) -> dict
105ParseTextures: TypeAlias = Callable[..., dict[str, Any]]
107# Registry: {scope: {extension_name: handler}}
108_handlers: dict[str, dict[str, Handler]] = {}
111def _deep_merge(target: dict, source: dict) -> None:
112 """
113 Recursively merge source dict into target dict.
115 Parameters
116 ----------
117 target
118 Dict to merge into (modified in place)
119 source
120 Dict to merge from
121 """
122 for key, value in source.items():
123 if isinstance(value, dict) and key in target and isinstance(target[key], dict):
124 # Both are dicts - recurse
125 _deep_merge(target[key], value)
126 else:
127 # Overwrite or set new key
128 target[key] = value
131def register_handler(name: str, scope: Scope) -> Callable[[Handler], Handler]:
132 """
133 Decorator to register a handler for a glTF extension.
135 Parameters
136 ----------
137 name
138 Extension name, e.g. "KHR_materials_pbrSpecularGlossiness".
139 scope
140 Handler scope, e.g. "material", "texture_source", "primitive".
142 Returns
143 -------
144 decorator
145 Function that registers the handler and returns it unchanged.
147 Example
148 -------
149 >>> @register_handler("MY_extension", scope="material")
150 ... def my_handler(context: MaterialContext) -> dict | None:
151 ... data = context["data"]
152 ... images = context["images"]
153 ... return {"baseColorFactor": [1, 0, 0, 1]}
154 """
155 if scope not in _handlers:
156 _handlers[scope] = {}
158 def decorator(func: Handler) -> Handler:
159 _handlers[scope][name] = func
160 return func
162 return decorator
165def unregistered(extensions: Iterable[str], scope: Scope) -> set:
166 """
167 Find extension names with no registered handler for a scope.
169 Parameters
170 ----------
171 extensions
172 Extension names, i.e. the keys of a glTF "extensions" dict.
173 scope
174 Handler scope to check against.
176 Returns
177 -------
178 missing
179 Extension names with no handler registered for the scope.
180 """
181 return set(extensions) - _handlers.get(scope, {}).keys()
184def handle_extensions(
185 *,
186 extensions: dict[str, Any] | None,
187 scope: Scope,
188 failed: set | None = None,
189 **kwargs,
190) -> Any:
191 """
192 Process extensions dict for a given scope, calling registered handlers.
194 Parameters
195 ----------
196 extensions
197 The "extensions" dict from a glTF element, or None.
198 scope
199 Handler scope to invoke.
200 failed
201 If passed the name of any extension whose handler raised is added here.
202 **kwargs
203 Scope-specific arguments that will be combined with extension data
204 into a typed context dict. Required kwargs by scope:
205 - material: parse_textures, images
206 - texture_source: (none)
207 - primitive: primitive, mesh_kwargs, accessors
208 - primitive_preprocess: primitive, accessors, views
209 - primitive_export: mesh, name, tree, buffer_items, primitive, arrays
211 Returns
212 -------
213 results
214 Dict of {extension_name: result} for most scopes.
215 For scopes ending in "_source", returns first non-None result.
216 For "primitive" scope, automatically merges results into mesh_kwargs.
217 """
218 if not extensions or scope not in _handlers:
219 return {} if not scope.endswith("_source") else None
221 results = {}
222 for ext_name, data in extensions.items():
223 if ext_name not in _handlers[scope]:
224 continue
225 try:
226 # Build context dict with data + all kwargs
227 context = {"data": data, **kwargs}
228 if (result := _handlers[scope][ext_name](context)) is not None:
229 results[ext_name] = result
230 except Exception as e:
231 if failed is not None:
232 failed.add(ext_name)
233 log.warning(f"failed to process extension {ext_name}: {e}")
235 # for _source scopes return first result, otherwise return all results
236 if scope.endswith("_source"):
237 return next(iter(results.values()), None)
239 # for primitive scope, automatically merge results into mesh_kwargs
240 if scope == "primitive" and "mesh_kwargs" in kwargs:
241 mesh_kwargs = kwargs["mesh_kwargs"]
242 for ext_result in results.values():
243 if not isinstance(ext_result, dict):
244 continue
245 # merge extension results, recursively merging nested dicts
246 for key, value in ext_result.items():
247 if isinstance(value, dict):
248 if key not in mesh_kwargs:
249 mesh_kwargs[key] = {}
250 _deep_merge(mesh_kwargs[key], value)
251 else:
252 mesh_kwargs[key] = value
254 return results
257# ----------------------------------------------------------------------
258# Built-in handlers
259# ----------------------------------------------------------------------
262@register_handler("KHR_materials_pbrSpecularGlossiness", scope="material")
263def _specular_glossiness(context: MaterialContext) -> dict[str, Any] | None:
264 """
265 Convert specular-glossiness material to PBR metallic-roughness.
267 Parameters
268 ----------
269 context
270 MaterialContext with extension data, parse_textures function, and images.
272 Returns
273 -------
274 pbr_dict
275 PBR metallic-roughness parameters, or None on failure.
276 """
277 try:
278 from ...visual.gloss import specular_to_pbr
280 return specular_to_pbr(**context["parse_textures"](data=context["data"]))
281 except Exception:
282 log.debug("failed to convert specular-glossiness", exc_info=True)
283 return None
286@register_handler("EXT_texture_webp", scope="texture_source")
287def _texture_webp_source(context: TextureSourceContext) -> int | None:
288 """
289 Return image source index from EXT_texture_webp.
291 Parameters
292 ----------
293 context
294 TextureSourceContext with extension data.
296 Returns
297 -------
298 source_index
299 Index into glTF images array, or None if not present.
300 """
301 return context["data"].get("source")
304# the optional attributes draco can absorb, as
305# (glTF name, `DracoPy.encode` keyword, `DracoPy.AttributeType`)
306_draco_optional = (
307 ("COLOR_0", "colors", "COLOR"),
308 ("TEXCOORD_0", "tex_coord", "TEX_COORD"),
309 ("NORMAL", "normals", "NORMAL"),
310)
312# how hard draco tries, which trades export time for size and is not lossy
313_DRACO_COMPRESSION = 6
315# bits draco quantizes positions onto, which is what blender emits: a vertex
316# moves up to half a step, i.e. `mesh.extents.max() * 2**-(bits + 1)`
317_DRACO_QUANTIZATION = 14
320@register_handler("KHR_draco_mesh_compression", scope="primitive_preprocess")
321def draco_decode(context: PrimitivePreprocessContext) -> None:
322 """
323 Replace a primitive's placeholder accessors with decompressed draco data.
325 The accessors of a draco-compressed primitive have no `bufferView`, so the
326 loader filled them with zeros before calling us. All of the geometry is in
327 a single opaque blob, and the extension carries the indirection we need to
328 unpack it: a mapping of glTF attribute name to draco attribute id.
330 Parameters
331 ----------
332 context
333 PrimitivePreprocessContext, whose `accessors` we mutate in-place.
334 """
335 import DracoPy
337 data = context["data"]
338 accessors = context["accessors"]
339 attributes = context["primitive"].get("attributes", {})
341 # one blob holds every compressed attribute for this primitive
342 decoded = DracoPy.decode(context["views"][data["bufferView"]])
344 # the extension stores name -> draco id and we look up the other way
345 names = {ident: name for name, ident in data["attributes"].items()}
347 # overwrite the zero placeholders in-place by accessor index
348 for attr in decoded.attributes:
349 name = names.get(attr["unique_id"])
350 if name in attributes:
351 accessors[attributes[name]] = attr["data"]
353 # indices aren't an attribute so they're not in the extension mapping
354 indices = context["primitive"].get("indices")
355 faces = getattr(decoded, "faces", None)
356 if indices is not None and faces is not None:
357 accessors[indices] = faces
360@register_handler("KHR_draco_mesh_compression", scope="primitive_export")
361def draco_encode(context: PrimitiveExportContext) -> bool | None:
362 """
363 Compress a primitive's geometry into a single draco buffer.
365 Every array in `arrays` is absorbed, so the exporter left their accessors
366 with no `bufferView` rather than storing the same data twice. Returning
367 `None` makes it store them after all, so a failure here exports the
368 primitive uncompressed rather than pointing at accessors full of zeros.
370 Parameters
371 ----------
372 context
373 PrimitiveExportContext, whose `primitive` and `buffer_items` we mutate.
375 Returns
376 -------
377 compressed
378 True if the geometry is now inside a draco buffer, None if not.
379 """
380 import DracoPy
382 from . import _buffer_append
384 primitive = context["primitive"]
385 attributes = primitive["attributes"]
386 # the arrays the accessors were built from keyed by accessor index, so we
387 # never unpack bytes back into numpy and can only claim what was recorded
388 arrays = context["arrays"]
390 # the optional attributes this primitive actually has: anything the
391 # exporter stored itself, like a custom `_ATTRIBUTE`, is not in `arrays`
392 absorb, optional = [], {}
393 for name, keyword, kind in _draco_optional:
394 value = arrays.get(attributes.get(name))
395 if value is None:
396 continue
397 absorb.append((name, kind))
398 # DracoPy asserts on float64 for UV and normals, colors stay uint8
399 optional[keyword] = value if value.dtype == np.uint8 else value.astype(np.float64)
401 position = arrays[attributes["POSITION"]]
402 indices = arrays[primitive["indices"]]
404 # `preserve_order` is load-bearing: without it draco rewelds and permutes
405 # vertices, which would invalidate the count/min/max already written into
406 # every accessor and silently misindex any attribute we didn't compress
407 buffer = DracoPy.encode(
408 points=position,
409 faces=indices,
410 preserve_order=True,
411 quantization_bits=_DRACO_QUANTIZATION,
412 compression_level=_DRACO_COMPRESSION,
413 **optional,
414 )
416 # decode what we just encoded, which does two things we can't get otherwise:
417 # it reports the attribute ids draco assigned, and it proves the round trip
418 # preserved our counts before we let the exporter skip storing the source.
419 # DO NOT remove this check: it is all that stands between a lossy encoder
420 # and silently corrupt geometry in the exported file.
421 check = DracoPy.decode(buffer)
422 if len(check.points) != len(position) or len(check.faces) != len(indices):
423 log.warning("draco round trip changed the vertex count, not compressing")
424 return None
426 # draco identifies attributes by kind, the extension identifies them by id.
427 # an attribute draco silently dropped is missing from `ids`, and the
428 # `KeyError` leaves this primitive uncompressed rather than pointing at
429 # an accessor the exporter is about to stop storing
430 ids = {attr["attribute_type"]: attr["unique_id"] for attr in check.attributes}
431 absorbed = {
432 name: ids[getattr(DracoPy.AttributeType, kind)]
433 for name, kind in [("POSITION", "POSITION"), *absorb]
434 }
436 primitive.setdefault("extensions", {})["KHR_draco_mesh_compression"] = {
437 # identical geometry encodes identically so this dedupes for free
438 "bufferView": _buffer_append(context["buffer_items"], buffer),
439 "attributes": absorbed,
440 }
441 return True