import io import copy import uuid import numpy as np try: # pip install pycollada import collada except BaseException: collada = None 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 = u'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(u'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 = u'{}.{}'.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. """ from PIL import Image file_data = resolver.get(file_name) image = 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) # Compute opacity if (effect.transparent is not None and not isinstance(effect.transparent, collada.material.Map)): baseColorFactor = tuple( np.append(baseColorFactor[:3], float(int(255 * effect.transparent[3])))) 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