OpenMotionLab/MotionGPT
118
1"""Scenes, conforming to the glTF 2.0 standards as specified in2https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-scene3 4Author: Matthew Matl5"""6import numpy as np7import networkx as nx8import trimesh9 10from .mesh import Mesh11from .camera import Camera12from .light import Light, PointLight, DirectionalLight, SpotLight13from .node import Node14from .utils import format_color_vector15 16 17class Scene(object):18 """A hierarchical scene graph.19 20 Parameters21 ----------22 nodes : list of :class:`Node`23 The set of all nodes in the scene.24 bg_color : (4,) float, optional25 Background color of scene.26 ambient_light : (3,) float, optional27 Color of ambient light. Defaults to no ambient light.28 name : str, optional29 The user-defined name of this object.30 """31 32 def __init__(self,33 nodes=None,34 bg_color=None,35 ambient_light=None,36 name=None):37 38 if bg_color is None:39 bg_color = np.ones(4)40 else:41 bg_color = format_color_vector(bg_color, 4)42 43 if ambient_light is None:44 ambient_light = np.zeros(3)45 46 if nodes is None:47 nodes = set()48 self._nodes = set() # Will be added at the end of this function49 50 self.bg_color = bg_color51 self.ambient_light = ambient_light52 self.name = name53 54 self._name_to_nodes = {}55 self._obj_to_nodes = {}56 self._obj_name_to_nodes = {}57 self._mesh_nodes = set()58 self._point_light_nodes = set()59 self._spot_light_nodes = set()60 self._directional_light_nodes = set()61 self._camera_nodes = set()62 self._main_camera_node = None63 self._bounds = None64 65 # Transform tree66 self._digraph = nx.DiGraph()67 self._digraph.add_node('world')68 self._path_cache = {}69 70 # Find root nodes and add them71 if len(nodes) > 0:72 node_parent_map = {n: None for n in nodes}73 for node in nodes:74 for child in node.children:75 if node_parent_map[child] is not None:76 raise ValueError('Nodes may not have more than '77 'one parent')78 node_parent_map[child] = node79 for node in node_parent_map:80 if node_parent_map[node] is None:81 self.add_node(node)82 83 @property84 def name(self):85 """str : The user-defined name of this object.86 """87 return self._name88 89 @name.setter90 def name(self, value):91 if value is not None:92 value = str(value)93 self._name = value94 95 @property96 def nodes(self):97 """set of :class:`Node` : Set of nodes in the scene.98 """99 return self._nodes100 101 @property102 def bg_color(self):103 """(3,) float : The scene background color.104 """105 return self._bg_color106 107 @bg_color.setter108 def bg_color(self, value):109 if value is None:110 value = np.ones(4)111 else:112 value = format_color_vector(value, 4)113 self._bg_color = value114 115 @property116 def ambient_light(self):117 """(3,) float : The ambient light in the scene.118 """119 return self._ambient_light120 121 @ambient_light.setter122 def ambient_light(self, value):123 if value is None:124 value = np.zeros(3)125 else:126 value = format_color_vector(value, 3)127 self._ambient_light = value128 129 @property130 def meshes(self):131 """set of :class:`Mesh` : The meshes in the scene.132 """133 return set([n.mesh for n in self.mesh_nodes])134 135 @property136 def mesh_nodes(self):137 """set of :class:`Node` : The nodes containing meshes.138 """139 return self._mesh_nodes140 141 @property142 def lights(self):143 """set of :class:`Light` : The lights in the scene.144 """145 return self.point_lights | self.spot_lights | self.directional_lights146 147 @property148 def light_nodes(self):149 """set of :class:`Node` : The nodes containing lights.150 """151 return (self.point_light_nodes | self.spot_light_nodes |152 self.directional_light_nodes)153 154 @property155 def point_lights(self):156 """set of :class:`PointLight` : The point lights in the scene.157 """158 return set([n.light for n in self.point_light_nodes])159 160 @property161 def point_light_nodes(self):162 """set of :class:`Node` : The nodes containing point lights.163 """164 return self._point_light_nodes165 166 @property167 def spot_lights(self):168 """set of :class:`SpotLight` : The spot lights in the scene.169 """170 return set([n.light for n in self.spot_light_nodes])171 172 @property173 def spot_light_nodes(self):174 """set of :class:`Node` : The nodes containing spot lights.175 """176 return self._spot_light_nodes177 178 @property179 def directional_lights(self):180 """set of :class:`DirectionalLight` : The directional lights in181 the scene.182 """183 return set([n.light for n in self.directional_light_nodes])184 185 @property186 def directional_light_nodes(self):187 """set of :class:`Node` : The nodes containing directional lights.188 """189 return self._directional_light_nodes190 191 @property192 def cameras(self):193 """set of :class:`Camera` : The cameras in the scene.194 """195 return set([n.camera for n in self.camera_nodes])196 197 @property198 def camera_nodes(self):199 """set of :class:`Node` : The nodes containing cameras in the scene.200 """201 return self._camera_nodes202 203 @property204 def main_camera_node(self):205 """set of :class:`Node` : The node containing the main camera in the206 scene.207 """208 return self._main_camera_node209 210 @main_camera_node.setter211 def main_camera_node(self, value):212 if value not in self.nodes:213 raise ValueError('New main camera node must already be in scene')214 self._main_camera_node = value215 216 @property217 def bounds(self):218 """(2,3) float : The axis-aligned bounds of the scene.219 """220 if self._bounds is None:221 # Compute corners222 corners = []223 for mesh_node in self.mesh_nodes:224 mesh = mesh_node.mesh225 pose = self.get_pose(mesh_node)226 corners_local = trimesh.bounds.corners(mesh.bounds)227 corners_world = pose[:3,:3].dot(corners_local.T).T + pose[:3,3]228 corners.append(corners_world)229 if len(corners) == 0:230 self._bounds = np.zeros((2,3))231 else:232 corners = np.vstack(corners)233 self._bounds = np.array([np.min(corners, axis=0),234 np.max(corners, axis=0)])235 return self._bounds236 237 @property238 def centroid(self):239 """(3,) float : The centroid of the scene's axis-aligned bounding box240 (AABB).241 """242 return np.mean(self.bounds, axis=0)243 244 @property245 def extents(self):246 """(3,) float : The lengths of the axes of the scene's AABB.247 """248 return np.diff(self.bounds, axis=0).reshape(-1)249 250 @property251 def scale(self):252 """(3,) float : The length of the diagonal of the scene's AABB.253 """254 return np.linalg.norm(self.extents)255 256 def add(self, obj, name=None, pose=None,257 parent_node=None, parent_name=None):258 """Add an object (mesh, light, or camera) to the scene.259 260 Parameters261 ----------262 obj : :class:`Mesh`, :class:`Light`, or :class:`Camera`263 The object to add to the scene.264 name : str265 A name for the new node to be created.266 pose : (4,4) float267 The local pose of this node relative to its parent node.268 parent_node : :class:`Node`269 The parent of this Node. If None, the new node is a root node.270 parent_name : str271 The name of the parent node, can be specified instead of272 `parent_node`.273 274 Returns275 -------276 node : :class:`Node`277 The newly-created and inserted node.278 """279 if isinstance(obj, Mesh):280 node = Node(name=name, matrix=pose, mesh=obj)281 elif isinstance(obj, Light):282 node = Node(name=name, matrix=pose, light=obj)283 elif isinstance(obj, Camera):284 node = Node(name=name, matrix=pose, camera=obj)285 else:286 raise TypeError('Unrecognized object type')287 288 if parent_node is None and parent_name is not None:289 parent_nodes = self.get_nodes(name=parent_name)290 if len(parent_nodes) == 0:291 raise ValueError('No parent node with name {} found'292 .format(parent_name))293 elif len(parent_nodes) > 1:294 raise ValueError('More than one parent node with name {} found'295 .format(parent_name))296 parent_node = list(parent_nodes)[0]297 298 self.add_node(node, parent_node=parent_node)299 300 return node301 302 def get_nodes(self, node=None, name=None, obj=None, obj_name=None):303 """Search for existing nodes. Only nodes matching all specified304 parameters is returned, or None if no such node exists.305 306 Parameters307 ----------308 node : :class:`Node`, optional309 If present, returns this node if it is in the scene.310 name : str311 A name for the Node.312 obj : :class:`Mesh`, :class:`Light`, or :class:`Camera`313 An object that is attached to the node.314 obj_name : str315 The name of an object that is attached to the node.316 317 Returns318 -------319 nodes : set of :class:`.Node`320 The nodes that match all query terms.321 """322 if node is not None:323 if node in self.nodes:324 return set([node])325 else:326 return set()327 nodes = set(self.nodes)328 if name is not None:329 matches = set()330 if name in self._name_to_nodes:331 matches = self._name_to_nodes[name]332 nodes = nodes & matches333 if obj is not None:334 matches = set()335 if obj in self._obj_to_nodes:336 matches = self._obj_to_nodes[obj]337 nodes = nodes & matches338 if obj_name is not None:339 matches = set()340 if obj_name in self._obj_name_to_nodes:341 matches = self._obj_name_to_nodes[obj_name]342 nodes = nodes & matches343 344 return nodes345 346 def add_node(self, node, parent_node=None):347 """Add a Node to the scene.348 349 Parameters350 ----------351 node : :class:`Node`352 The node to be added.353 parent_node : :class:`Node`354 The parent of this Node. If None, the new node is a root node.355 """356 if node in self.nodes:357 raise ValueError('Node already in scene')358 self.nodes.add(node)359 360 # Add node to sets361 if node.name is not None:362 if node.name not in self._name_to_nodes:363 self._name_to_nodes[node.name] = set()364 self._name_to_nodes[node.name].add(node)365 for obj in [node.mesh, node.camera, node.light]:366 if obj is not None:367 if obj not in self._obj_to_nodes:368 self._obj_to_nodes[obj] = set()369 self._obj_to_nodes[obj].add(node)370 if obj.name is not None:371 if obj.name not in self._obj_name_to_nodes:372 self._obj_name_to_nodes[obj.name] = set()373 self._obj_name_to_nodes[obj.name].add(node)374 if node.mesh is not None:375 self._mesh_nodes.add(node)376 if node.light is not None:377 if isinstance(node.light, PointLight):378 self._point_light_nodes.add(node)379 if isinstance(node.light, SpotLight):380 self._spot_light_nodes.add(node)381 if isinstance(node.light, DirectionalLight):382 self._directional_light_nodes.add(node)383 if node.camera is not None:384 self._camera_nodes.add(node)385 if self._main_camera_node is None:386 self._main_camera_node = node387 388 if parent_node is None:389 parent_node = 'world'390 elif parent_node not in self.nodes:391 raise ValueError('Parent node must already be in scene')392 elif node not in parent_node.children:393 parent_node.children.append(node)394 395 # Create node in graph396 self._digraph.add_node(node)397 self._digraph.add_edge(node, parent_node)398 399 # Iterate over children400 for child in node.children:401 self.add_node(child, node)402 403 self._path_cache = {}404 self._bounds = None405 406 def has_node(self, node):407 """Check if a node is already in the scene.408 409 Parameters410 ----------411 node : :class:`Node`412 The node to be checked.413 414 Returns415 -------416 has_node : bool417 True if the node is already in the scene and false otherwise.418 """419 return node in self.nodes420 421 def remove_node(self, node):422 """Remove a node and all its children from the scene.423 424 Parameters425 ----------426 node : :class:`Node`427 The node to be removed.428 """429 # Disconnect self from parent who is staying in the graph430 parent = list(self._digraph.neighbors(node))[0]431 self._remove_node(node)432 if isinstance(parent, Node):433 parent.children.remove(node)434 self._path_cache = {}435 self._bounds = None436 437 def get_pose(self, node):438 """Get the world-frame pose of a node in the scene.439 440 Parameters441 ----------442 node : :class:`Node`443 The node to find the pose of.444 445 Returns446 -------447 pose : (4,4) float448 The transform matrix for this node.449 """450 if node not in self.nodes:451 raise ValueError('Node must already be in scene')452 if node in self._path_cache:453 path = self._path_cache[node]454 else:455 # Get path from from_frame to to_frame456 path = nx.shortest_path(self._digraph, node, 'world')457 self._path_cache[node] = path458 459 # Traverse from from_node to to_node460 pose = np.eye(4)461 for n in path[:-1]:462 pose = np.dot(n.matrix, pose)463 464 return pose465 466 def set_pose(self, node, pose):467 """Set the local-frame pose of a node in the scene.468 469 Parameters470 ----------471 node : :class:`Node`472 The node to set the pose of.473 pose : (4,4) float474 The pose to set the node to.475 """476 if node not in self.nodes:477 raise ValueError('Node must already be in scene')478 node._matrix = pose479 if node.mesh is not None:480 self._bounds = None481 482 def clear(self):483 """Clear out all nodes to form an empty scene.484 """485 self._nodes = set()486 487 self._name_to_nodes = {}488 self._obj_to_nodes = {}489 self._obj_name_to_nodes = {}490 self._mesh_nodes = set()491 self._point_light_nodes = set()492 self._spot_light_nodes = set()493 self._directional_light_nodes = set()494 self._camera_nodes = set()495 self._main_camera_node = None496 self._bounds = None497 498 # Transform tree499 self._digraph = nx.DiGraph()500 self._digraph.add_node('world')501 self._path_cache = {}502 503 def _remove_node(self, node):504 """Remove a node and all its children from the scene.505 506 Parameters507 ----------508 node : :class:`Node`509 The node to be removed.510 """511 512 # Remove self from nodes513 self.nodes.remove(node)514 515 # Remove children516 for child in node.children:517 self._remove_node(child)518 519 # Remove self from the graph520 self._digraph.remove_node(node)521 522 # Remove from maps523 if node.name in self._name_to_nodes:524 self._name_to_nodes[node.name].remove(node)525 if len(self._name_to_nodes[node.name]) == 0:526 self._name_to_nodes.pop(node.name)527 for obj in [node.mesh, node.camera, node.light]:528 if obj is None:529 continue530 self._obj_to_nodes[obj].remove(node)531 if len(self._obj_to_nodes[obj]) == 0:532 self._obj_to_nodes.pop(obj)533 if obj.name is not None:534 self._obj_name_to_nodes[obj.name].remove(node)535 if len(self._obj_name_to_nodes[obj.name]) == 0:536 self._obj_name_to_nodes.pop(obj.name)537 if node.mesh is not None:538 self._mesh_nodes.remove(node)539 if node.light is not None:540 if isinstance(node.light, PointLight):541 self._point_light_nodes.remove(node)542 if isinstance(node.light, SpotLight):543 self._spot_light_nodes.remove(node)544 if isinstance(node.light, DirectionalLight):545 self._directional_light_nodes.remove(node)546 if node.camera is not None:547 self._camera_nodes.remove(node)548 if self._main_camera_node == node:549 if len(self._camera_nodes) > 0:550 self._main_camera_node = next(iter(self._camera_nodes))551 else:552 self._main_camera_node = None553 554 @staticmethod555 def from_trimesh_scene(trimesh_scene,556 bg_color=None, ambient_light=None):557 """Create a :class:`.Scene` from a :class:`trimesh.scene.scene.Scene`.558 559 Parameters560 ----------561 trimesh_scene : :class:`trimesh.scene.scene.Scene`562 Scene with :class:~`trimesh.base.Trimesh` objects.563 bg_color : (4,) float564 Background color for the created scene.565 ambient_light : (3,) float or None566 Ambient light in the scene.567 568 Returns569 -------570 scene_pr : :class:`Scene`571 A scene containing the same geometry as the trimesh scene.572 """573 # convert trimesh geometries to pyrender geometries574 geometries = {name: Mesh.from_trimesh(geom)575 for name, geom in trimesh_scene.geometry.items()}576 577 # create the pyrender scene object578 scene_pr = Scene(bg_color=bg_color, ambient_light=ambient_light)579 580 # add every node with geometry to the pyrender scene581 for node in trimesh_scene.graph.nodes_geometry:582 pose, geom_name = trimesh_scene.graph[node]583 scene_pr.add(geometries[geom_name], pose=pose)584 585 return scene_pr586