/
dae.py
407 lines (345 loc) · 12.7 KB
/
dae.py
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
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
import io
import copy
import uuid
import numpy as np
try:
# pip install pycollada
import collada
except BaseException:
collada = None
try:
import PIL.Image
except ImportError:
pass
from .. import util
from .. import visual
from ..constants import log
def load_collada(file_obj, resolver=None, **kwargs):
"""
Load a COLLADA (.dae) file into a list of trimesh kwargs.
Parameters
----------
file_obj : file object
Containing a COLLADA file
resolver : trimesh.visual.Resolver or None
For loading referenced files, like texture images
kwargs : **
Passed to trimesh.Trimesh.__init__
Returns
-------
loaded : list of dict
kwargs for Trimesh constructor
"""
# load scene using pycollada
c = collada.Collada(file_obj)
# Create material map from Material ID to trimesh material
material_map = {}
for m in c.materials:
effect = m.effect
material_map[m.id] = _parse_material(effect, resolver)
# name : kwargs
meshes = {}
# list of dict
graph = []
for node in c.scene.nodes:
_parse_node(node=node,
parent_matrix=np.eye(4),
material_map=material_map,
meshes=meshes,
graph=graph,
resolver=resolver)
# create kwargs for load_kwargs
result = {'class': 'Scene',
'graph': graph,
'geometry': meshes}
return result
def export_collada(mesh, **kwargs):
"""
Export a mesh or a list of meshes as a COLLADA .dae file.
Parameters
-----------
mesh: Trimesh object or list of Trimesh objects
The mesh(es) to export.
Returns
-----------
export: str, string of COLLADA format output
"""
meshes = mesh
if not isinstance(mesh, (list, tuple, set, np.ndarray)):
meshes = [mesh]
c = collada.Collada()
nodes = []
for i, m in enumerate(meshes):
# Load uv, colors, materials
uv = None
colors = None
mat = _unparse_material(None)
if m.visual.defined:
if m.visual.kind == 'texture':
mat = _unparse_material(m.visual.material)
uv = m.visual.uv
elif m.visual.kind == 'vertex':
colors = (m.visual.vertex_colors / 255.0)[:, :3]
c.effects.append(mat.effect)
c.materials.append(mat)
# Create geometry object
vertices = collada.source.FloatSource(
'verts-array', m.vertices.flatten(), ('X', 'Y', 'Z'))
normals = collada.source.FloatSource(
'normals-array', m.vertex_normals.flatten(), ('X', 'Y', 'Z'))
input_list = collada.source.InputList()
input_list.addInput(0, 'VERTEX', '#verts-array')
input_list.addInput(1, 'NORMAL', '#normals-array')
arrays = [vertices, normals]
if uv is not None:
texcoords = collada.source.FloatSource(
'texcoords-array', uv.flatten(), ('U', 'V'))
input_list.addInput(2, 'TEXCOORD', '#texcoords-array')
arrays.append(texcoords)
if colors is not None:
idx = 2
if uv:
idx = 3
colors = collada.source.FloatSource('colors-array',
colors.flatten(), ('R', 'G', 'B'))
input_list.addInput(idx, 'COLOR', '#colors-array')
arrays.append(colors)
geom = collada.geometry.Geometry(
c, uuid.uuid4().hex, uuid.uuid4().hex, arrays
)
indices = np.repeat(m.faces.flatten(), len(arrays))
matref = 'material{}'.format(i)
triset = geom.createTriangleSet(indices, input_list, matref)
geom.primitives.append(triset)
c.geometries.append(geom)
matnode = collada.scene.MaterialNode(matref, mat, inputs=[])
geomnode = collada.scene.GeometryNode(geom, [matnode])
node = collada.scene.Node('node{}'.format(i), children=[geomnode])
nodes.append(node)
scene = collada.scene.Scene('scene', nodes)
c.scenes.append(scene)
c.scene = scene
b = io.BytesIO()
c.write(b)
b.seek(0)
return b.read()
def _parse_node(node,
parent_matrix,
material_map,
meshes,
graph,
resolver=None):
"""
Recursively parse COLLADA scene nodes.
"""
# Parse mesh node
if isinstance(node, collada.scene.GeometryNode):
geometry = node.geometry
# Create local material map from material symbol to actual material
local_material_map = {}
for mn in node.materials:
symbol = mn.symbol
m = mn.target
if m.id in material_map:
local_material_map[symbol] = material_map[m.id]
else:
local_material_map[symbol] = _parse_material(m, resolver)
# Iterate over primitives of geometry
for i, primitive in enumerate(geometry.primitives):
if isinstance(primitive, collada.polylist.Polylist):
primitive = primitive.triangleset()
if isinstance(primitive, collada.triangleset.TriangleSet):
vertex = primitive.vertex
vertex_index = primitive.vertex_index
vertices = vertex[vertex_index].reshape(
len(vertex_index) * 3, 3)
# Get normals if present
normals = None
if primitive.normal is not None:
normal = primitive.normal
normal_index = primitive.normal_index
normals = normal[normal_index].reshape(
len(normal_index) * 3, 3)
# Get colors if present
colors = None
s = primitive.sources
if ('COLOR' in s and len(s['COLOR'])
> 0 and len(primitive.index) > 0):
color = s['COLOR'][0][4].data
color_index = primitive.index[:, :, s['COLOR'][0][0]]
colors = color[color_index].reshape(
len(color_index) * 3, 3)
faces = np.arange(
vertices.shape[0]).reshape(
vertices.shape[0] // 3, 3)
# Get UV coordinates if possible
vis = None
if primitive.material in local_material_map:
material = copy.copy(
local_material_map[primitive.material])
uv = None
if len(primitive.texcoordset) > 0:
texcoord = primitive.texcoordset[0]
texcoord_index = primitive.texcoord_indexset[0]
uv = texcoord[texcoord_index].reshape(
(len(texcoord_index) * 3, 2))
vis = visual.texture.TextureVisuals(
uv=uv, material=material)
primid = '{}.{}'.format(geometry.id, i)
meshes[primid] = {
'vertices': vertices,
'faces': faces,
'vertex_normals': normals,
'vertex_colors': colors,
'visual': vis}
graph.append({'frame_to': primid,
'matrix': parent_matrix,
'geometry': primid})
# recurse down tree for nodes with children
elif isinstance(node, collada.scene.Node):
if node.children is not None:
for child in node.children:
# create the new matrix
matrix = np.dot(parent_matrix, node.matrix)
# parse the child node
_parse_node(
node=child,
parent_matrix=matrix,
material_map=material_map,
meshes=meshes,
graph=graph,
resolver=resolver)
elif isinstance(node, collada.scene.CameraNode):
# TODO: convert collada cameras to trimesh cameras
pass
elif isinstance(node, collada.scene.LightNode):
# TODO: convert collada lights to trimesh lights
pass
def _load_texture(file_name, resolver):
"""
Load a texture from a file into a PIL image.
"""
file_data = resolver.get(file_name)
image = PIL.Image.open(util.wrap_as_stream(file_data))
return image
def _parse_material(effect, resolver):
"""
Turn a COLLADA effect into a trimesh material.
"""
# Compute base color
baseColorFactor = np.ones(4)
baseColorTexture = None
if isinstance(effect.diffuse, collada.material.Map):
try:
baseColorTexture = _load_texture(
effect.diffuse.sampler.surface.image.path, resolver)
except BaseException:
log.warning('unable to load base texture',
exc_info=True)
elif effect.diffuse is not None:
baseColorFactor = effect.diffuse
# Compute emission color
emissiveFactor = np.zeros(3)
emissiveTexture = None
if isinstance(effect.emission, collada.material.Map):
try:
emissiveTexture = _load_texture(
effect.diffuse.sampler.surface.image.path, resolver)
except BaseException:
log.warning('unable to load emissive texture',
exc_info=True)
elif effect.emission is not None:
emissiveFactor = effect.emission[:3]
# Compute roughness
roughnessFactor = 1.0
if (not isinstance(effect.shininess, collada.material.Map)
and effect.shininess is not None):
roughnessFactor = np.sqrt(2.0 / (2.0 + effect.shininess))
# Compute metallic factor
metallicFactor = 0.0
# Compute normal texture
normalTexture = None
if effect.bumpmap is not None:
try:
normalTexture = _load_texture(
effect.bumpmap.sampler.surface.image.path, resolver)
except BaseException:
log.warning('unable to load bumpmap',
exc_info=True)
return visual.material.PBRMaterial(
emissiveFactor=emissiveFactor,
emissiveTexture=emissiveTexture,
normalTexture=normalTexture,
baseColorTexture=baseColorTexture,
baseColorFactor=baseColorFactor,
metallicFactor=metallicFactor,
roughnessFactor=roughnessFactor
)
def _unparse_material(material):
"""
Turn a trimesh material into a COLLADA material.
"""
# TODO EXPORT TEXTURES
if isinstance(material, visual.material.PBRMaterial):
diffuse = material.baseColorFactor
if diffuse is not None:
diffuse = list(diffuse)
emission = material.emissiveFactor
if emission is not None:
emission = [float(emission[0]), float(emission[1]),
float(emission[2]), 1.0]
shininess = material.roughnessFactor
if shininess is not None:
shininess = 2.0 / shininess**2 - 2.0
effect = collada.material.Effect(
uuid.uuid4().hex, params=[], shadingtype='phong',
diffuse=diffuse, emission=emission,
specular=[1.0, 1.0, 1.0, 1.0], shininess=float(shininess)
)
material = collada.material.Material(
uuid.uuid4().hex, 'pbrmaterial', effect
)
else:
effect = collada.material.Effect(
uuid.uuid4().hex, params=[], shadingtype='phong'
)
material = collada.material.Material(
uuid.uuid4().hex, 'defaultmaterial', effect
)
return material
def load_zae(file_obj, resolver=None, **kwargs):
"""
Load a ZAE file, which is just a zipped DAE file.
Parameters
-------------
file_obj : file object
Contains ZAE data
resolver : trimesh.visual.Resolver
Resolver to load additional assets
kwargs : dict
Passed to load_collada
Returns
------------
loaded : dict
Results of loading
"""
# a dict, {file name : file object}
archive = util.decompress(file_obj,
file_type='zip')
# load the first file with a .dae extension
file_name = next(i for i in archive.keys()
if i.lower().endswith('.dae'))
# a resolver so the loader can load textures / etc
resolver = visual.resolvers.ZipResolver(archive)
# run the regular collada loader
loaded = load_collada(archive[file_name],
resolver=resolver,
**kwargs)
return loaded
# only provide loaders if `pycollada` is installed
_collada_loaders = {}
_collada_exporters = {}
if collada is not None:
_collada_loaders['dae'] = load_collada
_collada_loaders['zae'] = load_zae
_collada_exporters['dae'] = export_collada