Coverage for trimesh/scene/transforms.py: 93%
313 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-31 23:55 +0000
1import collections
2import itertools
3from copy import deepcopy
5import numpy as np
7from .. import caching, util
8from ..caching import hash_fast
9from ..transformations import fix_rigid, quaternion_matrix, rotation_matrix
10from ..typed import ArrayLike, Floating, Hashable, NDArray, Sequence
12# we compare to identity a lot
13_identity = np.eye(4)
14_identity.flags["WRITEABLE"] = False
16# default name for the root frame of a scene graph
17DEFAULT_BASE_FRAME = "world"
20class SceneGraph:
21 """
22 Hold data about positions and instances of geometry
23 in a scene. This includes a forest (i.e. multi-root tree)
24 of transforms and information on which node is the base
25 frame, and which geometries are affiliated with which
26 nodes.
27 """
29 def __init__(
30 self, base_frame: Hashable | None = None, repair_rigid: Floating | None = 1e-5
31 ):
32 """
33 Create a scene graph, holding homogeneous transformation
34 matrices and instance information about geometry.
36 Parameters
37 -----------
38 base_frame
39 The root node transforms will be positioned from.
40 repair_rigid
41 If a float will attempt to repair rotation matrices
42 where `M @ M.T` differs from an identity matrix by
43 more than floating point zero but less than this value.
44 This can happen in a deep tree with a lot of matrix
45 multiplies.
46 """
48 # a graph structure, subclass of networkx DiGraph
49 self.transforms = EnforcedForest()
50 # hashable, the base or root frame — only replace None
51 # as falsy frames like `0` or `""` are valid node names
52 self.base_frame = base_frame if base_frame is not None else DEFAULT_BASE_FRAME
53 # if passed as a float try to repair rigid transforms
54 # that have accumulated floating point error
55 self.repair_rigid = repair_rigid
56 # cache transformation matrices keyed with tuples
57 self._cache = caching.Cache(self.__hash__)
59 def update(self, frame_to, frame_from=None, **kwargs):
60 """
61 Update a transform in the tree.
63 Parameters
64 ------------
65 frame_from : hashable object
66 Usually a string (eg 'world').
67 If left as None it will be set to self.base_frame
68 frame_to : hashable object
69 Usually a string (eg 'mesh_0')
70 matrix : (4,4) float
71 Homogeneous transformation matrix
72 quaternion : (4,) float
73 Quaternion ordered [w, x, y, z]
74 axis : (3,) float
75 Axis of rotation
76 angle : float
77 Angle of rotation, in radians
78 translation : (3,) float
79 Distance to translate
80 geometry : hashable
81 Geometry object name, e.g. 'mesh_0'
82 metadata: dictionary
83 Optional metadata attached to the new frame
84 (exports to glTF node 'extras').
85 """
86 # if no frame specified, use base frame
87 if frame_from is None:
88 frame_from = self.base_frame
90 # pass through
91 attr = {k: v for k, v in kwargs.items() if k in {"geometry", "metadata"}}
92 # convert various kwargs to a single matrix
93 attr["matrix"] = kwargs_to_matrix(**kwargs)
95 # add the edges for the transforms
96 # wi ll return if it changed anything
97 self.transforms.add_edge(frame_from, frame_to, **attr)
99 # set the node attribute with the geometry information
100 if "geometry" in kwargs:
101 self.transforms.node_data[frame_to]["geometry"] = kwargs["geometry"]
103 def get(
104 self, frame_to: Hashable, frame_from: Hashable | None = None
105 ) -> tuple[NDArray[np.float64], Hashable | None]:
106 """
107 Get the transform from one frame to another.
109 Parameters
110 ------------
111 frame_to : hashable
112 Node name, usually a string (eg 'mesh_0')
113 frame_from : hashable
114 Node name, usually a string (eg 'world').
115 If None it will be set to self.base_frame
117 Returns
118 ----------
119 transform : (4, 4) float
120 Homogeneous transformation matrix
121 geometry
122 The name of the geometry if it exists
124 Raises
125 -----------
126 ValueError
127 If the frames aren't connected.
128 """
130 # use base frame if not specified
131 if frame_from is None:
132 frame_from = self.base_frame
134 # look up transform to see if we have it already
135 key = (frame_from, frame_to)
136 if key in self._cache:
137 return self._cache[key]
139 # get the geometry at the final node if any
140 geometry = self.transforms.node_data[frame_to].get("geometry")
142 # get a local reference to edge data
143 data = self.transforms.edge_data
145 if frame_from == frame_to:
146 # if we're going from ourself return identity
147 matrix = _identity
148 elif key in data:
149 # if the path is just an edge return early
150 matrix = data[key]["matrix"]
151 else:
152 # we have a 3+ node path
153 # get the path from the forest always going from
154 # parent -> child -> child
155 path = self.transforms.shortest_path(frame_from, frame_to)
156 # the path should always start with `frame_from`
157 assert path[0] == frame_from
158 # and end with the `frame_to` node
159 assert path[-1] == frame_to
161 # loop through pairs of the path
162 matrices = []
163 for u, v in itertools.pairwise(path):
164 forward = data.get((u, v))
165 if forward is not None:
166 if "matrix" in forward:
167 # append the matrix from u to v
168 matrices.append(forward["matrix"])
169 continue
170 # since forwards didn't exist backward must
171 # exist otherwise this is a disconnected path
172 # and we should raise an error anyway
173 backward = data[(v, u)]
174 if "matrix" in backward:
175 # append the inverted backwards matrix
176 matrices.append(np.linalg.inv(backward["matrix"]))
177 # filter out any identity matrices
178 matrices = [m for m in matrices if np.abs(m - _identity).max() > 1e-8]
179 if len(matrices) == 0:
180 matrix = _identity
181 elif len(matrices) == 1:
182 matrix = matrices[0]
183 else:
184 # multiply matrices into single transform
185 matrix = util.multi_dot(matrices)
187 # if instructed to repair rigid transforms do it here
188 if self.repair_rigid is not None:
189 matrix = fix_rigid(matrix, max_deviance=self.repair_rigid)
191 # matrix being edited in-place leads to subtle bugs
192 matrix.flags["WRITEABLE"] = False
194 # store the result
195 self._cache[key] = (matrix, geometry)
197 return matrix, geometry
199 def __hash__(self):
200 return self.transforms.__hash__()
202 def copy(self):
203 """
204 Return a copy of the current TransformForest.
206 Returns
207 ------------
208 copied : TransformForest
209 Copy of current object.
210 """
211 # create a copy without transferring cache
212 copied = SceneGraph()
213 copied.base_frame = deepcopy(self.base_frame)
214 copied.transforms = deepcopy(self.transforms)
215 return copied
217 def to_flattened(self):
218 """
219 Export the current transform graph with all
220 transforms baked into world->instance.
222 Returns
223 ---------
224 flat : dict
225 Keyed {node : {transform, geometry}
226 """
227 flat = {}
228 base_frame = self.base_frame
229 for node in self.nodes:
230 if node == base_frame:
231 continue
232 # get the matrix and geometry name
233 matrix, geometry = self.get(frame_to=node, frame_from=base_frame)
234 # store matrix as list rather than numpy array
235 flat[node] = {"transform": matrix.tolist(), "geometry": geometry}
237 return flat
239 def to_gltf(self, scene, mesh_index=None):
240 """
241 Export a transforms as the 'nodes' section of the
242 GLTF header dict.
244 Parameters
245 ------------
246 scene : trimesh.Scene
247 Scene with geometry.
248 mesh_index : dict or None
249 Mapping { key in scene.geometry : int }
251 Returns
252 --------
253 gltf : dict
254 With 'nodes' referencing a list of dicts and 'scene_roots'
255 referencing the node indices the scene should start from.
256 """
258 if mesh_index is None:
259 # geometry is an OrderedDict
260 # map mesh name to index: {geometry key : index}
261 mesh_index = {name: i for i, name in enumerate(scene.geometry.keys())}
263 # get graph information into local scope before loop
264 graph = self.transforms
265 # get the stored node data
266 node_data = graph.node_data
267 edge_data = graph.edge_data
268 base_frame = self.base_frame
269 # does the scene have a defined camera to export
270 has_camera = scene.has_camera
271 children = graph.children
273 # the base frame is a synthetic wrapper: when it carries nothing
274 # export its children as the scene roots instead of writing it —
275 # a file node named like the base frame is renamed on import so
276 # writing one would grow a wrapper node on every round-trip
277 skip_base = (
278 "geometry" not in node_data.get(base_frame, {})
279 and not (has_camera and base_frame == scene.camera.name)
280 and len(children.get(base_frame, [])) > 0
281 )
283 # list of dict, in gltf format
284 # {node name : node index in gltf}
285 if skip_base:
286 result = []
287 lookup = {}
288 else:
289 result = [{"name": base_frame}]
290 lookup = {base_frame: 0}
292 # collect the nodes in order
293 for node in node_data.keys():
294 if node == base_frame:
295 continue
296 # assign the index to the node-name lookup
297 lookup[node] = len(result)
298 # populate a result at the correct index
299 result.append({"name": node})
301 extensions_used = set()
303 # then iterate through to collect data
304 for info in result:
305 # name of the scene node
306 node = info["name"]
308 # get the original node names for children
309 childs = children.get(node, [])
310 if len(childs) > 0:
311 info["children"] = [lookup[k] for k in childs]
313 # if we have a mesh store by index
314 if "geometry" in node_data[node]:
315 mesh_key = node_data[node]["geometry"]
316 if mesh_key in mesh_index:
317 info["mesh"] = mesh_index[mesh_key]
318 # check to see if we have camera node
319 if has_camera and node == scene.camera.name:
320 info["camera"] = 0
322 if node != base_frame:
323 parent = graph.parents[node]
324 node_edge = edge_data[(parent, node)]
326 # get the matrix from this edge
327 matrix = node_edge["matrix"]
328 # only include if it's not an identify matrix
329 if not util.allclose(matrix, _identity):
330 info["matrix"] = matrix.T.reshape(-1).tolist()
332 # if an extra was stored on this edge
333 extras = node_edge.get("metadata")
334 if extras:
335 extras = extras.copy()
337 # if extensionss were stored on this edge
338 extensions = extras.pop("gltf_extensions", None)
339 if isinstance(extensions, dict):
340 info["extensions"] = extensions
341 extensions_used = extensions_used.union(set(extensions.keys()))
343 # convert any numpy arrays to lists
344 extras.update(
345 {k: v.tolist() for k, v in extras.items() if hasattr(v, "tolist")}
346 )
347 info["extras"] = extras
349 if skip_base:
350 roots = [lookup[c] for c in children[base_frame]]
351 else:
352 roots = [lookup[base_frame]]
354 gltf = {"nodes": result, "scene_roots": roots}
355 if len(extensions_used) > 0:
356 gltf["extensionsUsed"] = list(extensions_used)
357 return gltf
359 def to_edgelist(self):
360 """
361 Export the current transforms as a list of
362 edge tuples, with each tuple having the format:
363 (node_a, node_b, {metadata})
365 Returns
366 ---------
367 edgelist : (n,) list
368 Of edge tuples
369 """
370 # save local reference to node_data
371 nodes = self.transforms.node_data
372 # save cleaned edges
373 export = []
374 # loop through (node, node, edge attributes)
375 for edge, attr in self.transforms.edge_data.items():
376 # node indexes from edge
377 a, b = edge
378 # geometry is a node property but save it to the
379 # edge so we don't need two dictionaries
380 b_attr = nodes[b]
381 # make sure we're not stomping on original
382 attr_new = attr.copy()
383 # apply node geometry to edge attributes
384 if "geometry" in b_attr:
385 attr_new["geometry"] = b_attr["geometry"]
386 # convert any numpy arrays to regular lists
387 attr_new.update(
388 {k: v.tolist() for k, v in attr_new.items() if hasattr(v, "tolist")}
389 )
390 export.append([a, b, attr_new])
391 return export
393 def from_edgelist(self, edges, strict=True):
394 """
395 Load transform data from an edge list into the current
396 scene graph.
398 Parameters
399 -------------
400 edgelist : (n,) tuples
401 Keyed (node_a, node_b, {key: value})
402 strict : bool
403 If True raise a ValueError when a
404 malformed edge is passed in a tuple.
405 """
407 # loop through each edge
408 for edge in edges:
409 # edge contains attributes
410 if len(edge) == 3:
411 self.update(edge[1], edge[0], **edge[2])
412 # edge just contains nodes
413 elif len(edge) == 2:
414 self.update(edge[1], edge[0])
415 # edge is broken
416 elif strict:
417 raise ValueError("edge incorrect shape: %s", str(edge))
419 def to_networkx(self):
420 """
421 Return a `networkx` copy of this graph.
423 Returns
424 ----------
425 graph : networkx.DiGraph
426 Directed graph.
427 """
428 import networkx
430 return networkx.from_edgelist(self.to_edgelist(), create_using=networkx.DiGraph)
432 def show(self, **kwargs):
433 """
434 Plot the scene graph using `networkx.draw_networkx`
435 which uses matplotlib to display the graph.
437 Parameters
438 -----------
439 kwargs : dict
440 Passed to `networkx.draw_networkx`
441 """
442 import matplotlib.pyplot as plt # noqa
443 import networkx
445 # default kwargs will only be set if not
446 # passed explicitly to the show command
447 defaults = {"with_labels": True}
448 kwargs.update(**{k: v for k, v in defaults.items() if k not in kwargs})
449 networkx.draw_networkx(G=self.to_networkx(), **kwargs)
451 plt.show()
453 def load(self, edgelist):
454 """
455 Load transform data from an edge list into the current
456 scene graph.
458 Parameters
459 -------------
460 edgelist : (n,) tuples
461 Structured (node_a, node_b, {key: value})
462 """
463 self.from_edgelist(edgelist, strict=True)
465 @caching.cache_decorator
466 def nodes(self):
467 """
468 A list of every node in the graph.
470 Returns
471 -------------
472 nodes : (n,) array
473 All node names.
474 """
475 return self.transforms.nodes
477 @caching.cache_decorator
478 def nodes_geometry(self):
479 """
480 The nodes in the scene graph with geometry attached.
482 Returns
483 ------------
484 nodes_geometry : (m,) array
485 Node names which have geometry associated
486 """
487 return [n for n, attr in self.transforms.node_data.items() if "geometry" in attr]
489 @caching.cache_decorator
490 def geometry_nodes(self):
491 """
492 Which nodes have this geometry? Inverse
493 of `nodes_geometry`.
495 Returns
496 ------------
497 geometry_nodes : dict
498 Keyed {geometry_name : node name}
499 """
500 res = collections.defaultdict(list)
501 for node, attr in self.transforms.node_data.items():
502 if "geometry" in attr:
503 res[attr["geometry"]].append(node)
504 return res
506 def remove_geometries(self, geometries: str | set | Sequence):
507 """
508 Remove the reference for specified geometries
509 from nodes without deleting the node.
511 Parameters
512 ------------
513 geometries : list or str
514 Name of scene.geometry to dereference.
515 """
516 # make sure we have a set of geometries to remove
517 if isinstance(geometries, str):
518 geometries = [geometries]
519 geometries = set(geometries)
521 # remove the geometry reference from the node without deleting nodes
522 # this lets us keep our cached paths, and will not screw up children
523 for attrib in self.transforms.node_data.values():
524 if "geometry" in attrib and attrib["geometry"] in geometries:
525 attrib.pop("geometry")
527 # it would be safer to just run _cache.clear
528 # but the only property using the geometry should be
529 # nodes_geometry: if this becomes not true change this to clear!
530 self._cache.cache.pop("nodes_geometry", None)
531 self.transforms._hash = None
533 def __contains__(self, key: Hashable) -> bool:
534 return key in self.transforms.node_data
536 def __getitem__(self, key: Hashable) -> tuple[NDArray[np.float64], Hashable | None]:
537 return self.get(key)
539 def __setitem__(self, key: Hashable, value: ArrayLike):
540 value = np.asanyarray(value, dtype=np.float64)
541 if value.shape != (4, 4):
542 raise ValueError("Matrix must be specified!")
543 return self.update(key, matrix=value)
545 def clear(self):
546 self.transforms = EnforcedForest()
547 self._cache.clear()
550class EnforcedForest:
551 """
552 A simple forest graph data structure: every node
553 is allowed to have exactly one parent. This makes
554 traversal and implementation much simpler than a
555 full graph data type; by storing only one parent
556 reference, it enforces the structure for "free."
557 """
559 def __init__(self):
560 # since every node can have only one parent
561 # this data structure transparently enforces
562 # the forest data structure without checks
563 # a dict {child : parent}
564 self.parents = {}
566 # store data for a particular edge keyed by tuple
567 # {(u, v) : data }
568 self.edge_data = collections.defaultdict(dict)
569 # {u: data}
570 self.node_data = collections.defaultdict(dict)
572 # if multiple calls are made for the same path
573 # but the connectivity hasn't changed return cached
574 self._cache = {}
576 def add_edge(self, u, v, **kwargs):
577 """
578 Add an edge to the forest cleanly.
580 Parameters
581 -----------
582 u : any
583 Hashable node key.
584 v : any
585 Hashable node key.
586 kwargs : dict
587 Stored as (u, v) edge data.
589 Returns
590 --------
591 changed : bool
592 Return if this operation changed anything.
593 """
594 self._hash = None
596 # topology has changed so clear cache
597 if (u, v) not in self.edge_data:
598 self._cache = {}
599 else:
600 # check to see if matrix and geometry are identical
601 edge = self.edge_data[(u, v)]
602 if util.allclose(
603 kwargs.get("matrix", _identity), edge.get("matrix", _identity), 1e-8
604 ) and (edge.get("geometry") == kwargs.get("geometry")):
605 return False
607 # store a parent reference for traversal
608 self.parents[v] = u
609 # store kwargs for edge data keyed with tuple
610 self.edge_data[(u, v)] = kwargs
611 # set empty node data
612 self.node_data[u].update({})
613 if "geometry" in kwargs:
614 self.node_data[v].update({"geometry": kwargs["geometry"]})
615 else:
616 self.node_data[v].update({})
618 return True
620 def remove_node(self, u):
621 """
622 Remove a node from the forest.
624 Parameters
625 -----------
626 u : any
627 Hashable node key.
629 Returns
630 --------
631 changed : bool
632 Return if this operation changed anything.
633 """
634 # check if node is part of forest
635 if u not in self.node_data:
636 return False
638 # topology will change so clear cache
639 self._cache = {}
640 self._hash = None
642 # delete all children's references and parent reference
643 children = [child for (child, parent) in self.parents.items() if parent == u]
644 for c in children:
645 del self.parents[c]
646 if u in self.parents:
647 del self.parents[u]
649 # delete edge data
650 edges = [(a, b) for (a, b) in self.edge_data if a == u or b == u]
651 for e in edges:
652 del self.edge_data[e]
654 # delete node data
655 del self.node_data[u]
657 return True
659 def shortest_path(self, u, v):
660 """
661 Find the shortest path between `u` and `v`, returning
662 a path where the first element is always `u` and the
663 last element is always `v`, disregarding edge direction.
665 Parameters
666 -----------
667 u : any
668 Hashable node key.
669 v : any
670 Hashable node key.
672 Returns
673 -----------
674 path : (n,)
675 Path between `u` and `v`
676 """
677 # see if we've already computed this path
678 if u == v:
679 # the path between itself is an edge case
680 return []
681 elif (u, v) in self._cache:
682 # return the same path for either direction
683 return self._cache[(u, v)]
684 elif (v, u) in self._cache:
685 return self._cache[(v, u)][::-1]
687 # local reference to parent dict for performance
688 parents = self.parents
689 # store both forward and backwards traversal
690 forward = [u]
691 backward = [v]
693 # cap iteration to number of total nodes
694 for _ in range(len(parents) + 1):
695 # store the parent both forwards and backwards
696 f = parents.get(forward[-1])
697 b = parents.get(backward[-1])
698 forward.append(f)
699 backward.append(b)
701 if f == v:
702 self._cache[(u, v)] = forward
703 return forward
704 elif b == u:
705 # return reversed path
706 backward = backward[::-1]
707 self._cache[(u, v)] = backward
708 return backward
709 elif (b in forward) or (f is None and b is None):
710 # we have a either a common node between both
711 # traversal directions or we have consumed the whole
712 # tree in both directions so try to find the common node
713 common = set(backward).intersection(forward).difference({None})
714 if len(common) == 0:
715 raise ValueError(f"No path from {u}->{v}!")
716 elif len(common) > 1:
717 # get the first occurring common element in "forward"
718 link = next(f for f in forward if f in common)
719 assert link in common
720 else:
721 # take the only common element
722 link = next(iter(common))
724 # combine the forward and backwards traversals
725 a = forward[: forward.index(link) + 1]
726 b = backward[: backward.index(link)]
727 path = a + b[::-1]
729 # verify we didn't screw up the order
730 assert path[0] == u
731 assert path[-1] == v
733 self._cache[(u, v)] = path
735 return path
737 raise ValueError("Iteration limit exceeded!")
739 @property
740 def nodes(self):
741 """
742 Get a set of every node.
744 Returns
745 -----------
746 nodes : set
747 Every node currently stored.
748 """
749 return self.node_data.keys()
751 @property
752 def children(self):
753 """
754 Get the children of each node.
756 Returns
757 ----------
758 children : dict
759 Keyed {node : [child, child, ...]}
760 """
761 if "children" in self._cache:
762 return self._cache["children"]
763 child = collections.defaultdict(list)
764 # append children to parent references
765 # skip self-references to avoid a node loop
766 [child[v].append(u) for u, v in self.parents.items() if u != v]
768 # cache and return as a vanilla dict
769 self._cache["children"] = dict(child)
770 return self._cache["children"]
772 def successors(self, node):
773 """
774 Get all nodes that are successors to specified node,
775 including the specified node.
777 Parameters
778 -------------
779 node : any
780 Hashable key for a node.
782 Returns
783 ------------
784 successors : set
785 Nodes that succeed specified node.
786 """
787 # get mapping of {parent : child}
788 children = self.children
789 # if node doesn't exist return early
790 if node not in children:
791 return {node}
793 # children we need to collect
794 queue = [node]
795 # start collecting values with children of source
796 collected = set(queue)
798 # cap maximum iterations
799 for _ in range(len(self.node_data) + 1):
800 if len(queue) == 0:
801 # no more nodes to visit so we're done
802 return collected
803 # add the children of this node to be processed
804 childs = children.get(queue.pop())
805 if childs is not None:
806 queue.extend(childs)
807 collected.update(childs)
808 return collected
810 def __hash__(self):
811 """
812 Actually hash all of the data, but use a "dirty" mechanism
813 in functions that modify the data, which MUST
814 # all invalidate the hash by setting `self._hash = None`
816 This was optimized a bit, and is evaluating on an
817 older laptop on a scene with 77 nodes and 76 edges
818 10,000 times in 0.7s which seems fast enough.
819 """
820 # see if there is an available hash value
821 # if you are seeing cache bugs this is the thing
822 # to try eliminating because it is very likely that
823 # someone somewhere is modifying the data without
824 # setting `self._hash = None`
825 hashed = getattr(self, "_hash", None)
826 if hashed is not None:
827 return hashed
829 hashed = hash_fast(
830 (
831 "".join(
832 str(hash(k)) + v.get("geometry", "")
833 for k, v in self.edge_data.items()
834 )
835 + "".join(
836 str(k) + v.get("geometry", "") for k, v in self.node_data.items()
837 )
838 ).encode("utf-8")
839 + b"".join(
840 v["matrix"].tobytes() for v in self.edge_data.values() if "matrix" in v
841 )
842 )
843 self._hash = hashed
844 return hashed
847def kwargs_to_matrix(
848 matrix=None, quaternion=None, translation=None, axis=None, angle=None, **kwargs
849):
850 """
851 Take multiple keyword arguments and parse them
852 into a homogeneous transformation matrix.
854 Returns
855 ---------
856 matrix : (4, 4) float
857 Homogeneous transformation matrix.
858 """
859 if matrix is not None:
860 # a matrix takes immediate precedence over other options
861 return np.array(matrix, dtype=np.float64)
862 elif quaternion is not None:
863 matrix = quaternion_matrix(quaternion)
864 elif axis is not None and angle is not None:
865 matrix = rotation_matrix(angle, axis)
866 else:
867 matrix = np.eye(4)
869 if translation is not None:
870 # translation can be used in conjunction with any
871 # of the methods specifying transforms
872 matrix[:3, 3] += translation
874 return matrix