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