init
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
from .cameras import Camera
|
||||
from .scene import Scene, split_scene
|
||||
|
||||
# add to __all__ as per pep8
|
||||
__all__ = ["Camera", "Scene", "split_scene"]
|
||||
@@ -0,0 +1,423 @@
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import util
|
||||
|
||||
|
||||
class Camera:
|
||||
def __init__(
|
||||
self, name=None, resolution=None, focal=None, fov=None, z_near=0.01, z_far=1000.0
|
||||
):
|
||||
"""
|
||||
Create a new Camera object that stores camera intrinsic
|
||||
and extrinsic parameters.
|
||||
|
||||
TODO: skew is not supported
|
||||
TODO: cx and cy that are not half of width and height
|
||||
|
||||
Parameters
|
||||
------------
|
||||
name : str or None
|
||||
Name for camera to be used as node name
|
||||
resolution : (2,) int
|
||||
Pixel size in (height, width)
|
||||
focal : (2,) float
|
||||
Focal length in pixels. Either pass this OR FOV
|
||||
but not both. focal = (K[0][0], K[1][1])
|
||||
fov : (2,) float
|
||||
Field of view (fovx, fovy) in degrees
|
||||
z_near : float
|
||||
What is the closest
|
||||
"""
|
||||
|
||||
if name is None:
|
||||
# if name is not passed, make it something unique
|
||||
self.name = f"camera_{util.unique_id(6).upper()}"
|
||||
else:
|
||||
# otherwise assign it
|
||||
self.name = name
|
||||
|
||||
if fov is None and focal is None:
|
||||
raise ValueError("either focal length or FOV required!")
|
||||
|
||||
# store whether or not we computed the focal length
|
||||
self._focal_computed = False
|
||||
|
||||
# set the passed (2,) float focal length
|
||||
self.focal = focal
|
||||
|
||||
# set the passed (2,) float FOV in degrees
|
||||
self.fov = fov
|
||||
|
||||
if resolution is None:
|
||||
# if unset make resolution 30 pixels per degree
|
||||
resolution = (self.fov * 30.0).round().astype(np.int64)
|
||||
self.resolution = resolution
|
||||
|
||||
# what is the farthest from the camera it should render
|
||||
self.z_far = float(z_far)
|
||||
# what is the closest to the camera it should render
|
||||
self.z_near = float(z_near)
|
||||
|
||||
def copy(self):
|
||||
"""
|
||||
Safely get a copy of the current camera.
|
||||
"""
|
||||
return Camera(
|
||||
name=copy.deepcopy(self.name),
|
||||
resolution=copy.deepcopy(self.resolution),
|
||||
focal=copy.deepcopy(self.focal),
|
||||
fov=copy.deepcopy(self.fov),
|
||||
)
|
||||
|
||||
@property
|
||||
def resolution(self):
|
||||
"""
|
||||
Get the camera resolution in pixels.
|
||||
|
||||
Returns
|
||||
------------
|
||||
resolution (2,) float
|
||||
Camera resolution in pixels
|
||||
"""
|
||||
return self._resolution
|
||||
|
||||
@resolution.setter
|
||||
def resolution(self, values):
|
||||
"""
|
||||
Set the camera resolution in pixels.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
resolution (2,) float
|
||||
Camera resolution in pixels
|
||||
"""
|
||||
values = np.asanyarray(values, dtype=np.int64)
|
||||
if values.shape != (2,):
|
||||
raise ValueError("resolution must be (2,) float")
|
||||
values.flags.writeable = False
|
||||
self._resolution = values
|
||||
# unset computed value that depends on the other plus resolution
|
||||
if self._focal_computed:
|
||||
self._focal = None
|
||||
else:
|
||||
# fov must be computed
|
||||
self._fov = None
|
||||
|
||||
@property
|
||||
def focal(self):
|
||||
"""
|
||||
Get the focal length in pixels for the camera.
|
||||
|
||||
Returns
|
||||
------------
|
||||
focal : (2,) float
|
||||
Focal length in pixels
|
||||
"""
|
||||
if self._focal is None:
|
||||
# calculate focal length from FOV
|
||||
focal = self._resolution / (2.0 * np.tan(np.radians(self._fov / 2.0)))
|
||||
focal.flags.writeable = False
|
||||
self._focal = focal
|
||||
|
||||
return self._focal
|
||||
|
||||
@focal.setter
|
||||
def focal(self, values):
|
||||
"""
|
||||
Set the focal length in pixels for the camera.
|
||||
|
||||
Returns
|
||||
------------
|
||||
focal : (2,) float
|
||||
Focal length in pixels.
|
||||
"""
|
||||
if values is None:
|
||||
self._focal = None
|
||||
else:
|
||||
# flag this as not computed (hence fov must be)
|
||||
# this is necessary so changes to resolution can reset the
|
||||
# computed quantity without changing the explicitly set quantity
|
||||
self._focal_computed = False
|
||||
values = np.asanyarray(values, dtype=np.float64)
|
||||
if values.shape != (2,):
|
||||
raise ValueError("focal length must be (2,) float")
|
||||
values.flags.writeable = False
|
||||
# assign passed values to focal length
|
||||
self._focal = values
|
||||
# focal overrides FOV
|
||||
self._fov = None
|
||||
|
||||
@property
|
||||
def K(self):
|
||||
"""
|
||||
Get the intrinsic matrix for the Camera object.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
K : (3, 3) float
|
||||
Intrinsic matrix for camera
|
||||
"""
|
||||
K = np.eye(3, dtype=np.float64)
|
||||
K[0, 0] = self.focal[0]
|
||||
K[1, 1] = self.focal[1]
|
||||
K[:2, 2] = self.resolution / 2.0
|
||||
return K
|
||||
|
||||
@K.setter
|
||||
def K(self, values):
|
||||
if values is None:
|
||||
return
|
||||
values = np.asanyarray(values, dtype=np.float64)
|
||||
if values.shape != (3, 3):
|
||||
raise ValueError("matrix must be (3,3)!")
|
||||
|
||||
if not np.allclose(values.flatten()[[1, 3, 6, 7, 8]], [0, 0, 0, 0, 1]):
|
||||
raise ValueError("matrix should only have focal length and resolution!")
|
||||
|
||||
# set focal length from matrix
|
||||
self.focal = [values[0, 0], values[1, 1]]
|
||||
# set resolution from matrix
|
||||
self.resolution = values[:2, 2] * 2
|
||||
|
||||
@property
|
||||
def fov(self):
|
||||
"""
|
||||
Get the field of view in degrees.
|
||||
|
||||
Returns
|
||||
-------------
|
||||
fov : (2,) float
|
||||
XY field of view in degrees
|
||||
"""
|
||||
if self._fov is None:
|
||||
fov = 2.0 * np.degrees(np.arctan((self._resolution / 2.0) / self._focal))
|
||||
fov.flags.writeable = False
|
||||
self._fov = fov
|
||||
return self._fov
|
||||
|
||||
@fov.setter
|
||||
def fov(self, values):
|
||||
"""
|
||||
Set the field of view in degrees.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
values : (2,) float
|
||||
Size of FOV to set in degrees
|
||||
"""
|
||||
if values is None:
|
||||
self._fov = None
|
||||
else:
|
||||
# flag this as computed (hence fov must not be)
|
||||
# this is necessary so changes to resolution can reset the
|
||||
# computed quantity without changing the explicitly set quantity
|
||||
self._focal_computed = True
|
||||
values = np.asanyarray(values, dtype=np.float64)
|
||||
if values.shape != (2,):
|
||||
raise ValueError("fov length must be (2,) int")
|
||||
values.flags.writeable = False
|
||||
# assign passed values to FOV
|
||||
self._fov = values
|
||||
# fov overrides focal
|
||||
self._focal = None
|
||||
|
||||
def to_rays(self):
|
||||
"""
|
||||
Calculate ray direction vectors.
|
||||
|
||||
Will return one ray per pixel, as set in self.resolution.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
vectors : (n, 3) float
|
||||
Ray direction vectors in camera frame with z == -1
|
||||
"""
|
||||
return camera_to_rays(self)
|
||||
|
||||
def angles(self):
|
||||
"""
|
||||
Get ray spherical coordinates in radians.
|
||||
|
||||
|
||||
Returns
|
||||
--------------
|
||||
angles : (n, 2) float
|
||||
Ray spherical coordinate angles in radians.
|
||||
"""
|
||||
return np.arctan(-ray_pixel_coords(self))
|
||||
|
||||
def look_at(self, points, **kwargs):
|
||||
"""
|
||||
Generate transform for a camera to keep a list
|
||||
of points in the camera's field of view.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
points : (n, 3) float
|
||||
Points in space
|
||||
rotation : None, or (4, 4) float
|
||||
Rotation matrix for initial rotation
|
||||
distance : None or float
|
||||
Distance from camera to center
|
||||
center : None, or (3,) float
|
||||
Center of field of view.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
transform : (4, 4) float
|
||||
Transformation matrix from world to camera
|
||||
"""
|
||||
return look_at(points, fov=self.fov, **kwargs)
|
||||
|
||||
def __repr__(self):
|
||||
return f"<trimesh.scene.Camera> FOV: {self.fov} Resolution: {self.resolution}"
|
||||
|
||||
|
||||
def look_at(points, fov, rotation=None, distance=None, center=None, pad=None):
|
||||
"""
|
||||
Generate transform for a camera to keep a list
|
||||
of points in the camera's field of view.
|
||||
|
||||
Examples
|
||||
------------
|
||||
```python
|
||||
points = np.array([0, 0, 0], [1, 1, 1])
|
||||
scene.camera_transform = scene.camera.look_at(points)
|
||||
```
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
points : (n, 3) float
|
||||
Points in space
|
||||
fov : (2,) float
|
||||
Field of view, in DEGREES
|
||||
rotation : None, or (4, 4) float
|
||||
Rotation matrix for initial rotation
|
||||
distance : None or float
|
||||
Distance from camera to center
|
||||
center : None, or (3,) float
|
||||
Center of field of view.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
transform : (4, 4) float
|
||||
Transformation matrix from world to camera
|
||||
"""
|
||||
|
||||
if rotation is None:
|
||||
rotation = np.eye(4)
|
||||
else:
|
||||
rotation = np.asanyarray(rotation, dtype=np.float64)
|
||||
points = np.asanyarray(points, dtype=np.float64)
|
||||
|
||||
# Transform points to camera frame (just use the rotation part)
|
||||
rinv = rotation[:3, :3].T
|
||||
points_c = rinv.dot(points.T).T
|
||||
|
||||
if center is None:
|
||||
# Find the center of the points' AABB in camera frame
|
||||
center_c = points_c.min(axis=0) + 0.5 * np.ptp(points_c, axis=0)
|
||||
else:
|
||||
# Transform center to camera frame
|
||||
center_c = rinv.dot(center)
|
||||
|
||||
# Re-center the points around the camera-frame origin
|
||||
points_c -= center_c
|
||||
|
||||
# Find the minimum distance for the camera from the origin
|
||||
# so that all points fit in the view frustum
|
||||
tfov = np.tan(np.radians(fov) / 2.0)
|
||||
|
||||
if distance is None:
|
||||
distance = np.max(np.abs(points_c[:, :2]) / tfov + points_c[:, 2][:, np.newaxis])
|
||||
|
||||
if pad is not None:
|
||||
distance *= pad
|
||||
|
||||
# set the pose translation
|
||||
center_w = rotation[:3, :3].dot(center_c)
|
||||
cam_pose = rotation.copy()
|
||||
cam_pose[:3, 3] = center_w + distance * cam_pose[:3, 2]
|
||||
|
||||
return cam_pose
|
||||
|
||||
|
||||
def ray_pixel_coords(camera):
|
||||
"""
|
||||
Get the x-y coordinates of rays in camera coordinates at
|
||||
z == -1.
|
||||
|
||||
One coordinate pair will be given for each pixel as defined in
|
||||
camera.resolution. If reshaped, the returned array corresponds
|
||||
to pixels of the rendered image.
|
||||
|
||||
Examples
|
||||
------------
|
||||
```python
|
||||
xy = ray_pixel_coords(camera).reshape(
|
||||
tuple(camera.coordinates) + (2,))
|
||||
top_left == xy[0, 0]
|
||||
bottom_right == xy[-1, -1]
|
||||
```
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
camera : trimesh.scene.Camera
|
||||
Camera object to generate rays from
|
||||
|
||||
Returns
|
||||
--------------
|
||||
xy : (n, 2) float
|
||||
x-y coordinates of intersection of each camera ray
|
||||
with the z == -1 frame
|
||||
"""
|
||||
# shorthand
|
||||
res = camera.resolution
|
||||
half_fov = np.radians(camera.fov) / 2.0
|
||||
|
||||
right_top = np.tan(half_fov)
|
||||
# move half a pixel width in
|
||||
right_top *= 1 - (1.0 / res)
|
||||
left_bottom = -right_top
|
||||
# we are looking down the negative z axis, so
|
||||
# right_top corresponds to maximum x/y values
|
||||
# bottom_left corresponds to minimum x/y values
|
||||
right, top = right_top
|
||||
left, bottom = left_bottom
|
||||
|
||||
# create a grid of vectors
|
||||
xy = util.grid_linspace(
|
||||
bounds=[[left, top], [right, bottom]], count=camera.resolution
|
||||
)
|
||||
|
||||
# create a matching array of pixel indexes for the rays
|
||||
pixels = util.grid_linspace(
|
||||
bounds=[[0, res[1] - 1], [res[0] - 1, 0]], count=res
|
||||
).astype(np.int64)
|
||||
assert xy.shape == pixels.shape
|
||||
|
||||
return xy, pixels
|
||||
|
||||
|
||||
def camera_to_rays(camera: Camera):
|
||||
"""
|
||||
Calculate the trimesh.scene.Camera object to direction vectors.
|
||||
|
||||
Will return one ray per pixel, as set in camera.resolution.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
camera : trimesh.scene.Camera
|
||||
|
||||
Returns
|
||||
--------------
|
||||
vectors : (n, 3) float
|
||||
Ray direction vectors in camera frame with z == -1
|
||||
"""
|
||||
# get the on-plane coordinates
|
||||
xy, pixels = ray_pixel_coords(camera)
|
||||
# convert vectors to 3D unit vectors
|
||||
vectors = util.unitize(np.column_stack((xy, -np.ones_like(xy[:, :1]))))
|
||||
return vectors, pixels
|
||||
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
lighting.py
|
||||
--------------
|
||||
|
||||
Hold basic information about lights.
|
||||
|
||||
Forked from the light model in `pyrender`:
|
||||
https://github.com/mmatl/pyrender
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import transformations, util, visual
|
||||
from ..typed import List, NDArray, Tuple, float64
|
||||
|
||||
# default light color
|
||||
_DEFAULT_RGBA = np.array([60, 60, 60, 255], dtype=np.uint8)
|
||||
|
||||
|
||||
class Light(util.ABC):
|
||||
"""
|
||||
Base class for all light objects.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
name : str, optional
|
||||
Name of the light.
|
||||
color : (4,) uint8
|
||||
RGBA value for the light's color in linear space.
|
||||
intensity : float
|
||||
Brightness of light. The units that this is defined in depend
|
||||
on the type of light: point and spot lights use luminous intensity
|
||||
in candela (lm/sr) while directional lights use illuminance
|
||||
in lux (lm/m2).
|
||||
radius : float
|
||||
Cutoff distance at which light's intensity may be considered to
|
||||
have reached zero. Supported only for point and spot lights
|
||||
Must be > 0.0
|
||||
If None, the radius is assumed to be infinite.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, color=None, intensity=None, radius=None):
|
||||
if name is None:
|
||||
# if name is not passed, make it something unique
|
||||
self.name = f"light_{util.unique_id(6).upper()}"
|
||||
else:
|
||||
# otherwise assign it
|
||||
self.name = name
|
||||
|
||||
self.color = color
|
||||
self.intensity = intensity
|
||||
self.radius = radius
|
||||
|
||||
@property
|
||||
def color(self):
|
||||
return self._color
|
||||
|
||||
@color.setter
|
||||
def color(self, value):
|
||||
if value is None:
|
||||
self._color = _DEFAULT_RGBA
|
||||
else:
|
||||
value = visual.to_rgba(value)
|
||||
if len(value.shape) == 2:
|
||||
value = value[0]
|
||||
if value.shape != (4,):
|
||||
raise ValueError("couldn't convert color to RGBA!")
|
||||
# uint8 RGB color
|
||||
self._color = value
|
||||
|
||||
@property
|
||||
def intensity(self):
|
||||
return self._intensity
|
||||
|
||||
@intensity.setter
|
||||
def intensity(self, value):
|
||||
if value is not None:
|
||||
self._intensity = float(value)
|
||||
else:
|
||||
self._intensity = 1.0
|
||||
|
||||
@property
|
||||
def radius(self):
|
||||
return self._radius
|
||||
|
||||
@radius.setter
|
||||
def radius(self, value):
|
||||
if value is None or value < 0.0:
|
||||
self._radius = value
|
||||
else:
|
||||
self._radius = float(value)
|
||||
|
||||
|
||||
class DirectionalLight(Light):
|
||||
"""
|
||||
Directional lights are light sources that act as though they are
|
||||
infinitely far away and emit light in the direction of the local -z axis.
|
||||
This light type inherits the orientation of the node that it belongs to;
|
||||
position and scale are ignored except for their effect on the inherited
|
||||
node orientation. Because it is at an infinite distance, the light is
|
||||
not attenuated. Its intensity is defined in lumens per metre squared,
|
||||
or lux (lm/m2).
|
||||
|
||||
Attributes
|
||||
----------
|
||||
name : str, optional
|
||||
Name of the light.
|
||||
color : (4,) unit8
|
||||
RGBA value for the light's color in linear space.
|
||||
intensity : float
|
||||
Brightness of light. The units that this is defined in depend
|
||||
on the type of light.
|
||||
point and spot lights use luminous intensity in candela (lm/sr),
|
||||
while directional lights use illuminance in lux (lm/m2).
|
||||
radius : float
|
||||
Cutoff distance at which light's intensity may be considered to
|
||||
have reached zero. Supported only for point and spot lights, must be > 0.
|
||||
If None, the radius is assumed to be infinite.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, color=None, intensity=None, radius=None):
|
||||
super().__init__(name=name, color=color, intensity=intensity, radius=radius)
|
||||
|
||||
|
||||
class PointLight(Light):
|
||||
"""
|
||||
Point lights emit light in all directions from their position in space;
|
||||
rotation and scale are ignored except for their effect on the inherited
|
||||
node position. The brightness of the light attenuates in a physically
|
||||
correct manner as distance increases from the light's position (i.e.
|
||||
brightness goes like the inverse square of the distance). Point light
|
||||
intensity is defined in candela, which is lumens per square radian (lm/sr).
|
||||
|
||||
Attributes
|
||||
----------
|
||||
name : str, optional
|
||||
Name of the light.
|
||||
color : (4,) uint8
|
||||
RGBA value for the light's color in linear space.
|
||||
intensity : float
|
||||
Brightness of light. The units that this is defined in depend
|
||||
on the type of light.
|
||||
point and spot lights use luminous intensity in candela (lm/sr),
|
||||
while directional lights use illuminance in lux (lm/m2).
|
||||
radius : float
|
||||
Cutoff distance at which light's intensity may be considered to
|
||||
have reached zero. Supported only for point and spot lights, must be > 0.
|
||||
If None, the radius is assumed to be infinite.
|
||||
"""
|
||||
|
||||
def __init__(self, name=None, color=None, intensity=None, radius=None):
|
||||
super().__init__(name=name, color=color, intensity=intensity, radius=radius)
|
||||
|
||||
|
||||
class SpotLight(Light):
|
||||
"""
|
||||
Spot lights emit light in a cone in the direction of the local -z axis.
|
||||
The angle and falloff of the cone is defined using two numbers, the
|
||||
`innerConeAngle` and `outerConeAngle`. As with point lights, the brightness
|
||||
also attenuates in a physically correct manner as distance increases from
|
||||
the light's position (i.e. brightness goes like the inverse square of the
|
||||
distance). Spot light intensity refers to the brightness inside the
|
||||
`innerConeAngle` (and at the location of the light) and is defined in
|
||||
candela, which is lumens per square radian (lm/sr). A spot light's position
|
||||
and orientation are inherited from its node transform. Inherited scale does
|
||||
not affect cone shape, and is ignored except for its effect on position
|
||||
and orientation.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
name : str, optional
|
||||
Name of the light.
|
||||
color : (4,) uint8
|
||||
RGBA value for the light's color in linear space.
|
||||
intensity : float
|
||||
Brightness of light. The units that this is defined in depend
|
||||
on the type of light.
|
||||
point and spot lights use luminous intensity in candela (lm/sr),
|
||||
while directional lights use illuminance in lux (lm/m2).
|
||||
radius : float
|
||||
Cutoff distance at which light's intensity may be considered to
|
||||
have reached zero. Supported only for point and spot lights, must be > 0.
|
||||
If None, the radius is assumed to be infinite.
|
||||
innerConeAngle : float
|
||||
Angle, in radians, from centre of spotlight where falloff begins.
|
||||
Must be greater than or equal to `0` and less than `outerConeAngle`.
|
||||
outerConeAngle : float
|
||||
Angle, in radians, from centre of spotlight where falloff ends.
|
||||
Must be greater than `innerConeAngle` and less than or equal to `PI / 2.0`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name=None,
|
||||
color=None,
|
||||
intensity=None,
|
||||
radius=None,
|
||||
innerConeAngle=0.0,
|
||||
outerConeAngle=np.pi / 4.0,
|
||||
):
|
||||
super().__init__(name=name, color=color, intensity=intensity, radius=radius)
|
||||
self.outerConeAngle = outerConeAngle
|
||||
self.innerConeAngle = innerConeAngle
|
||||
|
||||
@property
|
||||
def innerConeAngle(self):
|
||||
return self._innerConeAngle
|
||||
|
||||
@innerConeAngle.setter
|
||||
def innerConeAngle(self, value):
|
||||
if value < 0.0 or value > self.outerConeAngle:
|
||||
raise ValueError("Invalid value for inner cone angle")
|
||||
self._innerConeAngle = float(value)
|
||||
|
||||
@property
|
||||
def outerConeAngle(self):
|
||||
return self._outerConeAngle
|
||||
|
||||
@outerConeAngle.setter
|
||||
def outerConeAngle(self, value):
|
||||
if value < 0.0 or value > np.pi / 2.0 + 1e-9:
|
||||
raise ValueError("Invalid value for outer cone angle")
|
||||
self._outerConeAngle = float(value)
|
||||
|
||||
|
||||
def autolight(scene) -> Tuple[List[Light], List[NDArray[float64]]]:
|
||||
"""
|
||||
Generate a list of lights for a scene that looks decent.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
scene : trimesh.Scene
|
||||
Scene with geometry
|
||||
|
||||
Returns
|
||||
--------------
|
||||
lights : [Light]
|
||||
List of light objects
|
||||
transforms : (len(lights), 4, 4) float
|
||||
Transformation matrices for light positions.
|
||||
"""
|
||||
|
||||
# start with empty lights and transforms
|
||||
lights = []
|
||||
transforms = []
|
||||
|
||||
# if there are no objects in the scene don't crash
|
||||
bounds = scene.bounds
|
||||
if bounds is not None:
|
||||
# create two translation matrices for bounds corners
|
||||
transforms.extend(transformations.translation_matrix(b) for b in bounds)
|
||||
lights.extend(PointLight() for _ in range(2))
|
||||
|
||||
return lights, transforms
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,851 @@
|
||||
import collections
|
||||
from copy import deepcopy
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import caching, util
|
||||
from ..caching import hash_fast
|
||||
from ..transformations import fix_rigid, quaternion_matrix, rotation_matrix
|
||||
from ..typed import ArrayLike, Hashable, NDArray, Optional, Sequence, Tuple, Union
|
||||
|
||||
# we compare to identity a lot
|
||||
_identity = np.eye(4)
|
||||
_identity.flags["WRITEABLE"] = False
|
||||
|
||||
|
||||
class SceneGraph:
|
||||
"""
|
||||
Hold data about positions and instances of geometry
|
||||
in a scene. This includes a forest (i.e. multi-root tree)
|
||||
of transforms and information on which node is the base
|
||||
frame, and which geometries are affiliated with which
|
||||
nodes.
|
||||
"""
|
||||
|
||||
def __init__(self, base_frame="world", repair_rigid=1e-5):
|
||||
"""
|
||||
Create a scene graph, holding homogeneous transformation
|
||||
matrices and instance information about geometry.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
base_frame : any
|
||||
The root node transforms will be positioned from.
|
||||
repair_rigid : None or float
|
||||
If a float will attempt to repair rotation matrices
|
||||
where `M @ M.T` differs from an identity matrix by
|
||||
more than floating point zero but less than this value.
|
||||
This can happen in a deep tree with a lot of matrix
|
||||
multiplies.
|
||||
"""
|
||||
# a graph structure, subclass of networkx DiGraph
|
||||
self.transforms = EnforcedForest()
|
||||
# hashable, the base or root frame
|
||||
self.base_frame = base_frame
|
||||
# if passed as a float try to repair rigid transforms
|
||||
# that have accumulated floating point error
|
||||
self.repair_rigid = repair_rigid
|
||||
# cache transformation matrices keyed with tuples
|
||||
self._cache = caching.Cache(self.__hash__)
|
||||
|
||||
def update(self, frame_to, frame_from=None, **kwargs):
|
||||
"""
|
||||
Update a transform in the tree.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
frame_from : hashable object
|
||||
Usually a string (eg 'world').
|
||||
If left as None it will be set to self.base_frame
|
||||
frame_to : hashable object
|
||||
Usually a string (eg 'mesh_0')
|
||||
matrix : (4,4) float
|
||||
Homogeneous transformation matrix
|
||||
quaternion : (4,) float
|
||||
Quaternion ordered [w, x, y, z]
|
||||
axis : (3,) float
|
||||
Axis of rotation
|
||||
angle : float
|
||||
Angle of rotation, in radians
|
||||
translation : (3,) float
|
||||
Distance to translate
|
||||
geometry : hashable
|
||||
Geometry object name, e.g. 'mesh_0'
|
||||
metadata: dictionary
|
||||
Optional metadata attached to the new frame
|
||||
(exports to glTF node 'extras').
|
||||
"""
|
||||
# if no frame specified, use base frame
|
||||
if frame_from is None:
|
||||
frame_from = self.base_frame
|
||||
|
||||
# pass through
|
||||
attr = {k: v for k, v in kwargs.items() if k in {"geometry", "metadata"}}
|
||||
# convert various kwargs to a single matrix
|
||||
attr["matrix"] = kwargs_to_matrix(**kwargs)
|
||||
|
||||
# add the edges for the transforms
|
||||
# wi ll return if it changed anything
|
||||
self.transforms.add_edge(frame_from, frame_to, **attr)
|
||||
|
||||
# set the node attribute with the geometry information
|
||||
if "geometry" in kwargs:
|
||||
self.transforms.node_data[frame_to]["geometry"] = kwargs["geometry"]
|
||||
|
||||
def get(
|
||||
self, frame_to: Hashable, frame_from: Optional[Hashable] = None
|
||||
) -> Tuple[NDArray[np.float64], Optional[Hashable]]:
|
||||
"""
|
||||
Get the transform from one frame to another.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
frame_to : hashable
|
||||
Node name, usually a string (eg 'mesh_0')
|
||||
frame_from : hashable
|
||||
Node name, usually a string (eg 'world').
|
||||
If None it will be set to self.base_frame
|
||||
|
||||
Returns
|
||||
----------
|
||||
transform : (4, 4) float
|
||||
Homogeneous transformation matrix
|
||||
geometry
|
||||
The name of the geometry if it exists
|
||||
|
||||
Raises
|
||||
-----------
|
||||
ValueError
|
||||
If the frames aren't connected.
|
||||
"""
|
||||
|
||||
# use base frame if not specified
|
||||
if frame_from is None:
|
||||
frame_from = self.base_frame
|
||||
|
||||
# look up transform to see if we have it already
|
||||
key = (frame_from, frame_to)
|
||||
if key in self._cache:
|
||||
return self._cache[key]
|
||||
|
||||
# get the geometry at the final node if any
|
||||
geometry = self.transforms.node_data[frame_to].get("geometry")
|
||||
|
||||
# get a local reference to edge data
|
||||
data = self.transforms.edge_data
|
||||
|
||||
if frame_from == frame_to:
|
||||
# if we're going from ourself return identity
|
||||
matrix = _identity
|
||||
elif key in data:
|
||||
# if the path is just an edge return early
|
||||
matrix = data[key]["matrix"]
|
||||
else:
|
||||
# we have a 3+ node path
|
||||
# get the path from the forest always going from
|
||||
# parent -> child -> child
|
||||
path = self.transforms.shortest_path(frame_from, frame_to)
|
||||
# the path should always start with `frame_from`
|
||||
assert path[0] == frame_from
|
||||
# and end with the `frame_to` node
|
||||
assert path[-1] == frame_to
|
||||
|
||||
# loop through pairs of the path
|
||||
matrices = []
|
||||
for u, v in zip(path[:-1], path[1:]):
|
||||
forward = data.get((u, v))
|
||||
if forward is not None:
|
||||
if "matrix" in forward:
|
||||
# append the matrix from u to v
|
||||
matrices.append(forward["matrix"])
|
||||
continue
|
||||
# since forwards didn't exist backward must
|
||||
# exist otherwise this is a disconnected path
|
||||
# and we should raise an error anyway
|
||||
backward = data[(v, u)]
|
||||
if "matrix" in backward:
|
||||
# append the inverted backwards matrix
|
||||
matrices.append(np.linalg.inv(backward["matrix"]))
|
||||
# filter out any identity matrices
|
||||
matrices = [m for m in matrices if np.abs(m - _identity).max() > 1e-8]
|
||||
if len(matrices) == 0:
|
||||
matrix = _identity
|
||||
elif len(matrices) == 1:
|
||||
matrix = matrices[0]
|
||||
else:
|
||||
# multiply matrices into single transform
|
||||
matrix = util.multi_dot(matrices)
|
||||
|
||||
# if instructed to repair rigid transforms do it here
|
||||
if self.repair_rigid is not None:
|
||||
matrix = fix_rigid(matrix, max_deviance=self.repair_rigid)
|
||||
|
||||
# matrix being edited in-place leads to subtle bugs
|
||||
matrix.flags["WRITEABLE"] = False
|
||||
|
||||
# store the result
|
||||
self._cache[key] = (matrix, geometry)
|
||||
|
||||
return matrix, geometry
|
||||
|
||||
def __hash__(self):
|
||||
return self.transforms.__hash__()
|
||||
|
||||
def copy(self):
|
||||
"""
|
||||
Return a copy of the current TransformForest.
|
||||
|
||||
Returns
|
||||
------------
|
||||
copied : TransformForest
|
||||
Copy of current object.
|
||||
"""
|
||||
# create a copy without transferring cache
|
||||
copied = SceneGraph()
|
||||
copied.base_frame = deepcopy(self.base_frame)
|
||||
copied.transforms = deepcopy(self.transforms)
|
||||
return copied
|
||||
|
||||
def to_flattened(self):
|
||||
"""
|
||||
Export the current transform graph with all
|
||||
transforms baked into world->instance.
|
||||
|
||||
Returns
|
||||
---------
|
||||
flat : dict
|
||||
Keyed {node : {transform, geometry}
|
||||
"""
|
||||
flat = {}
|
||||
base_frame = self.base_frame
|
||||
for node in self.nodes:
|
||||
if node == base_frame:
|
||||
continue
|
||||
# get the matrix and geometry name
|
||||
matrix, geometry = self.get(frame_to=node, frame_from=base_frame)
|
||||
# store matrix as list rather than numpy array
|
||||
flat[node] = {"transform": matrix.tolist(), "geometry": geometry}
|
||||
|
||||
return flat
|
||||
|
||||
def to_gltf(self, scene, mesh_index=None):
|
||||
"""
|
||||
Export a transforms as the 'nodes' section of the
|
||||
GLTF header dict.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
scene : trimesh.Scene
|
||||
Scene with geometry.
|
||||
mesh_index : dict or None
|
||||
Mapping { key in scene.geometry : int }
|
||||
|
||||
Returns
|
||||
--------
|
||||
gltf : dict
|
||||
With 'nodes' referencing a list of dicts
|
||||
"""
|
||||
|
||||
if mesh_index is None:
|
||||
# geometry is an OrderedDict
|
||||
# map mesh name to index: {geometry key : index}
|
||||
mesh_index = {name: i for i, name in enumerate(scene.geometry.keys())}
|
||||
|
||||
# get graph information into local scope before loop
|
||||
graph = self.transforms
|
||||
# get the stored node data
|
||||
node_data = graph.node_data
|
||||
edge_data = graph.edge_data
|
||||
base_frame = self.base_frame
|
||||
|
||||
# list of dict, in gltf format
|
||||
# start with base frame as first node index
|
||||
result = [{"name": base_frame}]
|
||||
# {node name : node index in gltf}
|
||||
lookup = {base_frame: 0}
|
||||
|
||||
# collect the nodes in order
|
||||
for node in node_data.keys():
|
||||
if node == base_frame:
|
||||
continue
|
||||
# assign the index to the node-name lookup
|
||||
lookup[node] = len(result)
|
||||
# populate a result at the correct index
|
||||
result.append({"name": node})
|
||||
|
||||
# get generated properties outside of loop
|
||||
# does the scene have a defined camera to export
|
||||
has_camera = scene.has_camera
|
||||
children = graph.children
|
||||
|
||||
extensions_used = set()
|
||||
|
||||
# then iterate through to collect data
|
||||
for info in result:
|
||||
# name of the scene node
|
||||
node = info["name"]
|
||||
|
||||
# get the original node names for children
|
||||
childs = children.get(node, [])
|
||||
if len(childs) > 0:
|
||||
info["children"] = [lookup[k] for k in childs]
|
||||
|
||||
# if we have a mesh store by index
|
||||
if "geometry" in node_data[node]:
|
||||
mesh_key = node_data[node]["geometry"]
|
||||
if mesh_key in mesh_index:
|
||||
info["mesh"] = mesh_index[mesh_key]
|
||||
# check to see if we have camera node
|
||||
if has_camera and node == scene.camera.name:
|
||||
info["camera"] = 0
|
||||
|
||||
if node != base_frame:
|
||||
parent = graph.parents[node]
|
||||
node_edge = edge_data[(parent, node)]
|
||||
|
||||
# get the matrix from this edge
|
||||
matrix = node_edge["matrix"]
|
||||
# only include if it's not an identify matrix
|
||||
if not util.allclose(matrix, _identity):
|
||||
info["matrix"] = matrix.T.reshape(-1).tolist()
|
||||
|
||||
# if an extra was stored on this edge
|
||||
extras = node_edge.get("metadata")
|
||||
if extras:
|
||||
extras = extras.copy()
|
||||
|
||||
# if extensionss were stored on this edge
|
||||
extensions = extras.pop("gltf_extensions", None)
|
||||
if isinstance(extensions, dict):
|
||||
info["extensions"] = extensions
|
||||
extensions_used = extensions_used.union(set(extensions.keys()))
|
||||
|
||||
# convert any numpy arrays to lists
|
||||
extras.update(
|
||||
{k: v.tolist() for k, v in extras.items() if hasattr(v, "tolist")}
|
||||
)
|
||||
info["extras"] = extras
|
||||
|
||||
gltf = {"nodes": result}
|
||||
if len(extensions_used) > 0:
|
||||
gltf["extensionsUsed"] = list(extensions_used)
|
||||
return gltf
|
||||
|
||||
def to_edgelist(self):
|
||||
"""
|
||||
Export the current transforms as a list of
|
||||
edge tuples, with each tuple having the format:
|
||||
(node_a, node_b, {metadata})
|
||||
|
||||
Returns
|
||||
---------
|
||||
edgelist : (n,) list
|
||||
Of edge tuples
|
||||
"""
|
||||
# save local reference to node_data
|
||||
nodes = self.transforms.node_data
|
||||
# save cleaned edges
|
||||
export = []
|
||||
# loop through (node, node, edge attributes)
|
||||
for edge, attr in self.transforms.edge_data.items():
|
||||
# node indexes from edge
|
||||
a, b = edge
|
||||
# geometry is a node property but save it to the
|
||||
# edge so we don't need two dictionaries
|
||||
b_attr = nodes[b]
|
||||
# make sure we're not stomping on original
|
||||
attr_new = attr.copy()
|
||||
# apply node geometry to edge attributes
|
||||
if "geometry" in b_attr:
|
||||
attr_new["geometry"] = b_attr["geometry"]
|
||||
# convert any numpy arrays to regular lists
|
||||
attr_new.update(
|
||||
{k: v.tolist() for k, v in attr_new.items() if hasattr(v, "tolist")}
|
||||
)
|
||||
export.append([a, b, attr_new])
|
||||
return export
|
||||
|
||||
def from_edgelist(self, edges, strict=True):
|
||||
"""
|
||||
Load transform data from an edge list into the current
|
||||
scene graph.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
edgelist : (n,) tuples
|
||||
Keyed (node_a, node_b, {key: value})
|
||||
strict : bool
|
||||
If True raise a ValueError when a
|
||||
malformed edge is passed in a tuple.
|
||||
"""
|
||||
|
||||
# loop through each edge
|
||||
for edge in edges:
|
||||
# edge contains attributes
|
||||
if len(edge) == 3:
|
||||
self.update(edge[1], edge[0], **edge[2])
|
||||
# edge just contains nodes
|
||||
elif len(edge) == 2:
|
||||
self.update(edge[1], edge[0])
|
||||
# edge is broken
|
||||
elif strict:
|
||||
raise ValueError("edge incorrect shape: %s", str(edge))
|
||||
|
||||
def to_networkx(self):
|
||||
"""
|
||||
Return a `networkx` copy of this graph.
|
||||
|
||||
Returns
|
||||
----------
|
||||
graph : networkx.DiGraph
|
||||
Directed graph.
|
||||
"""
|
||||
import networkx
|
||||
|
||||
return networkx.from_edgelist(self.to_edgelist(), create_using=networkx.DiGraph)
|
||||
|
||||
def show(self, **kwargs):
|
||||
"""
|
||||
Plot the scene graph using `networkx.draw_networkx`
|
||||
which uses matplotlib to display the graph.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
kwargs : dict
|
||||
Passed to `networkx.draw_networkx`
|
||||
"""
|
||||
import matplotlib.pyplot as plt
|
||||
import networkx
|
||||
|
||||
# default kwargs will only be set if not
|
||||
# passed explicitly to the show command
|
||||
defaults = {"with_labels": True}
|
||||
kwargs.update(**{k: v for k, v in defaults.items() if k not in kwargs})
|
||||
networkx.draw_networkx(G=self.to_networkx(), **kwargs)
|
||||
|
||||
plt.show()
|
||||
|
||||
def load(self, edgelist):
|
||||
"""
|
||||
Load transform data from an edge list into the current
|
||||
scene graph.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
edgelist : (n,) tuples
|
||||
Structured (node_a, node_b, {key: value})
|
||||
"""
|
||||
self.from_edgelist(edgelist, strict=True)
|
||||
|
||||
@caching.cache_decorator
|
||||
def nodes(self):
|
||||
"""
|
||||
A list of every node in the graph.
|
||||
|
||||
Returns
|
||||
-------------
|
||||
nodes : (n,) array
|
||||
All node names.
|
||||
"""
|
||||
return self.transforms.nodes
|
||||
|
||||
@caching.cache_decorator
|
||||
def nodes_geometry(self):
|
||||
"""
|
||||
The nodes in the scene graph with geometry attached.
|
||||
|
||||
Returns
|
||||
------------
|
||||
nodes_geometry : (m,) array
|
||||
Node names which have geometry associated
|
||||
"""
|
||||
return [n for n, attr in self.transforms.node_data.items() if "geometry" in attr]
|
||||
|
||||
@caching.cache_decorator
|
||||
def geometry_nodes(self):
|
||||
"""
|
||||
Which nodes have this geometry? Inverse
|
||||
of `nodes_geometry`.
|
||||
|
||||
Returns
|
||||
------------
|
||||
geometry_nodes : dict
|
||||
Keyed {geometry_name : node name}
|
||||
"""
|
||||
res = collections.defaultdict(list)
|
||||
for node, attr in self.transforms.node_data.items():
|
||||
if "geometry" in attr:
|
||||
res[attr["geometry"]].append(node)
|
||||
return res
|
||||
|
||||
def remove_geometries(self, geometries: Union[str, set, Sequence]):
|
||||
"""
|
||||
Remove the reference for specified geometries
|
||||
from nodes without deleting the node.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
geometries : list or str
|
||||
Name of scene.geometry to dereference.
|
||||
"""
|
||||
# make sure we have a set of geometries to remove
|
||||
if isinstance(geometries, str):
|
||||
geometries = [geometries]
|
||||
geometries = set(geometries)
|
||||
|
||||
# remove the geometry reference from the node without deleting nodes
|
||||
# this lets us keep our cached paths, and will not screw up children
|
||||
for attrib in self.transforms.node_data.values():
|
||||
if "geometry" in attrib and attrib["geometry"] in geometries:
|
||||
attrib.pop("geometry")
|
||||
|
||||
# it would be safer to just run _cache.clear
|
||||
# but the only property using the geometry should be
|
||||
# nodes_geometry: if this becomes not true change this to clear!
|
||||
self._cache.cache.pop("nodes_geometry", None)
|
||||
self.transforms._hash = None
|
||||
|
||||
def __contains__(self, key: Hashable) -> bool:
|
||||
return key in self.transforms.node_data
|
||||
|
||||
def __getitem__(
|
||||
self, key: Hashable
|
||||
) -> Tuple[NDArray[np.float64], Optional[Hashable]]:
|
||||
return self.get(key)
|
||||
|
||||
def __setitem__(self, key: Hashable, value: ArrayLike):
|
||||
value = np.asanyarray(value, dtype=np.float64)
|
||||
if value.shape != (4, 4):
|
||||
raise ValueError("Matrix must be specified!")
|
||||
return self.update(key, matrix=value)
|
||||
|
||||
def clear(self):
|
||||
self.transforms = EnforcedForest()
|
||||
self._cache.clear()
|
||||
|
||||
|
||||
class EnforcedForest:
|
||||
"""
|
||||
A simple forest graph data structure: every node
|
||||
is allowed to have exactly one parent. This makes
|
||||
traversal and implementation much simpler than a
|
||||
full graph data type; by storing only one parent
|
||||
reference, it enforces the structure for "free."
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# since every node can have only one parent
|
||||
# this data structure transparently enforces
|
||||
# the forest data structure without checks
|
||||
# a dict {child : parent}
|
||||
self.parents = {}
|
||||
|
||||
# store data for a particular edge keyed by tuple
|
||||
# {(u, v) : data }
|
||||
self.edge_data = collections.defaultdict(dict)
|
||||
# {u: data}
|
||||
self.node_data = collections.defaultdict(dict)
|
||||
|
||||
# if multiple calls are made for the same path
|
||||
# but the connectivity hasn't changed return cached
|
||||
self._cache = {}
|
||||
|
||||
def add_edge(self, u, v, **kwargs):
|
||||
"""
|
||||
Add an edge to the forest cleanly.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
u : any
|
||||
Hashable node key.
|
||||
v : any
|
||||
Hashable node key.
|
||||
kwargs : dict
|
||||
Stored as (u, v) edge data.
|
||||
|
||||
Returns
|
||||
--------
|
||||
changed : bool
|
||||
Return if this operation changed anything.
|
||||
"""
|
||||
self._hash = None
|
||||
|
||||
# topology has changed so clear cache
|
||||
if (u, v) not in self.edge_data:
|
||||
self._cache = {}
|
||||
else:
|
||||
# check to see if matrix and geometry are identical
|
||||
edge = self.edge_data[(u, v)]
|
||||
if util.allclose(
|
||||
kwargs.get("matrix", _identity), edge.get("matrix", _identity), 1e-8
|
||||
) and (edge.get("geometry") == kwargs.get("geometry")):
|
||||
return False
|
||||
|
||||
# store a parent reference for traversal
|
||||
self.parents[v] = u
|
||||
# store kwargs for edge data keyed with tuple
|
||||
self.edge_data[(u, v)] = kwargs
|
||||
# set empty node data
|
||||
self.node_data[u].update({})
|
||||
if "geometry" in kwargs:
|
||||
self.node_data[v].update({"geometry": kwargs["geometry"]})
|
||||
else:
|
||||
self.node_data[v].update({})
|
||||
|
||||
return True
|
||||
|
||||
def remove_node(self, u):
|
||||
"""
|
||||
Remove a node from the forest.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
u : any
|
||||
Hashable node key.
|
||||
|
||||
Returns
|
||||
--------
|
||||
changed : bool
|
||||
Return if this operation changed anything.
|
||||
"""
|
||||
# check if node is part of forest
|
||||
if u not in self.node_data:
|
||||
return False
|
||||
|
||||
# topology will change so clear cache
|
||||
self._cache = {}
|
||||
self._hash = None
|
||||
|
||||
# delete all children's references and parent reference
|
||||
children = [child for (child, parent) in self.parents.items() if parent == u]
|
||||
for c in children:
|
||||
del self.parents[c]
|
||||
if u in self.parents:
|
||||
del self.parents[u]
|
||||
|
||||
# delete edge data
|
||||
edges = [(a, b) for (a, b) in self.edge_data if a == u or b == u]
|
||||
for e in edges:
|
||||
del self.edge_data[e]
|
||||
|
||||
# delete node data
|
||||
del self.node_data[u]
|
||||
|
||||
return True
|
||||
|
||||
def shortest_path(self, u, v):
|
||||
"""
|
||||
Find the shortest path between `u` and `v`, returning
|
||||
a path where the first element is always `u` and the
|
||||
last element is always `v`, disregarding edge direction.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
u : any
|
||||
Hashable node key.
|
||||
v : any
|
||||
Hashable node key.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
path : (n,)
|
||||
Path between `u` and `v`
|
||||
"""
|
||||
# see if we've already computed this path
|
||||
if u == v:
|
||||
# the path between itself is an edge case
|
||||
return []
|
||||
elif (u, v) in self._cache:
|
||||
# return the same path for either direction
|
||||
return self._cache[(u, v)]
|
||||
elif (v, u) in self._cache:
|
||||
return self._cache[(v, u)][::-1]
|
||||
|
||||
# local reference to parent dict for performance
|
||||
parents = self.parents
|
||||
# store both forward and backwards traversal
|
||||
forward = [u]
|
||||
backward = [v]
|
||||
|
||||
# cap iteration to number of total nodes
|
||||
for _ in range(len(parents) + 1):
|
||||
# store the parent both forwards and backwards
|
||||
f = parents.get(forward[-1])
|
||||
b = parents.get(backward[-1])
|
||||
forward.append(f)
|
||||
backward.append(b)
|
||||
|
||||
if f == v:
|
||||
self._cache[(u, v)] = forward
|
||||
return forward
|
||||
elif b == u:
|
||||
# return reversed path
|
||||
backward = backward[::-1]
|
||||
self._cache[(u, v)] = backward
|
||||
return backward
|
||||
elif (b in forward) or (f is None and b is None):
|
||||
# we have a either a common node between both
|
||||
# traversal directions or we have consumed the whole
|
||||
# tree in both directions so try to find the common node
|
||||
common = set(backward).intersection(forward).difference({None})
|
||||
if len(common) == 0:
|
||||
raise ValueError(f"No path from {u}->{v}!")
|
||||
elif len(common) > 1:
|
||||
# get the first occurring common element in "forward"
|
||||
link = next(f for f in forward if f in common)
|
||||
assert link in common
|
||||
else:
|
||||
# take the only common element
|
||||
link = next(iter(common))
|
||||
|
||||
# combine the forward and backwards traversals
|
||||
a = forward[: forward.index(link) + 1]
|
||||
b = backward[: backward.index(link)]
|
||||
path = a + b[::-1]
|
||||
|
||||
# verify we didn't screw up the order
|
||||
assert path[0] == u
|
||||
assert path[-1] == v
|
||||
|
||||
self._cache[(u, v)] = path
|
||||
|
||||
return path
|
||||
|
||||
raise ValueError("Iteration limit exceeded!")
|
||||
|
||||
@property
|
||||
def nodes(self):
|
||||
"""
|
||||
Get a set of every node.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
nodes : set
|
||||
Every node currently stored.
|
||||
"""
|
||||
return self.node_data.keys()
|
||||
|
||||
@property
|
||||
def children(self):
|
||||
"""
|
||||
Get the children of each node.
|
||||
|
||||
Returns
|
||||
----------
|
||||
children : dict
|
||||
Keyed {node : [child, child, ...]}
|
||||
"""
|
||||
if "children" in self._cache:
|
||||
return self._cache["children"]
|
||||
child = collections.defaultdict(list)
|
||||
# append children to parent references
|
||||
# skip self-references to avoid a node loop
|
||||
[child[v].append(u) for u, v in self.parents.items() if u != v]
|
||||
|
||||
# cache and return as a vanilla dict
|
||||
self._cache["children"] = dict(child)
|
||||
return self._cache["children"]
|
||||
|
||||
def successors(self, node):
|
||||
"""
|
||||
Get all nodes that are successors to specified node,
|
||||
including the specified node.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
node : any
|
||||
Hashable key for a node.
|
||||
|
||||
Returns
|
||||
------------
|
||||
successors : set
|
||||
Nodes that succeed specified node.
|
||||
"""
|
||||
# get mapping of {parent : child}
|
||||
children = self.children
|
||||
# if node doesn't exist return early
|
||||
if node not in children:
|
||||
return {node}
|
||||
|
||||
# children we need to collect
|
||||
queue = [node]
|
||||
# start collecting values with children of source
|
||||
collected = set(queue)
|
||||
|
||||
# cap maximum iterations
|
||||
for _ in range(len(self.node_data) + 1):
|
||||
if len(queue) == 0:
|
||||
# no more nodes to visit so we're done
|
||||
return collected
|
||||
# add the children of this node to be processed
|
||||
childs = children.get(queue.pop())
|
||||
if childs is not None:
|
||||
queue.extend(childs)
|
||||
collected.update(childs)
|
||||
return collected
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
Actually hash all of the data, but use a "dirty" mechanism
|
||||
in functions that modify the data, which MUST
|
||||
# all invalidate the hash by setting `self._hash = None`
|
||||
|
||||
This was optimized a bit, and is evaluating on an
|
||||
older laptop on a scene with 77 nodes and 76 edges
|
||||
10,000 times in 0.7s which seems fast enough.
|
||||
"""
|
||||
# see if there is an available hash value
|
||||
# if you are seeing cache bugs this is the thing
|
||||
# to try eliminating because it is very likely that
|
||||
# someone somewhere is modifying the data without
|
||||
# setting `self._hash = None`
|
||||
hashed = getattr(self, "_hash", None)
|
||||
if hashed is not None:
|
||||
return hashed
|
||||
|
||||
hashed = hash_fast(
|
||||
(
|
||||
"".join(
|
||||
str(hash(k)) + v.get("geometry", "")
|
||||
for k, v in self.edge_data.items()
|
||||
)
|
||||
+ "".join(
|
||||
str(k) + v.get("geometry", "") for k, v in self.node_data.items()
|
||||
)
|
||||
).encode("utf-8")
|
||||
+ b"".join(
|
||||
v["matrix"].tobytes() for v in self.edge_data.values() if "matrix" in v
|
||||
)
|
||||
)
|
||||
self._hash = hashed
|
||||
return hashed
|
||||
|
||||
|
||||
def kwargs_to_matrix(
|
||||
matrix=None, quaternion=None, translation=None, axis=None, angle=None, **kwargs
|
||||
):
|
||||
"""
|
||||
Take multiple keyword arguments and parse them
|
||||
into a homogeneous transformation matrix.
|
||||
|
||||
Returns
|
||||
---------
|
||||
matrix : (4, 4) float
|
||||
Homogeneous transformation matrix.
|
||||
"""
|
||||
if matrix is not None:
|
||||
# a matrix takes immediate precedence over other options
|
||||
return np.array(matrix, dtype=np.float64)
|
||||
elif quaternion is not None:
|
||||
matrix = quaternion_matrix(quaternion)
|
||||
elif axis is not None and angle is not None:
|
||||
matrix = rotation_matrix(angle, axis)
|
||||
else:
|
||||
matrix = np.eye(4)
|
||||
|
||||
if translation is not None:
|
||||
# translation can be used in conjunction with any
|
||||
# of the methods specifying transforms
|
||||
matrix[:3, 3] += translation
|
||||
|
||||
return matrix
|
||||
Reference in New Issue
Block a user