1import io
2import copy
3import uuid
4
5import numpy as np
6
7try:
8    # pip install pycollada
9    import collada
10except BaseException:
11    collada = None
12
13try:
14    import PIL.Image
15except ImportError:
16    pass
17
18from .. import util
19from .. import visual
20
21from ..constants import log
22
23
24def load_collada(file_obj, resolver=None, **kwargs):
25    """
26    Load a COLLADA (.dae) file into a list of trimesh kwargs.
27
28    Parameters
29    ----------
30    file_obj : file object
31      Containing a COLLADA file
32    resolver : trimesh.visual.Resolver or None
33      For loading referenced files, like texture images
34    kwargs : **
35      Passed to trimesh.Trimesh.__init__
36
37    Returns
38    -------
39    loaded : list of dict
40      kwargs for Trimesh constructor
41    """
42    # load scene using pycollada
43    c = collada.Collada(file_obj)
44
45    # Create material map from Material ID to trimesh material
46    material_map = {}
47    for m in c.materials:
48        effect = m.effect
49        material_map[m.id] = _parse_material(effect, resolver)
50
51    # name : kwargs
52    meshes = {}
53    # list of dict
54    graph = []
55
56    for node in c.scene.nodes:
57        _parse_node(node=node,
58                    parent_matrix=np.eye(4),
59                    material_map=material_map,
60                    meshes=meshes,
61                    graph=graph,
62                    resolver=resolver)
63
64    # create kwargs for load_kwargs
65    result = {'class': 'Scene',
66              'graph': graph,
67              'geometry': meshes}
68
69    return result
70
71
72def export_collada(mesh, **kwargs):
73    """
74    Export a mesh or a list of meshes as a COLLADA .dae file.
75
76    Parameters
77    -----------
78    mesh: Trimesh object or list of Trimesh objects
79        The mesh(es) to export.
80
81    Returns
82    -----------
83    export: str, string of COLLADA format output
84    """
85    meshes = mesh
86    if not isinstance(mesh, (list, tuple, set, np.ndarray)):
87        meshes = [mesh]
88
89    c = collada.Collada()
90    nodes = []
91    for i, m in enumerate(meshes):
92
93        # Load uv, colors, materials
94        uv = None
95        colors = None
96        mat = _unparse_material(None)
97        if m.visual.defined:
98            if m.visual.kind == 'texture':
99                mat = _unparse_material(m.visual.material)
100                uv = m.visual.uv
101            elif m.visual.kind == 'vertex':
102                colors = (m.visual.vertex_colors / 255.0)[:, :3]
103        c.effects.append(mat.effect)
104        c.materials.append(mat)
105
106        # Create geometry object
107        vertices = collada.source.FloatSource(
108            'verts-array', m.vertices.flatten(), ('X', 'Y', 'Z'))
109        normals = collada.source.FloatSource(
110            'normals-array', m.vertex_normals.flatten(), ('X', 'Y', 'Z'))
111        input_list = collada.source.InputList()
112        input_list.addInput(0, 'VERTEX', '#verts-array')
113        input_list.addInput(1, 'NORMAL', '#normals-array')
114        arrays = [vertices, normals]
115        if uv is not None:
116            texcoords = collada.source.FloatSource(
117                'texcoords-array', uv.flatten(), ('U', 'V'))
118            input_list.addInput(2, 'TEXCOORD', '#texcoords-array')
119            arrays.append(texcoords)
120        if colors is not None:
121            idx = 2
122            if uv:
123                idx = 3
124            colors = collada.source.FloatSource('colors-array',
125                                                colors.flatten(), ('R', 'G', 'B'))
126            input_list.addInput(idx, 'COLOR', '#colors-array')
127            arrays.append(colors)
128        geom = collada.geometry.Geometry(
129            c, uuid.uuid4().hex, uuid.uuid4().hex, arrays
130        )
131        indices = np.repeat(m.faces.flatten(), len(arrays))
132
133        matref = u'material{}'.format(i)
134        triset = geom.createTriangleSet(indices, input_list, matref)
135        geom.primitives.append(triset)
136        c.geometries.append(geom)
137
138        matnode = collada.scene.MaterialNode(matref, mat, inputs=[])
139        geomnode = collada.scene.GeometryNode(geom, [matnode])
140        node = collada.scene.Node(u'node{}'.format(i), children=[geomnode])
141        nodes.append(node)
142    scene = collada.scene.Scene('scene', nodes)
143    c.scenes.append(scene)
144    c.scene = scene
145
146    b = io.BytesIO()
147    c.write(b)
148    b.seek(0)
149    return b.read()
150
151
152def _parse_node(node,
153                parent_matrix,
154                material_map,
155                meshes,
156                graph,
157                resolver=None):
158    """
159    Recursively parse COLLADA scene nodes.
160    """
161
162    # Parse mesh node
163    if isinstance(node, collada.scene.GeometryNode):
164        geometry = node.geometry
165
166        # Create local material map from material symbol to actual material
167        local_material_map = {}
168        for mn in node.materials:
169            symbol = mn.symbol
170            m = mn.target
171            if m.id in material_map:
172                local_material_map[symbol] = material_map[m.id]
173            else:
174                local_material_map[symbol] = _parse_material(m, resolver)
175
176        # Iterate over primitives of geometry
177        for i, primitive in enumerate(geometry.primitives):
178            if isinstance(primitive, collada.polylist.Polylist):
179                primitive = primitive.triangleset()
180            if isinstance(primitive, collada.triangleset.TriangleSet):
181                vertex = primitive.vertex
182                vertex_index = primitive.vertex_index
183                vertices = vertex[vertex_index].reshape(
184                    len(vertex_index) * 3, 3)
185
186                # Get normals if present
187                normals = None
188                if primitive.normal is not None:
189                    normal = primitive.normal
190                    normal_index = primitive.normal_index
191                    normals = normal[normal_index].reshape(
192                        len(normal_index) * 3, 3)
193
194                # Get colors if present
195                colors = None
196                s = primitive.sources
197                if ('COLOR' in s and len(s['COLOR'])
198                        > 0 and len(primitive.index) > 0):
199                    color = s['COLOR'][0][4].data
200                    color_index = primitive.index[:, :, s['COLOR'][0][0]]
201                    colors = color[color_index].reshape(
202                        len(color_index) * 3, 3)
203
204                faces = np.arange(
205                    vertices.shape[0]).reshape(
206                    vertices.shape[0] // 3, 3)
207
208                # Get UV coordinates if possible
209                vis = None
210                if primitive.material in local_material_map:
211                    material = copy.copy(
212                        local_material_map[primitive.material])
213                    uv = None
214                    if len(primitive.texcoordset) > 0:
215                        texcoord = primitive.texcoordset[0]
216                        texcoord_index = primitive.texcoord_indexset[0]
217                        uv = texcoord[texcoord_index].reshape(
218                            (len(texcoord_index) * 3, 2))
219                    vis = visual.texture.TextureVisuals(
220                        uv=uv, material=material)
221
222                primid = u'{}.{}'.format(geometry.id, i)
223                meshes[primid] = {
224                    'vertices': vertices,
225                    'faces': faces,
226                    'vertex_normals': normals,
227                    'vertex_colors': colors,
228                    'visual': vis}
229
230                graph.append({'frame_to': primid,
231                              'matrix': parent_matrix,
232                              'geometry': primid})
233
234    # recurse down tree for nodes with children
235    elif isinstance(node, collada.scene.Node):
236        if node.children is not None:
237            for child in node.children:
238                # create the new matrix
239                matrix = np.dot(parent_matrix, node.matrix)
240                # parse the child node
241                _parse_node(
242                    node=child,
243                    parent_matrix=matrix,
244                    material_map=material_map,
245                    meshes=meshes,
246                    graph=graph,
247                    resolver=resolver)
248
249    elif isinstance(node, collada.scene.CameraNode):
250        # TODO: convert collada cameras to trimesh cameras
251        pass
252    elif isinstance(node, collada.scene.LightNode):
253        # TODO: convert collada lights to trimesh lights
254        pass
255
256
257def _load_texture(file_name, resolver):
258    """
259    Load a texture from a file into a PIL image.
260    """
261    file_data = resolver.get(file_name)
262    image = PIL.Image.open(util.wrap_as_stream(file_data))
263    return image
264
265
266def _parse_material(effect, resolver):
267    """
268    Turn a COLLADA effect into a trimesh material.
269    """
270
271    # Compute base color
272    baseColorFactor = np.ones(4)
273    baseColorTexture = None
274    if isinstance(effect.diffuse, collada.material.Map):
275        try:
276            baseColorTexture = _load_texture(
277                effect.diffuse.sampler.surface.image.path, resolver)
278        except BaseException:
279            log.warning('unable to load base texture',
280                        exc_info=True)
281    elif effect.diffuse is not None:
282        baseColorFactor = effect.diffuse
283
284    # Compute emission color
285    emissiveFactor = np.zeros(3)
286    emissiveTexture = None
287    if isinstance(effect.emission, collada.material.Map):
288        try:
289            emissiveTexture = _load_texture(
290                effect.diffuse.sampler.surface.image.path, resolver)
291        except BaseException:
292            log.warning('unable to load emissive texture',
293                        exc_info=True)
294    elif effect.emission is not None:
295        emissiveFactor = effect.emission[:3]
296
297    # Compute roughness
298    roughnessFactor = 1.0
299    if (not isinstance(effect.shininess, collada.material.Map)
300            and effect.shininess is not None):
301        roughnessFactor = np.sqrt(2.0 / (2.0 + effect.shininess))
302
303    # Compute metallic factor
304    metallicFactor = 0.0
305
306    # Compute normal texture
307    normalTexture = None
308    if effect.bumpmap is not None:
309        try:
310            normalTexture = _load_texture(
311                effect.bumpmap.sampler.surface.image.path, resolver)
312        except BaseException:
313            log.warning('unable to load bumpmap',
314                        exc_info=True)
315
316    return visual.material.PBRMaterial(
317        emissiveFactor=emissiveFactor,
318        emissiveTexture=emissiveTexture,
319        normalTexture=normalTexture,
320        baseColorTexture=baseColorTexture,
321        baseColorFactor=baseColorFactor,
322        metallicFactor=metallicFactor,
323        roughnessFactor=roughnessFactor
324    )
325
326
327def _unparse_material(material):
328    """
329    Turn a trimesh material into a COLLADA material.
330    """
331    # TODO EXPORT TEXTURES
332    if isinstance(material, visual.material.PBRMaterial):
333        diffuse = material.baseColorFactor
334        if diffuse is not None:
335            diffuse = list(diffuse)
336
337        emission = material.emissiveFactor
338        if emission is not None:
339            emission = [float(emission[0]), float(emission[1]),
340                        float(emission[2]), 1.0]
341
342        shininess = material.roughnessFactor
343        if shininess is not None:
344            shininess = 2.0 / shininess**2 - 2.0
345
346        effect = collada.material.Effect(
347            uuid.uuid4().hex, params=[], shadingtype='phong',
348            diffuse=diffuse, emission=emission,
349            specular=[1.0, 1.0, 1.0, 1.0], shininess=float(shininess)
350        )
351        material = collada.material.Material(
352            uuid.uuid4().hex, 'pbrmaterial', effect
353        )
354    else:
355        effect = collada.material.Effect(
356            uuid.uuid4().hex, params=[], shadingtype='phong'
357        )
358        material = collada.material.Material(
359            uuid.uuid4().hex, 'defaultmaterial', effect
360        )
361    return material
362
363
364def load_zae(file_obj, resolver=None, **kwargs):
365    """
366    Load a ZAE file, which is just a zipped DAE file.
367
368    Parameters
369    -------------
370    file_obj : file object
371      Contains ZAE data
372    resolver : trimesh.visual.Resolver
373      Resolver to load additional assets
374    kwargs : dict
375      Passed to load_collada
376
377    Returns
378    ------------
379    loaded : dict
380      Results of loading
381    """
382
383    # a dict, {file name : file object}
384    archive = util.decompress(file_obj,
385                              file_type='zip')
386
387    # load the first file with a .dae extension
388    file_name = next(i for i in archive.keys()
389                     if i.lower().endswith('.dae'))
390
391    # a resolver so the loader can load textures / etc
392    resolver = visual.resolvers.ZipResolver(archive)
393
394    # run the regular collada loader
395    loaded = load_collada(archive[file_name],
396                          resolver=resolver,
397                          **kwargs)
398    return loaded
399
400
401# only provide loaders if `pycollada` is installed
402_collada_loaders = {}
403_collada_exporters = {}
404if collada is not None:
405    _collada_loaders['dae'] = load_collada
406    _collada_loaders['zae'] = load_zae
407    _collada_exporters['dae'] = export_collada
408