import copy import io import uuid import numpy as np from .. import util, visual from ..constants import log from ..util import unique_name _EYE = np.eye(4) _EYE.flags.writeable = False def load_collada(file_obj, resolver=None, ignore_broken=True, **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 ignore_broken: bool Ignores broken references during loading: [collada.common.DaeUnsupportedError, collada.common.DaeBrokenRefError] kwargs : ** Passed to trimesh.Trimesh.__init__ Returns ------- loaded : list of dict kwargs for Trimesh constructor """ import collada if ignore_broken: ignores = [ collada.common.DaeError, collada.common.DaeIncompleteError, collada.common.DaeMalformedError, collada.common.DaeBrokenRefError, collada.common.DaeUnsupportedError, collada.common.DaeIncompleteError, ] else: ignores = None # load scene using pycollada c = collada.Collada(file_obj, ignore=ignores) # 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) unit = c.assetInfo.unitmeter if unit is None or np.isclose(unit, 1.0): metadata = {"units": "meters"} else: metadata = {"units": f"{unit} * meters"} # name : kwargs meshes = {} # increments to enable `unique_name` to avoid n^2 behavior meshes_count = {} # list of dict graph = [] for node in c.scene.nodes: _parse_node( node=node, parent_matrix=_EYE, material_map=material_map, meshes=meshes, meshes_count=meshes_count, graph=graph, resolver=resolver, metadata=metadata, ) return {"class": "Scene", "graph": graph, "geometry": meshes} 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 """ import collada 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] mat.effect.diffuse = np.array(m.visual.main_color) / 255.0 elif m.visual.kind == "face": mat.effect.diffuse = np.array(m.visual.main_color) / 255.0 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) and (len(uv) > 0): 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 = f"material{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(f"node{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, meshes_count, graph, resolver, metadata ): """ Recursively parse COLLADA scene nodes. """ import collada # 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 primitive in geometry.primitives: if isinstance(primitive, collada.polylist.Polylist): primitive = primitive.triangleset() if isinstance(primitive, collada.triangleset.TriangleSet): vertex = primitive.vertex if vertex is None: continue 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, -1) 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) geom_name = unique_name(geometry.id, contains=meshes, counts=meshes_count) meshes[geom_name] = { "vertices": vertices, "faces": faces, "vertex_normals": normals, "vertex_colors": colors, "visual": vis, "metadata": metadata, } graph.append( { "frame_to": geom_name, "matrix": parent_matrix, "geometry": geom_name, } ) # 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, meshes_count=meshes_count, graph=graph, resolver=resolver, metadata=metadata, ) 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. """ import collada # 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.debug("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 ): try: shininess_value = float(effect.shininess) roughnessFactor = np.sqrt(2.0 / (2.0 + shininess_value)) except (TypeError, ValueError): log.warning( f"Invalid shininess value: {effect.shininess}, using default roughness" ) # 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(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. """ import collada # TODO EXPORT TEXTURES if isinstance(material, visual.material.PBRMaterial): diffuse = material.baseColorFactor if diffuse is None: diffuse = np.array([255.0, 255.0, 255.0, 255.0]) diffuse = diffuse / 255.0 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 None: shininess = 1.0 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 util.has_module("collada"): _collada_loaders["dae"] = load_collada _collada_loaders["zae"] = load_zae _collada_exporters["dae"] = export_collada else: # store an exception to raise later from ..exceptions import ExceptionWrapper _exc = ExceptionWrapper(ImportError("missing `pip install pycollada`")) _collada_loaders.update({"dae": _exc, "zae": _exc}) _collada_exporters["dae"] = _exc