281 lines
7.4 KiB
Python
281 lines
7.4 KiB
Python
import copy
|
|
import tempfile
|
|
import collections
|
|
|
|
import numpy as np
|
|
|
|
from .. import util
|
|
|
|
|
|
def load_pyassimp(file_obj,
|
|
file_type=None,
|
|
resolver=None,
|
|
**kwargs):
|
|
"""
|
|
Use the pyassimp library to load a mesh from a file object
|
|
and type or file name if file_obj is a string
|
|
|
|
Parameters
|
|
---------
|
|
file_obj: str, or file object
|
|
File path or object containing mesh data
|
|
file_type : str
|
|
File extension, aka 'stl'
|
|
resolver : trimesh.visual.resolvers.Resolver
|
|
Used to load referenced data (like texture files)
|
|
kwargs : dict
|
|
Passed through to mesh constructor
|
|
|
|
Returns
|
|
---------
|
|
scene : trimesh.Scene
|
|
Native trimesh copy of assimp scene
|
|
"""
|
|
|
|
def LP_to_TM(lp):
|
|
# try to get the vertex colors attribute
|
|
colors = (np.reshape(lp.colors, (-1, 4))
|
|
[:, :3] * 255).round().astype(np.uint8)
|
|
# If no vertex colors, try to extract them from the material
|
|
if len(colors) == 0:
|
|
if 'diffuse' in lp.material.properties.keys():
|
|
colors = np.array(lp.material.properties['diffuse'])
|
|
|
|
# pass kwargs through to mesh constructor
|
|
mesh_kwargs = copy.deepcopy(kwargs)
|
|
# add data from the LP_Mesh
|
|
mesh_kwargs.update({'vertices': lp.vertices,
|
|
'vertex_normals': lp.normals,
|
|
'vertex_colors': colors,
|
|
'faces': lp.faces})
|
|
|
|
return mesh_kwargs
|
|
|
|
# did we open the file inside this function
|
|
opened = False
|
|
# not a file object
|
|
if not hasattr(file_obj, 'read'):
|
|
# if there is no read attribute
|
|
# we assume we've been passed a file name
|
|
file_type = (str(file_obj).split('.')[-1]).lower()
|
|
file_obj = open(file_obj, 'rb')
|
|
opened = True
|
|
# we need files to be bytes
|
|
elif not hasattr(file_obj, 'mode') or file_obj.mode != 'rb':
|
|
# assimp will crash on anything that isn't binary
|
|
# so if we have a text mode file or anything else
|
|
# grab the data, encode as bytes, and then use stream
|
|
data = file_obj.read()
|
|
if hasattr(data, 'encode'):
|
|
data = data.encode('utf-8')
|
|
file_obj = util.wrap_as_stream(data)
|
|
|
|
# load the scene using pyassimp
|
|
scene = pyassimp.load(file_obj,
|
|
file_type=file_type)
|
|
|
|
# save a record of mesh names used so we
|
|
# don't have to do queries on mesh_id.values()
|
|
mesh_names = set()
|
|
# save a mapping for {id(mesh) : name}
|
|
mesh_id = {}
|
|
# save results as {name : Trimesh}
|
|
meshes = {}
|
|
# loop through scene LPMesh objects
|
|
for m in scene.meshes:
|
|
# skip meshes without tri/quad faces
|
|
if m.faces.shape[1] not in [3, 4]:
|
|
continue
|
|
# if this mesh has the name of an existing mesh
|
|
if m.name in mesh_names:
|
|
# make it the name plus the unique ID of the object
|
|
name = m.name + str(id(m))
|
|
else:
|
|
# otherwise just use the name it calls itself by
|
|
name = m.name
|
|
|
|
# save the name to mark as consumed
|
|
mesh_names.add(name)
|
|
# save the id:name mapping
|
|
mesh_id[id(m)] = name
|
|
# convert the mesh to a trimesh object
|
|
meshes[name] = LP_to_TM(m)
|
|
|
|
# now go through and collect the transforms from the scene
|
|
# we are going to save them as a list of dict kwargs
|
|
transforms = []
|
|
# nodes as (parent, node) tuples
|
|
# use deque so we can pop from both ends
|
|
queue = collections.deque(
|
|
[('world', n) for
|
|
n in scene.rootnode.children])
|
|
|
|
# consume the queue
|
|
while len(queue) > 0:
|
|
# parent name, node object
|
|
parent, node = queue.pop()
|
|
|
|
# assimp uses weirdly duplicate node names
|
|
# object ID's are actually unique and consistent
|
|
node_name = id(node)
|
|
transforms.append({'frame_from': parent,
|
|
'frame_to': node_name,
|
|
'matrix': node.transformation})
|
|
|
|
# loop through meshes this node uses
|
|
# note that they are the SAME objects as converted
|
|
# above so we can find their reference using id()
|
|
for m in node.meshes:
|
|
if id(m) not in mesh_id:
|
|
continue
|
|
|
|
# create kwargs for graph.update
|
|
edge = {'frame_from': node_name,
|
|
'frame_to': str(id(m)) + str(node_name),
|
|
'matrix': np.eye(4),
|
|
'geometry': mesh_id[id(m)]}
|
|
transforms.append(edge)
|
|
|
|
# add any children to the queue to be visited
|
|
for child in node.children:
|
|
queue.appendleft((node_name, child))
|
|
|
|
# release the loaded scene
|
|
pyassimp.release(scene)
|
|
|
|
# if we opened the file in this function close it
|
|
if opened:
|
|
file_obj.close()
|
|
|
|
# create kwargs for trimesh.exchange.load.load_kwargs
|
|
result = {'class': 'Scene',
|
|
'geometry': meshes,
|
|
'graph': transforms,
|
|
'base_frame': 'world'}
|
|
|
|
return result
|
|
|
|
|
|
def load_cyassimp(file_obj,
|
|
file_type=None,
|
|
resolver=None,
|
|
**kwargs):
|
|
"""
|
|
Load a file using the cyassimp bindings.
|
|
|
|
The easiest way to install these is with conda:
|
|
conda install -c menpo/label/master cyassimp
|
|
|
|
Parameters
|
|
---------
|
|
file_obj: str, or file object
|
|
File path or object containing mesh data
|
|
file_type : str
|
|
File extension, aka 'stl'
|
|
resolver : trimesh.visual.resolvers.Resolver
|
|
Used to load referenced data (like texture files)
|
|
kwargs : dict
|
|
Passed through to mesh constructor
|
|
|
|
Returns
|
|
---------
|
|
meshes : (n,) list of dict
|
|
Contain kwargs for Trimesh constructor
|
|
"""
|
|
|
|
if hasattr(file_obj, 'read'):
|
|
# if it has a read attribute it is probably a file object
|
|
with tempfile.NamedTemporaryFile(
|
|
suffix=str(file_type)) as file_temp:
|
|
|
|
file_temp.write(file_obj.read())
|
|
# file name should be bytes
|
|
scene = cyassimp.AIImporter(
|
|
file_temp.name.encode('utf-8'))
|
|
scene.build_scene()
|
|
else:
|
|
scene = cyassimp.AIImporter(file_obj.encode('utf-8'))
|
|
scene.build_scene()
|
|
|
|
meshes = []
|
|
for m in scene.meshes:
|
|
mesh_kwargs = kwargs.copy()
|
|
mesh_kwargs.update({'vertices': m.points,
|
|
'faces': m.trilist})
|
|
meshes.append(mesh_kwargs)
|
|
|
|
if len(meshes) == 1:
|
|
return meshes[0]
|
|
return meshes
|
|
|
|
|
|
_assimp_formats = [
|
|
'fbx',
|
|
'dae',
|
|
'gltf',
|
|
'glb',
|
|
'blend',
|
|
'3ds',
|
|
'ase',
|
|
'obj',
|
|
'ifc',
|
|
'xgl',
|
|
'zgl',
|
|
'ply',
|
|
'lwo',
|
|
'lws',
|
|
'lxo',
|
|
'stl',
|
|
'x',
|
|
'ac',
|
|
'ms3d',
|
|
'cob',
|
|
'scn',
|
|
'bvh',
|
|
'csm',
|
|
'xml',
|
|
'irrmesh',
|
|
'irr',
|
|
'mdl',
|
|
'md2',
|
|
'md3',
|
|
'pk3',
|
|
'mdc',
|
|
'md5',
|
|
'smd',
|
|
'vta',
|
|
'ogex',
|
|
'3d',
|
|
'b3d',
|
|
'q3d',
|
|
'.q3s',
|
|
'nff',
|
|
'off',
|
|
'raw',
|
|
'ter',
|
|
'3dgs',
|
|
'mdl',
|
|
'hmp',
|
|
'ndo']
|
|
_assimp_loaders = {}
|
|
|
|
|
|
# try importing both assimp bindings but prefer cyassimp
|
|
loader = None
|
|
try:
|
|
import pyassimp
|
|
loader = load_pyassimp
|
|
except BaseException:
|
|
pass
|
|
|
|
try:
|
|
import cyassimp
|
|
loader = load_cyassimp
|
|
except BaseException:
|
|
pass
|
|
|
|
|
|
if loader:
|
|
_assimp_loaders.update(zip(_assimp_formats,
|
|
[loader] * len(_assimp_formats)))
|