Coverage for trimesh/exchange/dae.py: 86%
218 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-24 04:40 +0000
1import copy
2import io
3import uuid
5import numpy as np
7from .. import util, visual
8from ..constants import log
9from ..util import unique_name
11_EYE = np.eye(4)
12_EYE.flags.writeable = False
15def load_collada(file_obj, resolver=None, ignore_broken=True, **kwargs):
16 """
17 Load a COLLADA (.dae) file into a list of trimesh kwargs.
19 Parameters
20 ----------
21 file_obj : file object
22 Containing a COLLADA file
23 resolver : trimesh.visual.Resolver or None
24 For loading referenced files, like texture images
25 ignore_broken: bool
26 Ignores broken references during loading:
27 [collada.common.DaeUnsupportedError,
28 collada.common.DaeBrokenRefError]
29 kwargs : **
30 Passed to trimesh.Trimesh.__init__
32 Returns
33 -------
34 loaded : list of dict
35 kwargs for Trimesh constructor
36 """
37 import collada
39 if ignore_broken:
40 ignores = [
41 collada.common.DaeError,
42 collada.common.DaeIncompleteError,
43 collada.common.DaeMalformedError,
44 collada.common.DaeBrokenRefError,
45 collada.common.DaeUnsupportedError,
46 collada.common.DaeIncompleteError,
47 ]
48 else:
49 ignores = None
51 # load scene using pycollada
52 c = collada.Collada(file_obj, ignore=ignores)
54 # Create material map from Material ID to trimesh material
55 material_map = {}
56 for m in c.materials:
57 effect = m.effect
58 material_map[m.id] = _parse_material(effect, resolver)
60 unit = c.assetInfo.unitmeter
61 if unit is None or np.isclose(unit, 1.0):
62 metadata = {"units": "meters"}
63 else:
64 metadata = {"units": f"{unit} * meters"}
66 # name : kwargs
67 meshes = {}
68 # increments to enable `unique_name` to avoid n^2 behavior
69 meshes_count = {}
70 # list of dict
71 graph = []
73 for node in c.scene.nodes:
74 _parse_node(
75 node=node,
76 parent_matrix=_EYE,
77 material_map=material_map,
78 meshes=meshes,
79 meshes_count=meshes_count,
80 graph=graph,
81 resolver=resolver,
82 metadata=metadata,
83 )
85 return {"class": "Scene", "graph": graph, "geometry": meshes}
88def export_collada(mesh, **kwargs):
89 """
90 Export a mesh or a list of meshes as a COLLADA .dae file.
92 Parameters
93 -----------
94 mesh: Trimesh object or list of Trimesh objects
95 The mesh(es) to export.
97 Returns
98 -----------
99 export: str, string of COLLADA format output
100 """
101 import collada
103 meshes = mesh
104 if not isinstance(mesh, (list, tuple, set, np.ndarray)):
105 meshes = [mesh]
107 c = collada.Collada()
108 nodes = []
109 for i, m in enumerate(meshes):
110 # Load uv, colors, materials
111 uv = None
112 colors = None
113 mat = _unparse_material(None)
114 if m.visual.defined:
115 if m.visual.kind == "texture":
116 mat = _unparse_material(m.visual.material)
117 uv = m.visual.uv
118 elif m.visual.kind == "vertex":
119 colors = (m.visual.vertex_colors / 255.0)[:, :3]
120 mat.effect.diffuse = np.array(m.visual.main_color) / 255.0
121 elif m.visual.kind == "face":
122 mat.effect.diffuse = np.array(m.visual.main_color) / 255.0
123 c.effects.append(mat.effect)
124 c.materials.append(mat)
126 # Create geometry object
127 vertices = collada.source.FloatSource(
128 "verts-array", m.vertices.flatten(), ("X", "Y", "Z")
129 )
130 normals = collada.source.FloatSource(
131 "normals-array", m.vertex_normals.flatten(), ("X", "Y", "Z")
132 )
133 input_list = collada.source.InputList()
134 input_list.addInput(0, "VERTEX", "#verts-array")
135 input_list.addInput(1, "NORMAL", "#normals-array")
136 arrays = [vertices, normals]
137 if (uv is not None) and (len(uv) > 0):
138 texcoords = collada.source.FloatSource(
139 "texcoords-array", uv.flatten(), ("U", "V")
140 )
141 input_list.addInput(2, "TEXCOORD", "#texcoords-array")
142 arrays.append(texcoords)
143 if colors is not None:
144 idx = 2
145 if uv:
146 idx = 3
147 colors = collada.source.FloatSource(
148 "colors-array", colors.flatten(), ("R", "G", "B")
149 )
150 input_list.addInput(idx, "COLOR", "#colors-array")
151 arrays.append(colors)
152 geom = collada.geometry.Geometry(c, uuid.uuid4().hex, uuid.uuid4().hex, arrays)
153 indices = np.repeat(m.faces.flatten(), len(arrays))
155 matref = f"material{i}"
156 triset = geom.createTriangleSet(indices, input_list, matref)
157 geom.primitives.append(triset)
158 c.geometries.append(geom)
160 matnode = collada.scene.MaterialNode(matref, mat, inputs=[])
161 geomnode = collada.scene.GeometryNode(geom, [matnode])
162 node = collada.scene.Node(f"node{i}", children=[geomnode])
163 nodes.append(node)
164 scene = collada.scene.Scene("scene", nodes)
165 c.scenes.append(scene)
166 c.scene = scene
168 b = io.BytesIO()
169 c.write(b)
170 b.seek(0)
171 return b.read()
174def _parse_node(
175 node, parent_matrix, material_map, meshes, meshes_count, graph, resolver, metadata
176):
177 """
178 Recursively parse COLLADA scene nodes.
179 """
180 import collada
182 # Parse mesh node
183 if isinstance(node, collada.scene.GeometryNode):
184 geometry = node.geometry
186 # Create local material map from material symbol to actual material
187 local_material_map = {}
188 for mn in node.materials:
189 symbol = mn.symbol
190 m = mn.target
191 if m.id in material_map:
192 local_material_map[symbol] = material_map[m.id]
193 else:
194 local_material_map[symbol] = _parse_material(m, resolver)
196 # Iterate over primitives of geometry
197 for primitive in geometry.primitives:
198 if isinstance(primitive, collada.polylist.Polylist):
199 primitive = primitive.triangleset()
200 if isinstance(primitive, collada.triangleset.TriangleSet):
201 vertex = primitive.vertex
202 if vertex is None:
203 continue
204 vertex_index = primitive.vertex_index
205 vertices = vertex[vertex_index].reshape(len(vertex_index) * 3, 3)
207 # Get normals if present
208 normals = None
209 if primitive.normal is not None:
210 normal = primitive.normal
211 normal_index = primitive.normal_index
212 normals = normal[normal_index].reshape(len(normal_index) * 3, 3)
214 # Get colors if present
215 colors = None
216 s = primitive.sources
217 if "COLOR" in s and len(s["COLOR"]) > 0 and len(primitive.index) > 0:
218 color = s["COLOR"][0][4].data
219 color_index = primitive.index[:, :, s["COLOR"][0][0]]
220 colors = color[color_index].reshape(len(color_index) * 3, -1)
222 faces = np.arange(vertices.shape[0]).reshape(vertices.shape[0] // 3, 3)
224 # Get UV coordinates if possible
225 vis = None
226 if primitive.material in local_material_map:
227 material = copy.copy(local_material_map[primitive.material])
228 uv = None
229 if len(primitive.texcoordset) > 0:
230 texcoord = primitive.texcoordset[0]
231 texcoord_index = primitive.texcoord_indexset[0]
232 uv = texcoord[texcoord_index].reshape(
233 (len(texcoord_index) * 3, 2)
234 )
235 vis = visual.texture.TextureVisuals(uv=uv, material=material)
237 geom_name = unique_name(geometry.id, contains=meshes, counts=meshes_count)
238 meshes[geom_name] = {
239 "vertices": vertices,
240 "faces": faces,
241 "vertex_normals": normals,
242 "vertex_colors": colors,
243 "visual": vis,
244 "metadata": metadata,
245 }
247 graph.append(
248 {
249 "frame_to": geom_name,
250 "matrix": parent_matrix,
251 "geometry": geom_name,
252 }
253 )
255 # recurse down tree for nodes with children
256 elif isinstance(node, collada.scene.Node):
257 if node.children is not None:
258 for child in node.children:
259 # create the new matrix
260 matrix = np.dot(parent_matrix, node.matrix)
261 # parse the child node
262 _parse_node(
263 node=child,
264 parent_matrix=matrix,
265 material_map=material_map,
266 meshes=meshes,
267 meshes_count=meshes_count,
268 graph=graph,
269 resolver=resolver,
270 metadata=metadata,
271 )
273 elif isinstance(node, collada.scene.CameraNode):
274 # TODO: convert collada cameras to trimesh cameras
275 pass
276 elif isinstance(node, collada.scene.LightNode):
277 # TODO: convert collada lights to trimesh lights
278 pass
281def _load_texture(file_name, resolver):
282 """
283 Load a texture from a file into a PIL image.
284 """
285 from PIL import Image
287 file_data = resolver.get(file_name)
288 image = Image.open(util.wrap_as_stream(file_data))
289 return image
292def _parse_material(effect, resolver):
293 """
294 Turn a COLLADA effect into a trimesh material.
295 """
296 import collada
298 # Compute base color
299 baseColorFactor = np.ones(4)
300 baseColorTexture = None
301 if isinstance(effect.diffuse, collada.material.Map):
302 try:
303 baseColorTexture = _load_texture(
304 effect.diffuse.sampler.surface.image.path, resolver
305 )
306 except BaseException:
307 log.debug("unable to load base texture", exc_info=True)
308 elif effect.diffuse is not None:
309 baseColorFactor = effect.diffuse
311 # Compute emission color
312 emissiveFactor = np.zeros(3)
313 emissiveTexture = None
314 if isinstance(effect.emission, collada.material.Map):
315 try:
316 emissiveTexture = _load_texture(
317 effect.diffuse.sampler.surface.image.path, resolver
318 )
319 except BaseException:
320 log.warning("unable to load emissive texture", exc_info=True)
321 elif effect.emission is not None:
322 emissiveFactor = effect.emission[:3]
324 # Compute roughness
325 roughnessFactor = 1.0
326 if (
327 not isinstance(effect.shininess, collada.material.Map)
328 and effect.shininess is not None
329 ):
330 try:
331 shininess_value = float(effect.shininess)
332 roughnessFactor = np.sqrt(2.0 / (2.0 + shininess_value))
333 except (TypeError, ValueError):
334 log.warning(
335 f"Invalid shininess value: {effect.shininess}, using default roughness"
336 )
338 # Compute metallic factor
339 metallicFactor = 0.0
341 # Compute normal texture
342 normalTexture = None
343 if effect.bumpmap is not None:
344 try:
345 normalTexture = _load_texture(
346 effect.bumpmap.sampler.surface.image.path, resolver
347 )
348 except BaseException:
349 log.warning("unable to load bumpmap", exc_info=True)
351 # Compute opacity
352 if effect.transparent is not None and not isinstance(
353 effect.transparent, collada.material.Map
354 ):
355 baseColorFactor = tuple(
356 np.append(baseColorFactor[:3], float(effect.transparent[3]))
357 )
359 return visual.material.PBRMaterial(
360 emissiveFactor=emissiveFactor,
361 emissiveTexture=emissiveTexture,
362 normalTexture=normalTexture,
363 baseColorTexture=baseColorTexture,
364 baseColorFactor=baseColorFactor,
365 metallicFactor=metallicFactor,
366 roughnessFactor=roughnessFactor,
367 )
370def _unparse_material(material):
371 """
372 Turn a trimesh material into a COLLADA material.
373 """
374 import collada
376 # TODO EXPORT TEXTURES
377 if isinstance(material, visual.material.PBRMaterial):
378 diffuse = material.baseColorFactor
379 if diffuse is None:
380 diffuse = np.array([255.0, 255.0, 255.0, 255.0])
381 diffuse = diffuse / 255.0
382 if diffuse is not None:
383 diffuse = list(diffuse)
385 emission = material.emissiveFactor
386 if emission is not None:
387 emission = [float(emission[0]), float(emission[1]), float(emission[2]), 1.0]
389 shininess = material.roughnessFactor
390 if shininess is None:
391 shininess = 1.0
392 if shininess is not None:
393 shininess = 2.0 / shininess**2 - 2.0
395 effect = collada.material.Effect(
396 uuid.uuid4().hex,
397 params=[],
398 shadingtype="phong",
399 diffuse=diffuse,
400 emission=emission,
401 specular=[1.0, 1.0, 1.0, 1.0],
402 shininess=float(shininess),
403 )
404 material = collada.material.Material(uuid.uuid4().hex, "pbrmaterial", effect)
405 else:
406 effect = collada.material.Effect(uuid.uuid4().hex, params=[], shadingtype="phong")
407 material = collada.material.Material(uuid.uuid4().hex, "defaultmaterial", effect)
408 return material
411def load_zae(file_obj, resolver=None, **kwargs):
412 """
413 Load a ZAE file, which is just a zipped DAE file.
415 Parameters
416 -------------
417 file_obj : file object
418 Contains ZAE data
419 resolver : trimesh.visual.Resolver
420 Resolver to load additional assets
421 kwargs : dict
422 Passed to load_collada
424 Returns
425 ------------
426 loaded : dict
427 Results of loading
428 """
430 # a dict, {file name : file object}
431 archive = util.decompress(file_obj, file_type="zip")
433 # load the first file with a .dae extension
434 file_name = next(i for i in archive.keys() if i.lower().endswith(".dae"))
436 # a resolver so the loader can load textures / etc
437 resolver = visual.resolvers.ZipResolver(archive)
439 # run the regular collada loader
440 loaded = load_collada(archive[file_name], resolver=resolver, **kwargs)
441 return loaded
444# only provide loaders if `pycollada` is installed
445_collada_loaders = {}
446_collada_exporters = {}
447if util.has_module("collada"):
448 _collada_loaders["dae"] = load_collada
449 _collada_loaders["zae"] = load_zae
450 _collada_exporters["dae"] = export_collada
451else:
452 # store an exception to raise later
453 from ..exceptions import ExceptionWrapper
455 _exc = ExceptionWrapper(ImportError("missing `pip install pycollada`"))
456 _collada_loaders.update({"dae": _exc, "zae": _exc})
457 _collada_exporters["dae"] = _exc