init
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
from . import creation
|
||||
from .base import VoxelGrid
|
||||
|
||||
__all__ = ["VoxelGrid", "creation"]
|
||||
@@ -0,0 +1,427 @@
|
||||
"""
|
||||
voxel.py
|
||||
-----------
|
||||
|
||||
Convert meshes to a simple voxel data structure and back again.
|
||||
"""
|
||||
|
||||
from hashlib import sha256
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import bounds as bounds_module
|
||||
from .. import caching, util
|
||||
from .. import transformations as tr
|
||||
from ..constants import log
|
||||
from ..exchange.binvox import export_binvox
|
||||
from ..parent import Geometry
|
||||
from . import morphology, ops, transforms
|
||||
from .encoding import DenseEncoding, Encoding
|
||||
|
||||
|
||||
class VoxelGrid(Geometry):
|
||||
"""
|
||||
Store 3D voxels.
|
||||
"""
|
||||
|
||||
def __init__(self, encoding, transform=None, metadata=None):
|
||||
if transform is None:
|
||||
transform = np.eye(4)
|
||||
if isinstance(encoding, np.ndarray):
|
||||
encoding = DenseEncoding(encoding.astype(bool))
|
||||
if encoding.dtype != bool:
|
||||
raise ValueError("encoding must have dtype bool")
|
||||
self._data = caching.DataStore()
|
||||
self.encoding = encoding
|
||||
self._transform = transforms.Transform(transform, datastore=self._data)
|
||||
self._cache = caching.Cache(id_function=self._data.__hash__)
|
||||
|
||||
self.metadata = {}
|
||||
# update the mesh metadata with passed metadata
|
||||
if isinstance(metadata, dict):
|
||||
self.metadata.update(metadata)
|
||||
elif metadata is not None:
|
||||
raise ValueError(f"metadata should be a dict or None, got {metadata!s}")
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
Get the hash of the current transformation matrix.
|
||||
|
||||
Returns
|
||||
------------
|
||||
hash : str
|
||||
Hash of transformation matrix
|
||||
"""
|
||||
return self._data.__hash__()
|
||||
|
||||
@property
|
||||
def identifier_hash(self) -> str:
|
||||
return sha256(hash(self).to_bytes()).hexdigest()
|
||||
|
||||
@property
|
||||
def encoding(self):
|
||||
"""
|
||||
`Encoding` object providing the occupancy grid.
|
||||
|
||||
See `trimesh.voxel.encoding` for implementations.
|
||||
"""
|
||||
return self._data["encoding"]
|
||||
|
||||
@encoding.setter
|
||||
def encoding(self, encoding):
|
||||
if isinstance(encoding, np.ndarray):
|
||||
encoding = DenseEncoding(encoding)
|
||||
elif not isinstance(encoding, Encoding):
|
||||
raise ValueError(f"encoding must be an Encoding, got {encoding!s}")
|
||||
if len(encoding.shape) != 3:
|
||||
raise ValueError(f"encoding must be rank 3, got shape {encoding.shape!s}")
|
||||
if encoding.dtype != bool:
|
||||
raise ValueError(f"encoding must be binary, got {encoding.dtype}")
|
||||
self._data["encoding"] = encoding
|
||||
|
||||
@property
|
||||
def transform(self):
|
||||
"""4x4 homogeneous transformation matrix."""
|
||||
return self._transform.matrix
|
||||
|
||||
@transform.setter
|
||||
def transform(self, matrix):
|
||||
"""4x4 homogeneous transformation matrix."""
|
||||
self._transform.matrix = matrix
|
||||
|
||||
@property
|
||||
def translation(self):
|
||||
"""Location of voxel at [0, 0, 0]."""
|
||||
return self._transform.translation
|
||||
|
||||
@property
|
||||
def scale(self):
|
||||
"""
|
||||
3-element float representing per-axis scale.
|
||||
|
||||
Raises a `RuntimeError` if `self.transform` has rotation or
|
||||
shear components.
|
||||
"""
|
||||
return self._transform.scale
|
||||
|
||||
@property
|
||||
def pitch(self):
|
||||
"""
|
||||
Uniform scaling factor representing the side length of
|
||||
each voxel.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
pitch : float
|
||||
Pitch of the voxels.
|
||||
|
||||
Raises
|
||||
------------
|
||||
`RuntimeError`
|
||||
If `self.transformation` has rotation or shear
|
||||
components of has non-uniform scaling.
|
||||
"""
|
||||
return self._transform.pitch
|
||||
|
||||
@property
|
||||
def element_volume(self):
|
||||
return self._transform.unit_volume
|
||||
|
||||
def apply_transform(self, matrix):
|
||||
self._transform.apply_transform(matrix)
|
||||
return self
|
||||
|
||||
def strip(self):
|
||||
"""
|
||||
Mutate self by stripping leading/trailing planes of zeros.
|
||||
|
||||
Returns
|
||||
--------
|
||||
self after mutation occurs in-place
|
||||
"""
|
||||
encoding, padding = self.encoding.stripped
|
||||
self.encoding = encoding
|
||||
self._transform.matrix[:3, 3] = self.indices_to_points(padding[:, 0])
|
||||
return self
|
||||
|
||||
@caching.cache_decorator
|
||||
def bounds(self):
|
||||
indices = self.sparse_indices
|
||||
# get all 8 corners of the AABB
|
||||
corners = bounds_module.corners(
|
||||
[indices.min(axis=0) - 0.5, indices.max(axis=0) + 0.5]
|
||||
)
|
||||
# transform these corners to a new frame
|
||||
corners = self._transform.transform_points(corners)
|
||||
# get the AABB of corners in-frame
|
||||
bounds = np.array([corners.min(axis=0), corners.max(axis=0)])
|
||||
bounds.flags.writeable = False
|
||||
return bounds
|
||||
|
||||
@caching.cache_decorator
|
||||
def extents(self):
|
||||
bounds = self.bounds
|
||||
extents = bounds[1] - bounds[0]
|
||||
extents.flags.writeable = False
|
||||
return extents
|
||||
|
||||
@caching.cache_decorator
|
||||
def is_empty(self):
|
||||
return self.encoding.is_empty
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
"""3-tuple of ints denoting shape of occupancy grid."""
|
||||
return self.encoding.shape
|
||||
|
||||
@caching.cache_decorator
|
||||
def filled_count(self):
|
||||
"""int, number of occupied voxels in the grid."""
|
||||
return self.encoding.sum.item()
|
||||
|
||||
def is_filled(self, point):
|
||||
"""
|
||||
Query points to see if the voxel cells they lie in are
|
||||
filled or not.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
point : (n, 3) float
|
||||
Points in space
|
||||
|
||||
Returns
|
||||
---------
|
||||
is_filled : (n,) bool
|
||||
Is cell occupied or not for each point
|
||||
"""
|
||||
point = np.asanyarray(point)
|
||||
indices = self.points_to_indices(point)
|
||||
in_range = np.logical_and(
|
||||
np.all(indices < np.array(self.shape), axis=-1), np.all(indices >= 0, axis=-1)
|
||||
)
|
||||
|
||||
is_filled = np.zeros_like(in_range)
|
||||
is_filled[in_range] = self.encoding.gather_nd(indices[in_range])
|
||||
return is_filled
|
||||
|
||||
def fill(self, method="holes", **kwargs):
|
||||
"""
|
||||
Mutates self by filling in the encoding according
|
||||
to `morphology.fill`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
method : hashable
|
||||
Implementation key, one of
|
||||
`trimesh.voxel.morphology.fill.fillers` keys
|
||||
**kwargs : dict
|
||||
Additional kwargs passed through to
|
||||
the keyed implementation.
|
||||
|
||||
Returns
|
||||
----------
|
||||
self : VoxelGrid
|
||||
After replacing encoding with a filled version.
|
||||
"""
|
||||
self.encoding = morphology.fill(self.encoding, method=method, **kwargs)
|
||||
return self
|
||||
|
||||
def hollow(self):
|
||||
"""
|
||||
Mutates self by removing internal voxels
|
||||
leaving only surface elements.
|
||||
|
||||
Surviving elements are those in encoding that are
|
||||
adjacent to an empty voxel where adjacency is
|
||||
controlled by `structure`.
|
||||
|
||||
Returns
|
||||
----------
|
||||
self : VoxelGrid
|
||||
After replacing encoding with a surface version.
|
||||
"""
|
||||
self.encoding = morphology.surface(self.encoding)
|
||||
return self
|
||||
|
||||
@caching.cache_decorator
|
||||
def marching_cubes(self):
|
||||
"""
|
||||
A marching cubes Trimesh representation of the voxels.
|
||||
|
||||
No effort was made to clean or smooth the result in any way;
|
||||
it is merely the result of applying the scikit-image
|
||||
measure.marching_cubes function to self.encoding.dense.
|
||||
|
||||
Returns
|
||||
---------
|
||||
meshed : trimesh.Trimesh
|
||||
Representing the current voxel
|
||||
object as returned by marching cubes algorithm.
|
||||
"""
|
||||
return ops.matrix_to_marching_cubes(matrix=self.matrix)
|
||||
|
||||
@property
|
||||
def matrix(self):
|
||||
"""
|
||||
Return a DENSE matrix of the current voxel encoding.
|
||||
|
||||
Returns
|
||||
-------------
|
||||
dense : (a, b, c) bool
|
||||
Numpy array of dense matrix
|
||||
Shortcut to voxel.encoding.dense
|
||||
"""
|
||||
return self.encoding.dense
|
||||
|
||||
@caching.cache_decorator
|
||||
def volume(self):
|
||||
"""
|
||||
What is the volume of the filled cells in the current
|
||||
voxel object.
|
||||
|
||||
Returns
|
||||
---------
|
||||
volume : float
|
||||
Volume of filled cells.
|
||||
"""
|
||||
return self.filled_count * self.element_volume
|
||||
|
||||
@caching.cache_decorator
|
||||
def points(self):
|
||||
"""
|
||||
The center of each filled cell as a list of points.
|
||||
|
||||
Returns
|
||||
----------
|
||||
points : (self.filled, 3) float
|
||||
Points in space.
|
||||
"""
|
||||
return self._transform.transform_points(self.sparse_indices.astype(float))
|
||||
|
||||
@property
|
||||
def sparse_indices(self):
|
||||
"""(n, 3) int array of sparse indices of occupied voxels."""
|
||||
return self.encoding.sparse_indices
|
||||
|
||||
def as_boxes(self, colors=None, **kwargs):
|
||||
"""
|
||||
A rough Trimesh representation of the voxels with a box
|
||||
for each filled voxel.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
colors : None, (3,) or (4,) float or uint8
|
||||
(X, Y, Z, 3) or (X, Y, Z, 4) float or uint8
|
||||
Where matrix.shape == (X, Y, Z)
|
||||
|
||||
Returns
|
||||
---------
|
||||
mesh : trimesh.Trimesh
|
||||
Mesh with one box per filled cell.
|
||||
"""
|
||||
|
||||
if colors is not None:
|
||||
colors = np.asanyarray(colors)
|
||||
if colors.ndim == 4:
|
||||
encoding = self.encoding
|
||||
if colors.shape[:3] == encoding.shape:
|
||||
# TODO jackd: more efficient implementation?
|
||||
# encoding.as_mask?
|
||||
colors = colors[encoding.dense]
|
||||
else:
|
||||
log.warning("colors incorrect shape!")
|
||||
colors = None
|
||||
elif colors.shape not in ((3,), (4,)):
|
||||
log.warning("colors incorrect shape!")
|
||||
colors = None
|
||||
|
||||
mesh = ops.multibox(centers=self.sparse_indices.astype(float), colors=colors)
|
||||
|
||||
mesh = mesh.apply_transform(self.transform)
|
||||
return mesh
|
||||
|
||||
def points_to_indices(self, points):
|
||||
"""
|
||||
Convert points to indices in the matrix array.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points: (n, 3) float, point in space
|
||||
|
||||
Returns
|
||||
---------
|
||||
indices: (n, 3) int array of indices into self.encoding
|
||||
"""
|
||||
points = self._transform.inverse_transform_points(points)
|
||||
return np.round(points).astype(int)
|
||||
|
||||
def indices_to_points(self, indices):
|
||||
return self._transform.transform_points(indices.astype(float))
|
||||
|
||||
def show(self, *args, **kwargs):
|
||||
"""
|
||||
Convert the current set of voxels into a trimesh for visualization
|
||||
and show that via its built- in preview method.
|
||||
"""
|
||||
return self.as_boxes(kwargs.pop("colors", None)).show(*args, **kwargs)
|
||||
|
||||
def copy(self):
|
||||
return VoxelGrid(self.encoding.copy(), self._transform.matrix.copy())
|
||||
|
||||
def export(self, file_obj=None, file_type=None, **kwargs):
|
||||
"""
|
||||
Export the current VoxelGrid.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
file_obj : file-like or str
|
||||
File or file-name to export to.
|
||||
file_type : None or str
|
||||
Only 'binvox' currently supported.
|
||||
|
||||
Returns
|
||||
---------
|
||||
export : bytes
|
||||
Value of export.
|
||||
"""
|
||||
if isinstance(file_obj, str) and file_type is None:
|
||||
file_type = util.split_extension(file_obj).lower()
|
||||
|
||||
if file_type != "binvox":
|
||||
raise ValueError("only binvox exports supported!")
|
||||
|
||||
exported = export_binvox(self, **kwargs)
|
||||
if hasattr(file_obj, "write"):
|
||||
file_obj.write(exported)
|
||||
elif isinstance(file_obj, str):
|
||||
with open(file_obj, "wb") as f:
|
||||
f.write(exported)
|
||||
return exported
|
||||
|
||||
def revoxelized(self, shape):
|
||||
"""
|
||||
Create a new VoxelGrid without rotations, reflections
|
||||
or shearing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : (3, int)
|
||||
The shape of the returned VoxelGrid.
|
||||
|
||||
Returns
|
||||
----------
|
||||
vox : VoxelGrid
|
||||
Of the given shape with possibly non-uniform
|
||||
scale and translation transformation matrix.
|
||||
"""
|
||||
shape = tuple(shape)
|
||||
bounds = self.bounds.copy()
|
||||
extents = self.extents
|
||||
points = util.grid_linspace(bounds, shape).reshape(shape + (3,))
|
||||
dense = self.is_filled(points)
|
||||
scale = extents / np.asanyarray(shape)
|
||||
translate = bounds[0]
|
||||
return VoxelGrid(dense, transform=tr.scale_and_translate(scale, translate))
|
||||
|
||||
def __add__(self, other):
|
||||
raise NotImplementedError("TODO : implement voxel concatenation")
|
||||
@@ -0,0 +1,318 @@
|
||||
import numpy as np
|
||||
|
||||
from .. import grouping, remesh, util
|
||||
from .. import transformations as tr
|
||||
from ..constants import log_time
|
||||
from ..typed import ArrayLike, Integer, Number, Optional, VoxelizationMethodsType
|
||||
from . import base
|
||||
from . import encoding as enc
|
||||
|
||||
|
||||
@log_time
|
||||
def voxelize_subdivide(
|
||||
mesh, pitch: Number, max_iter: Optional[Integer] = 10, edge_factor: Number = 2.0
|
||||
) -> base.VoxelGrid:
|
||||
"""
|
||||
Voxelize a surface by subdividing a mesh until every edge is
|
||||
shorter than: (pitch / edge_factor)
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
mesh : trimesh.Trimesh
|
||||
Source mesh
|
||||
pitch
|
||||
Side length of a single voxel cube
|
||||
max_iter
|
||||
Cap maximum subdivisions or None for no limit.
|
||||
edge_factor
|
||||
Proportion of pitch maximum edge length.
|
||||
|
||||
Returns
|
||||
-----------
|
||||
VoxelGrid instance representing the voxelized mesh.
|
||||
"""
|
||||
max_edge = pitch / edge_factor
|
||||
|
||||
if max_iter is None:
|
||||
longest_edge = np.linalg.norm(
|
||||
mesh.vertices[mesh.edges[:, 0]] - mesh.vertices[mesh.edges[:, 1]], axis=1
|
||||
).max()
|
||||
max_iter = max(int(np.ceil(np.log2(longest_edge / max_edge))), 0)
|
||||
|
||||
# get the same mesh sudivided so every edge is shorter
|
||||
# than a factor of our pitch
|
||||
v, _f, _idx = remesh.subdivide_to_size(
|
||||
mesh.vertices, mesh.faces, max_edge=max_edge, max_iter=max_iter, return_index=True
|
||||
)
|
||||
|
||||
# convert the vertices to their voxel grid position
|
||||
# Provided edge_factor > 1 and max_iter is large enough, this is
|
||||
# sufficient to preserve 6-connectivity at the level of voxels.
|
||||
hit = np.round(v / pitch).astype(int)
|
||||
|
||||
# remove duplicates
|
||||
unique, _inverse = grouping.unique_rows(hit)
|
||||
|
||||
# get the voxel centers in model space
|
||||
occupied_index = hit[unique]
|
||||
|
||||
origin_index = occupied_index.min(axis=0)
|
||||
origin_position = origin_index * pitch
|
||||
|
||||
return base.VoxelGrid(
|
||||
enc.SparseBinaryEncoding(occupied_index - origin_index),
|
||||
transform=tr.scale_and_translate(scale=pitch, translate=origin_position),
|
||||
)
|
||||
|
||||
|
||||
def local_voxelize(
|
||||
mesh,
|
||||
point: ArrayLike,
|
||||
pitch: Number,
|
||||
radius: Number,
|
||||
fill: bool = True,
|
||||
**kwargs,
|
||||
) -> Optional[base.VoxelGrid]:
|
||||
"""
|
||||
Voxelize a mesh in the region of a cube around a point. When fill=True,
|
||||
uses proximity.contains to fill the resulting voxels so may be meaningless
|
||||
for non-watertight meshes. Useful to reduce memory cost for small values of
|
||||
pitch as opposed to global voxelization.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
mesh : trimesh.Trimesh
|
||||
Source geometry
|
||||
point : (3, ) float
|
||||
Point in space to voxelize around
|
||||
pitch
|
||||
Side length of a single voxel cube
|
||||
radius
|
||||
Number of voxel cubes to return in each direction.
|
||||
kwargs
|
||||
Parameters to pass to voxelize_subdivide
|
||||
|
||||
Returns
|
||||
-----------
|
||||
voxels : VoxelGrid instance with resolution (m, m, m) where m=2*radius+1
|
||||
or None if the volume is empty
|
||||
"""
|
||||
from scipy import ndimage
|
||||
|
||||
# make sure point is correct type/shape
|
||||
point = np.asanyarray(point, dtype=np.float64).reshape(3)
|
||||
# this is a gotcha- radius sounds a lot like it should be in
|
||||
# float model space, not int voxel space so check
|
||||
if not isinstance(radius, int):
|
||||
raise ValueError("radius needs to be an integer number of cubes!")
|
||||
|
||||
# Bounds of region
|
||||
bounds = np.concatenate(
|
||||
(point - (radius + 0.5) * pitch, point + (radius + 0.5) * pitch)
|
||||
)
|
||||
|
||||
# faces that intersect axis aligned bounding box
|
||||
faces = list(mesh.triangles_tree.intersection(bounds))
|
||||
|
||||
# didn't hit anything so exit
|
||||
if len(faces) == 0:
|
||||
return None
|
||||
|
||||
local = mesh.submesh([[f] for f in faces], append=True)
|
||||
|
||||
# Translate mesh so point is at 0,0,0
|
||||
local.apply_translation(-point)
|
||||
|
||||
# sparse, origin = voxelize_subdivide(local, pitch, **kwargs)
|
||||
vox = voxelize_subdivide(local, pitch, **kwargs)
|
||||
origin = vox.transform[:3, 3]
|
||||
matrix = vox.encoding.dense
|
||||
|
||||
# Find voxel index for point
|
||||
center = np.round(-origin / pitch).astype(np.int64)
|
||||
|
||||
# pad matrix if necessary
|
||||
prepad = np.maximum(radius - center, 0)
|
||||
postpad = np.maximum(center + radius + 1 - matrix.shape, 0)
|
||||
|
||||
matrix = np.pad(matrix, np.stack((prepad, postpad), axis=-1), mode="constant")
|
||||
center += prepad
|
||||
|
||||
# Extract voxels within the bounding box
|
||||
voxels = matrix[
|
||||
center[0] - radius : center[0] + radius + 1,
|
||||
center[1] - radius : center[1] + radius + 1,
|
||||
center[2] - radius : center[2] + radius + 1,
|
||||
]
|
||||
local_origin = point - radius * pitch # origin of local voxels
|
||||
|
||||
# Fill internal regions
|
||||
if fill:
|
||||
regions, n = ndimage.label(~voxels)
|
||||
distance = ndimage.distance_transform_cdt(~voxels)
|
||||
representatives = [
|
||||
np.unravel_index((distance * (regions == i)).argmax(), distance.shape)
|
||||
for i in range(1, n + 1)
|
||||
]
|
||||
contains = mesh.contains(np.asarray(representatives) * pitch + local_origin)
|
||||
where = np.where(contains)[0] + 1
|
||||
internal = np.isin(regions.flatten(), where).reshape(regions.shape)
|
||||
voxels = np.logical_or(voxels, internal)
|
||||
|
||||
return base.VoxelGrid(voxels, tr.translation_matrix(local_origin))
|
||||
|
||||
|
||||
@log_time
|
||||
def voxelize_ray(
|
||||
mesh, pitch: Number, per_cell: Optional[ArrayLike] = None
|
||||
) -> base.VoxelGrid:
|
||||
"""
|
||||
Voxelize a mesh using ray queries.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
mesh
|
||||
Mesh to be voxelized
|
||||
pitch
|
||||
Length of voxel cube
|
||||
per_cell : (2,) int
|
||||
How many ray queries to make per cell
|
||||
|
||||
Returns
|
||||
-------------
|
||||
grid
|
||||
VoxelGrid instance representing the voxelized mesh.
|
||||
"""
|
||||
if per_cell is None:
|
||||
# how many rays per cell
|
||||
per_cell = np.array([2, 2], dtype=np.int64)
|
||||
else:
|
||||
per_cell = np.array(per_cell, dtype=np.int64).reshape(2)
|
||||
|
||||
# edge length of cube voxels
|
||||
pitch = float(pitch)
|
||||
|
||||
# create the ray origins in a grid
|
||||
bounds = mesh.bounds[:, :2].copy()
|
||||
# offset start so we get the requested number per cell
|
||||
bounds[0] += pitch / (1.0 + per_cell)
|
||||
# offset end so arange doesn't short us
|
||||
bounds[1] += pitch
|
||||
# on X we are doing multiple rays per voxel step
|
||||
step = pitch / per_cell
|
||||
# 2D grid
|
||||
ray_ori = util.grid_arange(bounds, step=step)
|
||||
# a Z position below the mesh
|
||||
z = np.ones(len(ray_ori)) * (mesh.bounds[0][2] - pitch)
|
||||
ray_ori = np.column_stack((ray_ori, z))
|
||||
# all rays are along positive Z
|
||||
ray_dir = np.ones_like(ray_ori) * [0, 0, 1]
|
||||
|
||||
# if you have pyembree this should be decently fast
|
||||
hits = mesh.ray.intersects_location(ray_ori, ray_dir)[0]
|
||||
|
||||
# just convert hit locations to integer positions
|
||||
voxels = np.round(hits / pitch).astype(np.int64)
|
||||
|
||||
# offset voxels by min, so matrix isn't huge
|
||||
origin_index = voxels.min(axis=0)
|
||||
voxels -= origin_index
|
||||
encoding = enc.SparseBinaryEncoding(voxels)
|
||||
origin_position = origin_index * pitch
|
||||
return base.VoxelGrid(
|
||||
encoding, tr.scale_and_translate(scale=pitch, translate=origin_position)
|
||||
)
|
||||
|
||||
|
||||
@log_time
|
||||
def voxelize_binvox(
|
||||
mesh,
|
||||
pitch: Optional[Number] = None,
|
||||
dimension: Optional[Integer] = None,
|
||||
bounds: Optional[ArrayLike] = None,
|
||||
**binvoxer_kwargs,
|
||||
) -> base.VoxelGrid:
|
||||
"""
|
||||
Voxelize via binvox tool.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
mesh : trimesh.Trimesh
|
||||
Mesh to voxelize
|
||||
pitch : float
|
||||
Side length of each voxel. Ignored if dimension is provided
|
||||
dimension: int
|
||||
Number of voxels along each dimension. If not provided, this is
|
||||
calculated based on pitch and bounds/mesh extents
|
||||
bounds: (2, 3) float
|
||||
min/max values of the returned `VoxelGrid` in each instance. Uses
|
||||
`mesh.bounds` if not provided.
|
||||
**binvoxer_kwargs:
|
||||
Passed to `trimesh.exchange.binvox.Binvoxer`.
|
||||
Should not contain `bounding_box` if bounds is not None.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
grid
|
||||
`VoxelGrid` instance
|
||||
|
||||
Raises
|
||||
--------------
|
||||
`ValueError` if `bounds is not None and 'bounding_box' in binvoxer_kwargs`.
|
||||
"""
|
||||
from trimesh.exchange import binvox
|
||||
|
||||
if dimension is None:
|
||||
# pitch must be provided
|
||||
if bounds is None:
|
||||
extents = mesh.extents
|
||||
else:
|
||||
mins, maxs = bounds
|
||||
extents = maxs - mins
|
||||
dimension = int(np.ceil(np.max(extents) / pitch))
|
||||
if bounds is not None:
|
||||
if "bounding_box" in binvoxer_kwargs:
|
||||
raise ValueError("Cannot provide both bounds and bounding_box")
|
||||
binvoxer_kwargs["bounding_box"] = np.asanyarray(bounds).flatten()
|
||||
|
||||
binvoxer = binvox.Binvoxer(dimension=dimension, **binvoxer_kwargs)
|
||||
return binvox.voxelize_mesh(mesh, binvoxer)
|
||||
|
||||
|
||||
voxelizers = util.FunctionRegistry(
|
||||
ray=voxelize_ray, subdivide=voxelize_subdivide, binvox=voxelize_binvox
|
||||
)
|
||||
|
||||
|
||||
def voxelize(
|
||||
mesh,
|
||||
pitch: Optional[Number],
|
||||
method: VoxelizationMethodsType = "subdivide",
|
||||
**kwargs,
|
||||
) -> Optional[base.VoxelGrid]:
|
||||
"""
|
||||
Voxelize the given mesh using the specified implementation.
|
||||
|
||||
See `voxelizers` for available implementations or to add your own, e.g. via
|
||||
`voxelizers['custom_key'] = custom_fn`.
|
||||
|
||||
`custom_fn` should have signature `(mesh, pitch, **kwargs) -> VoxelGrid`
|
||||
and should not modify encoding.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
mesh
|
||||
Geometry to voxelize
|
||||
pitch
|
||||
Side length of each voxel.
|
||||
method
|
||||
Which voxelization method to use.
|
||||
kwargs
|
||||
Passed through to the specified implementation.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
grid
|
||||
A VoxelGrid instance.
|
||||
"""
|
||||
return voxelizers(method, mesh=mesh, pitch=pitch, **kwargs)
|
||||
@@ -0,0 +1,979 @@
|
||||
"""OO interfaces to encodings for ND arrays which caching."""
|
||||
|
||||
import abc
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import caching
|
||||
from ..util import ABC
|
||||
from . import runlength
|
||||
|
||||
try:
|
||||
from scipy import sparse as sp
|
||||
except BaseException as E:
|
||||
from ..exceptions import ExceptionWrapper
|
||||
|
||||
sp = ExceptionWrapper(E)
|
||||
|
||||
|
||||
def _empty_stripped(shape):
|
||||
num_dims = len(shape)
|
||||
encoding = DenseEncoding(np.zeros(shape=(0,) * num_dims, dtype=bool))
|
||||
padding = np.zeros(shape=(num_dims, 2), dtype=int)
|
||||
padding[:, 1] = shape
|
||||
return encoding, padding
|
||||
|
||||
|
||||
class Encoding(ABC):
|
||||
"""
|
||||
Base class for objects that implement a specific subset of of ndarray ops.
|
||||
|
||||
This presents a unified interface for various different ways of encoding
|
||||
conceptually dense arrays and to interoperate between them.
|
||||
|
||||
Example implementations are ND sparse arrays, run length encoded arrays
|
||||
and dense encodings (wrappers around np.ndarrays).
|
||||
"""
|
||||
|
||||
def __init__(self, data):
|
||||
self._data = data
|
||||
self._cache = caching.Cache(id_function=self._data.__hash__)
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def dtype(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def shape(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def sum(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def size(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def sparse_indices(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def sparse_values(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def dense(self):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def gather_nd(self, indices):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def mask(self, mask):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def get_value(self, index):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def copy(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def is_empty(self):
|
||||
return self.sparse_indices[self.sparse_values != 0].size == 0
|
||||
|
||||
@caching.cache_decorator
|
||||
def stripped(self):
|
||||
"""
|
||||
Get encoding with all zeros stripped from the start and end
|
||||
of each axis.
|
||||
|
||||
Returns
|
||||
------------
|
||||
encoding: ?
|
||||
padding : (n, 2) int
|
||||
Padding at the start and end that was stripped
|
||||
"""
|
||||
if self.is_empty:
|
||||
return _empty_stripped(self.shape)
|
||||
dense = self.dense
|
||||
shape = dense.shape
|
||||
ndims = len(shape)
|
||||
padding = []
|
||||
slices = []
|
||||
for dim, size in enumerate(shape):
|
||||
axis = tuple(range(dim)) + tuple(range(dim + 1, ndims))
|
||||
filled = np.any(dense, axis=axis)
|
||||
(indices,) = np.nonzero(filled)
|
||||
lower = indices.min()
|
||||
upper = indices.max() + 1
|
||||
padding.append([lower, size - upper])
|
||||
slices.append(slice(lower, upper))
|
||||
return DenseEncoding(dense[tuple(slices)]), np.array(padding, int)
|
||||
|
||||
def _flip(self, axes):
|
||||
return FlippedEncoding(self, axes)
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
Get the hash of the current transformation matrix.
|
||||
|
||||
Returns
|
||||
------------
|
||||
hash : str
|
||||
Hash of transformation matrix
|
||||
"""
|
||||
return self._data.__hash__()
|
||||
|
||||
@property
|
||||
def ndims(self):
|
||||
return len(self.shape)
|
||||
|
||||
def reshape(self, shape):
|
||||
return self.flat if len(shape) == 1 else ShapedEncoding(self, shape)
|
||||
|
||||
@property
|
||||
def flat(self):
|
||||
return FlattenedEncoding(self)
|
||||
|
||||
def flip(self, axis=0):
|
||||
return _flipped(self, axis)
|
||||
|
||||
@property
|
||||
def sparse_components(self):
|
||||
return self.sparse_indices, self.sparse_values
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
return self._data
|
||||
|
||||
def run_length_data(self, dtype=np.int64):
|
||||
if self.ndims != 1:
|
||||
raise ValueError("`run_length_data` only valid for flat encodings")
|
||||
return runlength.dense_to_rle(self.dense, dtype=dtype)
|
||||
|
||||
def binary_run_length_data(self, dtype=np.int64):
|
||||
if self.ndims != 1:
|
||||
raise ValueError("`run_length_data` only valid for flat encodings")
|
||||
return runlength.dense_to_brle(self.dense, dtype=dtype)
|
||||
|
||||
def transpose(self, perm):
|
||||
return _transposed(self, perm)
|
||||
|
||||
def _transpose(self, perm):
|
||||
return TransposedEncoding(self, perm)
|
||||
|
||||
@property
|
||||
def mutable(self):
|
||||
return self._data.mutable
|
||||
|
||||
@mutable.setter
|
||||
def mutable(self, value):
|
||||
self._data.mutable = value
|
||||
|
||||
|
||||
class DenseEncoding(Encoding):
|
||||
"""Simple `Encoding` implementation based on a numpy ndarray."""
|
||||
|
||||
def __init__(self, data):
|
||||
if not isinstance(data, caching.TrackedArray):
|
||||
if not isinstance(data, np.ndarray):
|
||||
raise ValueError("DenseEncoding data must be a numpy array")
|
||||
data = caching.tracked_array(data)
|
||||
super().__init__(data=data)
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._data.dtype
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._data.shape
|
||||
|
||||
@caching.cache_decorator
|
||||
def sum(self):
|
||||
return self._data.sum()
|
||||
|
||||
@caching.cache_decorator
|
||||
def is_empty(self):
|
||||
return not np.any(self._data)
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return self._data.size
|
||||
|
||||
@property
|
||||
def sparse_components(self):
|
||||
indices = self.sparse_indices
|
||||
values = self.gather(indices)
|
||||
return indices, values
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_indices(self):
|
||||
return np.column_stack(np.where(self._data))
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_values(self):
|
||||
return self.sparse_components[1]
|
||||
|
||||
def _flip(self, axes):
|
||||
dense = self.dense
|
||||
for a in axes:
|
||||
dense = np.flip(dense, a)
|
||||
return DenseEncoding(dense)
|
||||
|
||||
@property
|
||||
def dense(self):
|
||||
return self._data
|
||||
|
||||
def gather(self, indices):
|
||||
return self._data[indices]
|
||||
|
||||
def gather_nd(self, indices):
|
||||
return self._data[tuple(indices.T)]
|
||||
|
||||
def mask(self, mask):
|
||||
return self._data[mask if isinstance(mask, np.ndarray) else mask.dense]
|
||||
|
||||
def get_value(self, index):
|
||||
return self._data[tuple(index)]
|
||||
|
||||
def reshape(self, shape):
|
||||
return DenseEncoding(self._data.reshape(shape))
|
||||
|
||||
def _transpose(self, perm):
|
||||
return DenseEncoding(self._data.transpose(perm))
|
||||
|
||||
@property
|
||||
def flat(self):
|
||||
return DenseEncoding(self._data.reshape((-1,)))
|
||||
|
||||
def copy(self):
|
||||
return DenseEncoding(self._data.copy())
|
||||
|
||||
|
||||
class SparseEncoding(Encoding):
|
||||
"""
|
||||
`Encoding` implementation based on an ND sparse implementation.
|
||||
|
||||
Since the scipy.sparse implementations are for 2D arrays only, this
|
||||
implementation uses a single-column CSC matrix with index
|
||||
raveling/unraveling.
|
||||
"""
|
||||
|
||||
def __init__(self, indices, values, shape=None):
|
||||
"""
|
||||
Parameters
|
||||
------------
|
||||
indices: (m, n)-sized int array of indices
|
||||
values: (m, n)-sized dtype array of values at the specified indices
|
||||
shape: (n,) iterable of integers. If None, the maximum value of indices
|
||||
+ 1 is used.
|
||||
"""
|
||||
data = caching.DataStore()
|
||||
super().__init__(data)
|
||||
data["indices"] = indices
|
||||
data["values"] = values
|
||||
indices = data["indices"]
|
||||
if len(indices.shape) != 2:
|
||||
raise ValueError(f"indices must be 2D, got shaped {indices.shape!s}")
|
||||
if data["values"].shape != (indices.shape[0],):
|
||||
raise ValueError(
|
||||
"values and indices shapes inconsistent: {} and {}".format(
|
||||
data["values"], data["indices"]
|
||||
)
|
||||
)
|
||||
if shape is None:
|
||||
self._shape = tuple(data["indices"].max(axis=0) + 1)
|
||||
else:
|
||||
self._shape = tuple(shape)
|
||||
if not np.all(indices < self._shape):
|
||||
raise ValueError("all indices must be less than shape")
|
||||
if not np.all(indices >= 0):
|
||||
raise ValueError("all indices must be non-negative")
|
||||
|
||||
@staticmethod
|
||||
def from_dense(dense_data):
|
||||
sparse_indices = np.where(dense_data)
|
||||
values = dense_data[sparse_indices]
|
||||
return SparseEncoding(
|
||||
np.stack(sparse_indices, axis=-1), values, shape=dense_data.shape
|
||||
)
|
||||
|
||||
def copy(self):
|
||||
return SparseEncoding(
|
||||
indices=self.sparse_indices.copy(),
|
||||
values=self.sparse_values.copy(),
|
||||
shape=self.shape,
|
||||
)
|
||||
|
||||
@property
|
||||
def sparse_indices(self):
|
||||
return self._data["indices"]
|
||||
|
||||
@property
|
||||
def sparse_values(self):
|
||||
return self._data["values"]
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self.sparse_values.dtype
|
||||
|
||||
@caching.cache_decorator
|
||||
def sum(self):
|
||||
return self.sparse_values.sum()
|
||||
|
||||
@property
|
||||
def ndims(self):
|
||||
return self.sparse_indices.shape[-1]
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._shape
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return np.prod(self.shape)
|
||||
|
||||
@property
|
||||
def sparse_components(self):
|
||||
return self.sparse_indices, self.sparse_values
|
||||
|
||||
@caching.cache_decorator
|
||||
def dense(self):
|
||||
sparse = self._csc
|
||||
# sparse.todense gives an `np.matrix` which cannot be reshaped
|
||||
dense = np.zeros(shape=sparse.shape, dtype=sparse.dtype)
|
||||
sparse.todense(out=dense)
|
||||
return np.reshape(dense, self.shape)
|
||||
|
||||
@caching.cache_decorator
|
||||
def _csc(self):
|
||||
values = self.sparse_values
|
||||
indices = self._flat_indices(self.sparse_indices)
|
||||
indptr = [0, len(indices)]
|
||||
return sp.csc_matrix((values, indices, indptr), shape=(self.size, 1))
|
||||
|
||||
def _flat_indices(self, indices):
|
||||
assert indices.shape[1] == 3 and len(indices.shape) == 2
|
||||
return np.ravel_multi_index(indices.T, self.shape)
|
||||
|
||||
def _shaped_indices(self, flat_indices):
|
||||
return np.column_stack(np.unravel_index(flat_indices, self.shape))
|
||||
|
||||
def gather_nd(self, indices):
|
||||
mat = self._csc[self._flat_indices(indices)].todense()
|
||||
# mat is a np matrix, which stays rank 2 after squeeze
|
||||
# np.asarray changes this to a standard rank 2 array.
|
||||
return np.asarray(mat).squeeze(axis=-1)
|
||||
|
||||
def mask(self, mask):
|
||||
i, _ = np.where(self._csc[mask.reshape((-1,))])
|
||||
return self._shaped_indices(i)
|
||||
|
||||
def get_value(self, index):
|
||||
return self._gather_nd(np.expand_dims(index, axis=0))[0]
|
||||
|
||||
@caching.cache_decorator
|
||||
def stripped(self):
|
||||
"""
|
||||
Get encoding with all zeros stripped from the start/end of each axis.
|
||||
|
||||
Returns:
|
||||
encoding: SparseEncoding with same values but indices shifted down
|
||||
by padding[:, 0]
|
||||
padding: (n, 2) array of ints denoting padding at the start/end
|
||||
that was stripped
|
||||
"""
|
||||
if self.is_empty:
|
||||
return _empty_stripped(self.shape)
|
||||
indices = self.sparse_indices
|
||||
pad_left = np.min(indices, axis=0)
|
||||
pad_right = np.max(indices, axis=0)
|
||||
pad_right *= -1
|
||||
pad_right += self.shape
|
||||
padding = np.column_stack((pad_left, pad_right))
|
||||
return SparseEncoding(indices - pad_left, self.sparse_values), padding
|
||||
|
||||
|
||||
def SparseBinaryEncoding(indices, shape=None):
|
||||
"""
|
||||
Convenient factory constructor for SparseEncodings with values all ones.
|
||||
|
||||
Parameters
|
||||
------------
|
||||
indices: (m, n) sparse indices into conceptual rank-n array
|
||||
shape: length n iterable or None. If None, maximum of indices along first
|
||||
axis + 1 is used
|
||||
|
||||
Returns
|
||||
------------
|
||||
rank n bool `SparseEncoding` with True values at each index.
|
||||
"""
|
||||
return SparseEncoding(indices, np.ones(shape=(indices.shape[0],), dtype=bool), shape)
|
||||
|
||||
|
||||
class RunLengthEncoding(Encoding):
|
||||
"""1D run length encoding.
|
||||
|
||||
See `trimesh.voxel.runlength` documentation for implementation details.
|
||||
"""
|
||||
|
||||
def __init__(self, data, dtype=None):
|
||||
"""
|
||||
Parameters
|
||||
------------
|
||||
data: run length encoded data.
|
||||
dtype: dtype of encoded data. Each second value of data is cast will be
|
||||
cast to this dtype if provided.
|
||||
"""
|
||||
super().__init__(data=caching.tracked_array(data))
|
||||
if dtype is None:
|
||||
dtype = self._data.dtype
|
||||
if len(self._data.shape) != 1:
|
||||
raise ValueError("data must be 1D numpy array")
|
||||
self._dtype = dtype
|
||||
|
||||
@caching.cache_decorator
|
||||
def is_empty(self):
|
||||
return not np.any(np.logical_and(self._data[::2], self._data[1::2]))
|
||||
|
||||
@property
|
||||
def ndims(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return (self.size,)
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._dtype
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
Get the hash of the current transformation matrix.
|
||||
|
||||
Returns
|
||||
------------
|
||||
hash : str
|
||||
Hash of transformation matrix
|
||||
"""
|
||||
return self._data.__hash__()
|
||||
|
||||
@staticmethod
|
||||
def from_dense(dense_data, dtype=np.int64, encoding_dtype=np.int64):
|
||||
return RunLengthEncoding(
|
||||
runlength.dense_to_rle(dense_data, dtype=encoding_dtype), dtype=dtype
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_rle(rle_data, dtype=None):
|
||||
if dtype != rle_data.dtype:
|
||||
rle_data = runlength.rle_to_rle(rle_data, dtype=dtype)
|
||||
return RunLengthEncoding(rle_data)
|
||||
|
||||
@staticmethod
|
||||
def from_brle(brle_data, dtype=None):
|
||||
return RunLengthEncoding(runlength.brle_to_rle(brle_data, dtype=dtype))
|
||||
|
||||
@caching.cache_decorator
|
||||
def stripped(self):
|
||||
if self.is_empty:
|
||||
return _empty_stripped(self.shape)
|
||||
data, padding = runlength.rle_strip(self._data)
|
||||
if padding == (0, 0):
|
||||
encoding = self
|
||||
else:
|
||||
encoding = RunLengthEncoding(data, dtype=self._dtype)
|
||||
padding = np.expand_dims(padding, axis=0)
|
||||
return encoding, padding
|
||||
|
||||
@caching.cache_decorator
|
||||
def sum(self):
|
||||
return (self._data[::2] * self._data[1::2]).sum()
|
||||
|
||||
@caching.cache_decorator
|
||||
def size(self):
|
||||
return runlength.rle_length(self._data)
|
||||
|
||||
def _flip(self, axes):
|
||||
if axes != (0,):
|
||||
raise ValueError(f"encoding is 1D - cannot flip on axis {axes!s}")
|
||||
return RunLengthEncoding(runlength.rle_reverse(self._data))
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_components(self):
|
||||
return runlength.rle_to_sparse(self._data)
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_indices(self):
|
||||
return self.sparse_components[0]
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_values(self):
|
||||
return self.sparse_components[1]
|
||||
|
||||
@caching.cache_decorator
|
||||
def dense(self):
|
||||
return runlength.rle_to_dense(self._data, dtype=self._dtype)
|
||||
|
||||
def gather(self, indices):
|
||||
return runlength.rle_gather_1d(self._data, indices, dtype=self._dtype)
|
||||
|
||||
def gather_nd(self, indices):
|
||||
indices = np.squeeze(indices, axis=-1)
|
||||
return self.gather(indices)
|
||||
|
||||
def sorted_gather(self, ordered_indices):
|
||||
return np.array(
|
||||
tuple(runlength.sorted_rle_gather_1d(self._data, ordered_indices)),
|
||||
dtype=self._dtype,
|
||||
)
|
||||
|
||||
def mask(self, mask):
|
||||
return np.array(tuple(runlength.rle_mask(self._data, mask)), dtype=self._dtype)
|
||||
|
||||
def get_value(self, index):
|
||||
for value in self.sorted_gather((index,)):
|
||||
return np.asanyarray(value, dtype=self._dtype)
|
||||
|
||||
def copy(self):
|
||||
return RunLengthEncoding(self._data.copy(), dtype=self.dtype)
|
||||
|
||||
def run_length_data(self, dtype=np.int64):
|
||||
return runlength.rle_to_rle(self._data, dtype=dtype)
|
||||
|
||||
def binary_run_length_data(self, dtype=np.int64):
|
||||
return runlength.rle_to_brle(self._data, dtype=dtype)
|
||||
|
||||
|
||||
class BinaryRunLengthEncoding(RunLengthEncoding):
|
||||
"""1D binary run length encoding.
|
||||
|
||||
See `trimesh.voxel.runlength` documentation for implementation details.
|
||||
"""
|
||||
|
||||
def __init__(self, data):
|
||||
"""
|
||||
Parameters
|
||||
------------
|
||||
data: binary run length encoded data.
|
||||
"""
|
||||
super().__init__(data=data, dtype=bool)
|
||||
|
||||
@caching.cache_decorator
|
||||
def is_empty(self):
|
||||
return not np.any(self._data[1::2])
|
||||
|
||||
@staticmethod
|
||||
def from_dense(dense_data, encoding_dtype=np.int64):
|
||||
return BinaryRunLengthEncoding(
|
||||
runlength.dense_to_brle(dense_data, dtype=encoding_dtype)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def from_rle(rle_data, dtype=None):
|
||||
return BinaryRunLengthEncoding(runlength.rle_to_brle(rle_data, dtype=dtype))
|
||||
|
||||
@staticmethod
|
||||
def from_brle(brle_data, dtype=None):
|
||||
if dtype != brle_data.dtype:
|
||||
brle_data = runlength.brle_to_brle(brle_data, dtype=dtype)
|
||||
return BinaryRunLengthEncoding(brle_data)
|
||||
|
||||
@caching.cache_decorator
|
||||
def stripped(self):
|
||||
if self.is_empty:
|
||||
return _empty_stripped(self.shape)
|
||||
data, padding = runlength.rle_strip(self._data)
|
||||
if padding == (0, 0):
|
||||
encoding = self
|
||||
else:
|
||||
encoding = BinaryRunLengthEncoding(data)
|
||||
padding = np.expand_dims(padding, axis=0)
|
||||
return encoding, padding
|
||||
|
||||
@caching.cache_decorator
|
||||
def sum(self):
|
||||
return self._data[1::2].sum()
|
||||
|
||||
@caching.cache_decorator
|
||||
def size(self):
|
||||
return runlength.brle_length(self._data)
|
||||
|
||||
def _flip(self, axes):
|
||||
if axes != (0,):
|
||||
raise ValueError(f"encoding is 1D - cannot flip on axis {axes!s}")
|
||||
return BinaryRunLengthEncoding(runlength.brle_reverse(self._data))
|
||||
|
||||
@property
|
||||
def sparse_components(self):
|
||||
return self.sparse_indices, self.sparse_values
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_values(self):
|
||||
return np.ones(shape=(self.sum,), dtype=bool)
|
||||
|
||||
@caching.cache_decorator
|
||||
def sparse_indices(self):
|
||||
return runlength.brle_to_sparse(self._data)
|
||||
|
||||
@caching.cache_decorator
|
||||
def dense(self):
|
||||
return runlength.brle_to_dense(self._data)
|
||||
|
||||
def gather(self, indices):
|
||||
return runlength.brle_gather_1d(self._data, indices)
|
||||
|
||||
def gather_nd(self, indices):
|
||||
indices = np.squeeze(indices)
|
||||
return self.gather(indices)
|
||||
|
||||
def sorted_gather(self, ordered_indices):
|
||||
gen = runlength.sorted_brle_gather_1d(self._data, ordered_indices)
|
||||
return np.array(tuple(gen), dtype=bool)
|
||||
|
||||
def mask(self, mask):
|
||||
gen = runlength.brle_mask(self._data, mask)
|
||||
return np.array(tuple(gen), dtype=bool)
|
||||
|
||||
def copy(self):
|
||||
return BinaryRunLengthEncoding(self._data.copy())
|
||||
|
||||
def run_length_data(self, dtype=np.int64):
|
||||
return runlength.brle_to_rle(self._data, dtype=dtype)
|
||||
|
||||
def binary_run_length_data(self, dtype=np.int64):
|
||||
return runlength.brle_to_brle(self._data, dtype=dtype)
|
||||
|
||||
|
||||
class LazyIndexMap(Encoding):
|
||||
"""
|
||||
Abstract class for implementing lazy index mapping operations.
|
||||
|
||||
Implementations include transpose, flatten/reshaping and flipping
|
||||
|
||||
Derived classes must implement:
|
||||
* _to_base_indices(indices)
|
||||
* _from_base_indices(base_indices)
|
||||
* shape
|
||||
* dense
|
||||
* mask(mask)
|
||||
"""
|
||||
|
||||
@abc.abstractmethod
|
||||
def _to_base_indices(self, indices):
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
def _from_base_indices(self, base_indices):
|
||||
pass
|
||||
|
||||
@property
|
||||
def is_empty(self):
|
||||
return self._data.is_empty
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._data.dtype
|
||||
|
||||
@property
|
||||
def sum(self):
|
||||
return self._data.sum
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return self._data.size
|
||||
|
||||
@property
|
||||
def sparse_indices(self):
|
||||
return self._from_base_indices(self._data.sparse_indices)
|
||||
|
||||
@property
|
||||
def sparse_values(self):
|
||||
return self._data.sparse_values
|
||||
|
||||
def gather_nd(self, indices):
|
||||
return self._data.gather_nd(self._to_base_indices(indices))
|
||||
|
||||
def get_value(self, index):
|
||||
return self._data[tuple(self._to_base_indices(index))]
|
||||
|
||||
|
||||
class FlattenedEncoding(LazyIndexMap):
|
||||
"""
|
||||
Lazily flattened encoding.
|
||||
|
||||
Dense equivalent is np.reshape(data, (-1,)) (np.flatten creates a copy).
|
||||
"""
|
||||
|
||||
def _to_base_indices(self, indices):
|
||||
return np.column_stack(np.unravel_index(indices, self._data.shape))
|
||||
|
||||
def _from_base_indices(self, base_indices):
|
||||
return np.expand_dims(
|
||||
np.ravel_multi_index(base_indices.T, self._data.shape), axis=-1
|
||||
)
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return (self.size,)
|
||||
|
||||
@property
|
||||
def dense(self):
|
||||
return self._data.dense.reshape((-1,))
|
||||
|
||||
def mask(self, mask):
|
||||
return self._data.mask(mask.reshape(self._data.shape))
|
||||
|
||||
@property
|
||||
def flat(self):
|
||||
return self
|
||||
|
||||
def copy(self):
|
||||
return FlattenedEncoding(self._data.copy())
|
||||
|
||||
|
||||
class ShapedEncoding(LazyIndexMap):
|
||||
"""
|
||||
Lazily reshaped encoding.
|
||||
|
||||
Numpy equivalent is `np.reshape`
|
||||
"""
|
||||
|
||||
def __init__(self, encoding, shape):
|
||||
if isinstance(encoding, Encoding):
|
||||
if encoding.ndims != 1:
|
||||
encoding = encoding.flat
|
||||
else:
|
||||
raise ValueError("encoding must be an Encoding")
|
||||
super().__init__(data=encoding)
|
||||
self._shape = tuple(shape)
|
||||
nn = self._shape.count(-1)
|
||||
size = np.prod(self._shape)
|
||||
if nn == 1:
|
||||
size = np.abs(size)
|
||||
if self._data.size % size != 0:
|
||||
raise ValueError(
|
||||
"cannot reshape encoding of size %d into shape %s",
|
||||
self._data.size,
|
||||
str(self._shape),
|
||||
)
|
||||
|
||||
rem = self._data.size // size
|
||||
self._shape = tuple(rem if s == -1 else s for s in self._shape)
|
||||
elif nn > 2:
|
||||
raise ValueError("shape cannot have more than one -1 value")
|
||||
elif np.prod(self._shape) != self._data.size:
|
||||
raise ValueError(
|
||||
"cannot reshape encoding of size %d into shape %s",
|
||||
self._data.size,
|
||||
str(self._shape),
|
||||
)
|
||||
|
||||
def _from_base_indices(self, base_indices):
|
||||
return np.column_stack(np.unravel_index(base_indices, self.shape))
|
||||
|
||||
def _to_base_indices(self, indices):
|
||||
return np.expand_dims(np.ravel_multi_index(indices.T, self.shape), axis=-1)
|
||||
|
||||
@property
|
||||
def flat(self):
|
||||
return self._data
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._shape
|
||||
|
||||
@property
|
||||
def dense(self):
|
||||
return self._data.dense.reshape(self.shape)
|
||||
|
||||
def mask(self, mask):
|
||||
return self._data.mask(mask.flat)
|
||||
|
||||
def copy(self):
|
||||
return ShapedEncoding(encoding=self._data.copy(), shape=self.shape)
|
||||
|
||||
|
||||
class TransposedEncoding(LazyIndexMap):
|
||||
"""
|
||||
Lazily transposed encoding
|
||||
|
||||
Dense equivalent is `np.transpose`
|
||||
"""
|
||||
|
||||
def __init__(self, base_encoding, perm):
|
||||
if not isinstance(base_encoding, Encoding):
|
||||
raise ValueError(f"base_encoding must be an Encoding, got {base_encoding!s}")
|
||||
if len(base_encoding.shape) != len(perm):
|
||||
raise ValueError(
|
||||
"base_encoding has %d ndims - cannot transpose with perm %s",
|
||||
base_encoding.ndims,
|
||||
str(perm),
|
||||
)
|
||||
|
||||
super().__init__(base_encoding)
|
||||
perm = np.array(perm, dtype=np.int64)
|
||||
if not all(i in perm for i in range(base_encoding.ndims)):
|
||||
raise ValueError(f"perm {perm!s} is not a valid permutation")
|
||||
inv_perm = np.zeros_like(perm)
|
||||
inv_perm[perm] = np.arange(base_encoding.ndims)
|
||||
self._perm = perm
|
||||
self._inv_perm = inv_perm
|
||||
|
||||
def transpose(self, perm):
|
||||
return _transposed(self._data, [self._perm[p] for p in perm])
|
||||
|
||||
def _transpose(self, perm):
|
||||
raise RuntimeError("Should not be here")
|
||||
|
||||
@property
|
||||
def perm(self):
|
||||
return self._perm
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
shape = self._data.shape
|
||||
return tuple(shape[p] for p in self._perm)
|
||||
|
||||
def _to_base_indices(self, indices):
|
||||
return np.take(indices, self._perm, axis=-1)
|
||||
|
||||
def _from_base_indices(self, base_indices):
|
||||
try:
|
||||
return np.take(base_indices, self._inv_perm, axis=-1)
|
||||
except TypeError:
|
||||
# windows sometimes tries to use wrong dtypes
|
||||
return np.take(
|
||||
base_indices.astype(np.int64), self._inv_perm.astype(np.int64), axis=-1
|
||||
)
|
||||
|
||||
@property
|
||||
def dense(self):
|
||||
return self._data.dense.transpose(self._perm)
|
||||
|
||||
def gather(self, indices):
|
||||
return self._data.gather(self._base_indices(indices))
|
||||
|
||||
def mask(self, mask):
|
||||
return self._data.mask(mask.transpose(self._inv_perm)).transpose(self._perm)
|
||||
|
||||
def get_value(self, index):
|
||||
return self._data[tuple(self._base_indices(index))]
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
return self._data
|
||||
|
||||
def copy(self):
|
||||
return TransposedEncoding(base_encoding=self._data.copy(), perm=self._perm)
|
||||
|
||||
|
||||
class FlippedEncoding(LazyIndexMap):
|
||||
"""
|
||||
Encoding with entries flipped along one or more axes.
|
||||
|
||||
Dense equivalent is `np.flip`
|
||||
"""
|
||||
|
||||
def __init__(self, encoding, axes):
|
||||
ndims = encoding.ndims
|
||||
if isinstance(axes, np.ndarray) and axes.size == 1:
|
||||
axes = (axes.item(),)
|
||||
elif isinstance(axes, int):
|
||||
axes = (axes,)
|
||||
axes = tuple(a + ndims if a < 0 else a for a in axes)
|
||||
self._axes = tuple(sorted(axes))
|
||||
if len(set(self._axes)) != len(self._axes):
|
||||
raise ValueError(f"Axes cannot contain duplicates, got {self._axes!s}")
|
||||
super().__init__(encoding)
|
||||
if not all(0 <= a < self._data.ndims for a in axes):
|
||||
raise ValueError(
|
||||
"Invalid axes %s for %d-d encoding", str(axes), self._data.ndims
|
||||
)
|
||||
|
||||
def _to_base_indices(self, indices):
|
||||
indices = indices.copy()
|
||||
shape = self.shape
|
||||
for a in self._axes:
|
||||
indices[:, a] *= -1
|
||||
indices[:, a] += shape
|
||||
return indices
|
||||
|
||||
def _from_base_indices(self, base_indices):
|
||||
return self._to_base_indices(base_indices)
|
||||
|
||||
@property
|
||||
def shape(self):
|
||||
return self._data.shape
|
||||
|
||||
@property
|
||||
def dense(self):
|
||||
dense = self._data.dense
|
||||
for a in self._axes:
|
||||
dense = np.flip(dense, a)
|
||||
return dense
|
||||
|
||||
def mask(self, mask):
|
||||
if not isinstance(mask, Encoding):
|
||||
mask = DenseEncoding(mask)
|
||||
mask = mask.flip(self._axes)
|
||||
return self._data.mask(mask).flip(self._axes)
|
||||
|
||||
def copy(self):
|
||||
return FlippedEncoding(self._data.copy(), self._axes)
|
||||
|
||||
def flip(self, axis=0):
|
||||
if isinstance(axis, np.ndarray):
|
||||
if axis.size == 1:
|
||||
axis = (axis.item(),)
|
||||
else:
|
||||
axis = tuple(axis)
|
||||
elif isinstance(axis, int):
|
||||
axes = (axis,)
|
||||
else:
|
||||
axes = tuple(axis)
|
||||
return _flipped(self, self._axes + axes)
|
||||
|
||||
def _flip(self, axes):
|
||||
raise RuntimeError("Should not be here")
|
||||
|
||||
|
||||
def _flipped(encoding, axes):
|
||||
if not hasattr(axes, "__iter__"):
|
||||
axes = (axes,)
|
||||
unique_ax = set()
|
||||
ndims = encoding.ndims
|
||||
axes = tuple(a + ndims if a < 0 else a for a in axes)
|
||||
for a in axes:
|
||||
if a in unique_ax:
|
||||
unique_ax.remove(a)
|
||||
else:
|
||||
unique_ax.add(a)
|
||||
if len(unique_ax) == 0:
|
||||
return encoding
|
||||
else:
|
||||
return encoding._flip(tuple(sorted(unique_ax)))
|
||||
|
||||
|
||||
def _transposed(encoding, perm):
|
||||
ndims = encoding.ndims
|
||||
perm = tuple(p + ndims if p < 0 else p for p in perm)
|
||||
if np.all(np.arange(ndims) == perm):
|
||||
return encoding
|
||||
else:
|
||||
return encoding._transpose(perm)
|
||||
@@ -0,0 +1,185 @@
|
||||
"""Basic morphology operations that create new encodings."""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from .. import util
|
||||
from ..constants import log_time
|
||||
from . import encoding as enc
|
||||
from . import ops
|
||||
|
||||
try:
|
||||
from scipy import ndimage
|
||||
except BaseException as E:
|
||||
# scipy is a soft dependency
|
||||
from ..exceptions import ExceptionWrapper
|
||||
|
||||
ndimage = ExceptionWrapper(E)
|
||||
|
||||
|
||||
def _dense(encoding, rank=None):
|
||||
if isinstance(encoding, np.ndarray):
|
||||
dense = encoding
|
||||
elif isinstance(encoding, enc.Encoding):
|
||||
dense = encoding.dense
|
||||
else:
|
||||
raise ValueError(f"encoding must be np.ndarray or Encoding, got {encoding!s}")
|
||||
if rank:
|
||||
_assert_rank(dense, rank)
|
||||
return dense
|
||||
|
||||
|
||||
def _sparse_indices(encoding, rank=None):
|
||||
if isinstance(encoding, np.ndarray):
|
||||
sparse_indices = encoding
|
||||
elif isinstance(encoding, enc.Encoding):
|
||||
sparse_indices = encoding.sparse_indices
|
||||
else:
|
||||
raise ValueError(f"encoding must be np.ndarray or Encoding, got {encoding!s}")
|
||||
|
||||
_assert_sparse_rank(sparse_indices, 3)
|
||||
return sparse_indices
|
||||
|
||||
|
||||
def _assert_rank(value, rank):
|
||||
if len(value.shape) != rank:
|
||||
raise ValueError("Expected rank %d, got shape %s", rank, str(value.shape))
|
||||
|
||||
|
||||
def _assert_sparse_rank(value, rank=None):
|
||||
if len(value.shape) != 2:
|
||||
raise ValueError(f"sparse_indices must be rank 2, got shape {value.shape!s}")
|
||||
if rank is not None:
|
||||
if value.shape[-1] != rank:
|
||||
raise ValueError(
|
||||
"sparse_indices.shape[1] must be %d, got %d", rank, value.shape[-1]
|
||||
)
|
||||
|
||||
|
||||
@log_time
|
||||
def fill_base(encoding):
|
||||
"""
|
||||
Given a sparse surface voxelization, fill in between columns.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
encoding: Encoding object or sparse array with shape (?, 3)
|
||||
|
||||
Returns
|
||||
--------------
|
||||
A new filled encoding object.
|
||||
"""
|
||||
return enc.SparseBinaryEncoding(ops.fill_base(_sparse_indices(encoding, rank=3)))
|
||||
|
||||
|
||||
@log_time
|
||||
def fill_orthographic(encoding):
|
||||
"""
|
||||
Fill the given encoding by orthographic projection method.
|
||||
|
||||
Any voxel in the dense representation with no free ray along the x, y, z
|
||||
axes in each direction is assigned filled. This is likely faster than fill
|
||||
holes, and is more stable with regards to small holes.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
encoding: Encoding object or dense rank-3 array.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
A new filled encoding object.
|
||||
"""
|
||||
return enc.DenseEncoding(ops.fill_orthographic(_dense(encoding, rank=3)))
|
||||
|
||||
|
||||
@log_time
|
||||
def fill_holes(encoding, **kwargs):
|
||||
"""
|
||||
Encoding wrapper around scipy.ndimage.morphology.binary_fill_holes.
|
||||
|
||||
https://docs.scipy.org/doc/scipy-0.15.1/reference/generated/scipy.ndimage.morphology.binary_fill_holes.html#scipy.ndimage.morphology.binary_fill_holes
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
encoding: Encoding object or dense rank-3 array.
|
||||
**kwargs: see scipy.ndimage.morphology.binary_fill_holes.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
A new filled in encoding object.
|
||||
"""
|
||||
return enc.DenseEncoding(
|
||||
ndimage.binary_fill_holes(_dense(encoding, rank=3), **kwargs)
|
||||
)
|
||||
|
||||
|
||||
fillers = util.FunctionRegistry(
|
||||
base=fill_base,
|
||||
orthographic=fill_orthographic,
|
||||
holes=fill_holes,
|
||||
)
|
||||
|
||||
|
||||
def fill(encoding, method="base", **kwargs):
|
||||
"""
|
||||
Fill the given encoding using the specified implementation.
|
||||
|
||||
See `fillers` for available implementations or to add your own, e.g. via
|
||||
`fillers['custom_key'] = custom_fn`.
|
||||
|
||||
`custom_fn` should have signature `(encoding, **kwargs) -> filled_encoding`
|
||||
and should not modify encoding.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
encoding: Encoding object (left unchanged).
|
||||
method: method present in `fillers`.
|
||||
**kwargs: additional kwargs passed to the specified implementation.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
A new filled Encoding object.
|
||||
"""
|
||||
return fillers(method, encoding=encoding, **kwargs)
|
||||
|
||||
|
||||
def binary_dilation(encoding, **kwargs):
|
||||
"""
|
||||
Encoding wrapper around scipy.ndimage.morphology.binary_dilation.
|
||||
|
||||
https://docs.scipy.org/doc/scipy-0.15.1/reference/generated/scipy.ndimage.morphology.binary_dilation.html#scipy.ndimage.morphology.binary_dilation
|
||||
"""
|
||||
return enc.DenseEncoding(ndimage.binary_dilation(_dense(encoding, rank=3), **kwargs))
|
||||
|
||||
|
||||
def binary_closing(encoding, **kwargs):
|
||||
"""
|
||||
Encoding wrapper around scipy.ndimage.morphology.binary_closing.
|
||||
|
||||
https://docs.scipy.org/doc/scipy-0.15.1/reference/generated/scipy.ndimage.morphology.binary_closing.html#scipy.ndimage.morphology.binary_closing
|
||||
"""
|
||||
return enc.DenseEncoding(ndimage.binary_closing(_dense(encoding, rank=3), **kwargs))
|
||||
|
||||
|
||||
def surface(encoding, structure=None):
|
||||
"""
|
||||
Get elements on the surface of encoding.
|
||||
|
||||
A surface element is any one in encoding that is adjacent to an empty
|
||||
voxel.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
encoding: Encoding or dense rank-3 array
|
||||
structure: adjacency structure. If None, square connectivity is used.
|
||||
|
||||
Returns
|
||||
--------------
|
||||
new surface Encoding.
|
||||
"""
|
||||
dense = _dense(encoding, rank=3)
|
||||
# padding/unpadding resolves issues with occupied voxels on the boundary
|
||||
dense = np.pad(dense, np.ones((3, 2), dtype=int), mode="constant")
|
||||
empty = np.logical_not(dense)
|
||||
dilated = ndimage.binary_dilation(empty, structure=structure)
|
||||
surface = np.logical_and(dense, dilated)[1:-1, 1:-1, 1:-1]
|
||||
return enc.DenseEncoding(surface)
|
||||
@@ -0,0 +1,454 @@
|
||||
import numpy as np
|
||||
|
||||
from .. import util
|
||||
from ..constants import log
|
||||
from ..typed import ArrayLike, Number, Optional, Union
|
||||
|
||||
|
||||
def fill_orthographic(dense):
|
||||
shape = dense.shape
|
||||
indices = np.stack(
|
||||
np.meshgrid(*(np.arange(s) for s in shape), indexing="ij"), axis=-1
|
||||
)
|
||||
empty = np.logical_not(dense)
|
||||
|
||||
def fill_axis(axis):
|
||||
base_local_indices = indices[..., axis]
|
||||
local_indices = base_local_indices.copy()
|
||||
local_indices[empty] = shape[axis]
|
||||
mins = np.min(local_indices, axis=axis, keepdims=True)
|
||||
local_indices = base_local_indices.copy()
|
||||
local_indices[empty] = -1
|
||||
maxs = np.max(local_indices, axis=axis, keepdims=True)
|
||||
|
||||
return np.logical_and(
|
||||
base_local_indices >= mins,
|
||||
base_local_indices <= maxs,
|
||||
)
|
||||
|
||||
filled = fill_axis(axis=0)
|
||||
for axis in range(1, len(shape)):
|
||||
filled = np.logical_and(filled, fill_axis(axis))
|
||||
return filled
|
||||
|
||||
|
||||
def fill_base(sparse_indices):
|
||||
"""
|
||||
Given a sparse surface voxelization, fill in between columns.
|
||||
|
||||
Parameters
|
||||
--------------
|
||||
sparse_indices: (n, 3) int, location of filled cells
|
||||
|
||||
Returns
|
||||
--------------
|
||||
filled: (m, 3) int, location of filled cells
|
||||
"""
|
||||
# validate inputs
|
||||
sparse_indices = np.asanyarray(sparse_indices, dtype=np.int64)
|
||||
if not util.is_shape(sparse_indices, (-1, 3)):
|
||||
raise ValueError("incorrect shape")
|
||||
|
||||
# create grid and mark inner voxels
|
||||
max_value = sparse_indices.max() + 3
|
||||
|
||||
grid = np.zeros((max_value, max_value, max_value), bool)
|
||||
voxels_sparse = np.add(sparse_indices, 1)
|
||||
|
||||
grid[tuple(voxels_sparse.T)] = 1
|
||||
|
||||
for i in range(max_value):
|
||||
check_dir2 = False
|
||||
for j in range(0, max_value - 1):
|
||||
idx = []
|
||||
# find transitions first
|
||||
# transition positions are from 0 to 1 and from 1 to 0
|
||||
eq = np.equal(grid[i, j, :-1], grid[i, j, 1:])
|
||||
idx = np.where(np.logical_not(eq))[0] + 1
|
||||
c = len(idx)
|
||||
check_dir2 = (c % 4) > 0 and c > 4
|
||||
if c < 4:
|
||||
continue
|
||||
for s in range(0, c - c % 4, 4):
|
||||
grid[i, j, idx[s] : idx[s + 3]] = 1
|
||||
if not check_dir2:
|
||||
continue
|
||||
|
||||
# check another direction for robustness
|
||||
for k in range(0, max_value - 1):
|
||||
idx = []
|
||||
# find transitions first
|
||||
eq = np.equal(grid[i, :-1, k], grid[i, 1:, k])
|
||||
idx = np.where(np.logical_not(eq))[0] + 1
|
||||
c = len(idx)
|
||||
if c < 4:
|
||||
continue
|
||||
for s in range(0, c - c % 4, 4):
|
||||
grid[i, idx[s] : idx[s + 3], k] = 1
|
||||
|
||||
# generate new voxels
|
||||
filled = np.column_stack(np.where(grid))
|
||||
filled -= 1
|
||||
|
||||
return filled
|
||||
|
||||
|
||||
fill_voxelization = fill_base
|
||||
|
||||
|
||||
def matrix_to_marching_cubes(
|
||||
matrix: ArrayLike,
|
||||
pitch: Union[Number, ArrayLike] = 1.0,
|
||||
threshold: Optional[Number] = None,
|
||||
):
|
||||
"""
|
||||
Convert an (n, m, p) matrix into a mesh, using marching_cubes.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
matrix : (n, m, p) bool
|
||||
Occupancy array
|
||||
pitch : float or length-3 tuple of floats, optional
|
||||
Voxel spacing in each dimension
|
||||
threshold : float or None, optional
|
||||
If specified, converts the input into a boolean
|
||||
matrix by considering values above `threshold` as True
|
||||
|
||||
|
||||
Returns
|
||||
----------
|
||||
mesh : trimesh.Trimesh
|
||||
Mesh generated by meshing voxels using
|
||||
the marching cubes algorithm in skimage
|
||||
"""
|
||||
from skimage import measure
|
||||
|
||||
from ..base import Trimesh
|
||||
|
||||
if threshold is not None:
|
||||
matrix = np.asanyarray(matrix) > threshold
|
||||
else:
|
||||
matrix = np.asanyarray(matrix, dtype=bool)
|
||||
|
||||
rev_matrix = np.logical_not(matrix) # Takes set about 0.
|
||||
# Add in padding so marching cubes can function properly with
|
||||
# voxels on edge of AABB
|
||||
pad_width = 1
|
||||
rev_matrix = np.pad(
|
||||
rev_matrix, pad_width=(pad_width), mode="constant", constant_values=(1)
|
||||
)
|
||||
|
||||
# pick between old and new API
|
||||
if hasattr(measure, "marching_cubes_lewiner"):
|
||||
func = measure.marching_cubes_lewiner
|
||||
else:
|
||||
func = measure.marching_cubes
|
||||
|
||||
# Run marching cubes.
|
||||
pitch = np.asanyarray(pitch)
|
||||
if pitch.size == 1:
|
||||
pitch = (pitch,) * 3
|
||||
meshed = func(
|
||||
volume=rev_matrix,
|
||||
level=0.5,
|
||||
spacing=pitch, # it is a boolean voxel grid
|
||||
)
|
||||
|
||||
# allow results from either marching cubes function in skimage
|
||||
# binaries available for python 3.3 and 3.4 appear to use the classic
|
||||
# method
|
||||
if len(meshed) == 2:
|
||||
log.warning("using old marching cubes, may not be watertight!")
|
||||
vertices, faces = meshed
|
||||
normals = None
|
||||
elif len(meshed) == 4:
|
||||
vertices, faces, normals, _vals = meshed
|
||||
|
||||
# Return to the origin, add in the pad_width
|
||||
vertices = np.subtract(vertices, pad_width * pitch)
|
||||
# create the mesh
|
||||
mesh = Trimesh(vertices=vertices, faces=faces, vertex_normals=normals)
|
||||
return mesh
|
||||
|
||||
|
||||
def sparse_to_matrix(sparse):
|
||||
"""
|
||||
Take a sparse (n,3) list of integer indexes of filled cells,
|
||||
turn it into a dense (m,o,p) matrix.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
sparse : (n, 3) int
|
||||
Index of filled cells
|
||||
|
||||
Returns
|
||||
------------
|
||||
dense : (m, o, p) bool
|
||||
Matrix of filled cells
|
||||
"""
|
||||
|
||||
sparse = np.asanyarray(sparse, dtype=np.int64)
|
||||
if not util.is_shape(sparse, (-1, 3)):
|
||||
raise ValueError("sparse must be (n,3)!")
|
||||
|
||||
shape = sparse.max(axis=0) + 1
|
||||
matrix = np.zeros(np.prod(shape), dtype=bool)
|
||||
multiplier = np.array([np.prod(shape[1:]), shape[2], 1])
|
||||
|
||||
index = (sparse * multiplier).sum(axis=1)
|
||||
matrix[index] = True
|
||||
|
||||
dense = matrix.reshape(shape)
|
||||
return dense
|
||||
|
||||
|
||||
def points_to_marching_cubes(points, pitch=1.0):
|
||||
"""
|
||||
Mesh points by assuming they fill a voxel box, and then
|
||||
running marching cubes on them
|
||||
|
||||
Parameters
|
||||
------------
|
||||
points : (n, 3) float
|
||||
Points in 3D space
|
||||
|
||||
Returns
|
||||
-------------
|
||||
mesh : trimesh.Trimesh
|
||||
Points meshed using marching cubes
|
||||
"""
|
||||
# make sure inputs are as expected
|
||||
points = np.asanyarray(points, dtype=np.float64)
|
||||
pitch = np.asanyarray(pitch, dtype=float)
|
||||
|
||||
# find the minimum value of points for origin
|
||||
origin = points.min(axis=0)
|
||||
# convert points to occupied voxel cells
|
||||
index = ((points - origin) / pitch).round().astype(np.int64)
|
||||
|
||||
# convert voxel indices to a matrix
|
||||
matrix = sparse_to_matrix(index)
|
||||
|
||||
# run marching cubes on the matrix to generate a mesh
|
||||
mesh = matrix_to_marching_cubes(matrix, pitch=pitch)
|
||||
mesh.vertices += origin
|
||||
|
||||
return mesh
|
||||
|
||||
|
||||
def multibox(centers, pitch=1.0, colors=None, remove_internal_faces=False):
|
||||
"""
|
||||
Return a Trimesh object with a box at every center.
|
||||
|
||||
Doesn't do anything nice or fancy.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
centers : (n, 3) float
|
||||
Center of boxes that are occupied
|
||||
pitch : float
|
||||
The edge length of a voxel
|
||||
colors : (3,) or (4,) or (n,3) or (n, 4) float
|
||||
Color of boxes
|
||||
remove_internal_faces : bool
|
||||
If True, removes internal faces shared between adjacent boxes
|
||||
|
||||
Returns
|
||||
---------
|
||||
rough : Trimesh
|
||||
Mesh object representing inputs
|
||||
"""
|
||||
from .. import primitives
|
||||
from ..base import Trimesh
|
||||
|
||||
# get centers as numpy array
|
||||
centers = np.asanyarray(centers, dtype=np.float64)
|
||||
|
||||
# get a basic box
|
||||
b = primitives.Box()
|
||||
# apply the pitch
|
||||
b.apply_scale(float(pitch))
|
||||
# tile into one box vertex per center
|
||||
v = np.tile(centers, (1, len(b.vertices))).reshape((-1, 3))
|
||||
# offset to centers
|
||||
v += np.tile(b.vertices, (len(centers), 1))
|
||||
|
||||
f = np.tile(b.faces, (len(centers), 1))
|
||||
f += np.repeat(np.arange(len(centers)) * len(b.vertices), len(b.faces))[:, None]
|
||||
|
||||
if remove_internal_faces:
|
||||
# Get 12 unit normals (1 per triangle face) indicating face direction
|
||||
base_normals = np.round(b.face_normals).astype(int) # (12, 3)
|
||||
# Expand those directions across all voxel boxes so as to check neighbor presence
|
||||
face_normals = np.tile(base_normals, (len(centers), 1)) # (len(centers) * 12, 3)
|
||||
|
||||
# Maps each face to the voxel box it came from
|
||||
face_voxel_idx = np.repeat(
|
||||
np.arange(len(centers)), len(b.faces)
|
||||
) # (len(centers) * 12, )
|
||||
# Converts voxel centers to discrete grid coordinates
|
||||
voxel_coords = np.round(centers / pitch).astype(int) # (len(centers), 3)
|
||||
# Creates a fast lookup structure for checking voxel neighbors
|
||||
voxel_set = set(map(tuple, voxel_coords))
|
||||
|
||||
# Gets the grid coordinate of the voxel that owns each face
|
||||
voxel_face_coords = voxel_coords[face_voxel_idx]
|
||||
# Computes the adjacent voxel coordinate in the face direction
|
||||
neighbor_coords = voxel_face_coords + face_normals
|
||||
# Keeps only faces whose neighboring voxel does not exist
|
||||
keep_mask = np.array([tuple(c) not in voxel_set for c in neighbor_coords])
|
||||
else:
|
||||
keep_mask = np.ones(len(f), dtype=bool)
|
||||
|
||||
face_colors = None
|
||||
if colors is not None:
|
||||
colors = np.asarray(colors)
|
||||
if colors.ndim == 1:
|
||||
colors = colors[None].repeat(len(centers), axis=0)
|
||||
if colors.ndim == 2 and len(colors) == len(centers):
|
||||
face_colors = colors.repeat(12, axis=0)[keep_mask]
|
||||
|
||||
mesh = Trimesh(vertices=v, faces=f[keep_mask], face_colors=face_colors)
|
||||
|
||||
return mesh
|
||||
|
||||
|
||||
def boolean_sparse(a, b, operation=np.logical_and):
|
||||
"""
|
||||
Find common rows between two arrays very quickly
|
||||
using 3D boolean sparse matrices.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
a: (n, d) int, coordinates in space
|
||||
b: (m, d) int, coordinates in space
|
||||
operation: numpy operation function, ie:
|
||||
np.logical_and
|
||||
np.logical_or
|
||||
|
||||
Returns
|
||||
-----------
|
||||
coords: (q, d) int, coordinates in space
|
||||
"""
|
||||
# 3D sparse arrays, using wrapped scipy.sparse
|
||||
# pip install sparse
|
||||
import sparse
|
||||
|
||||
# find the bounding box of both arrays
|
||||
extrema = np.array([a.min(axis=0), a.max(axis=0), b.min(axis=0), b.max(axis=0)])
|
||||
origin = extrema.min(axis=0) - 1
|
||||
size = tuple(np.ptp(extrema, axis=0) + 2)
|
||||
|
||||
# put nearby voxel arrays into same shape sparse array
|
||||
sp_a = sparse.COO((a - origin).T, data=np.ones(len(a), dtype=bool), shape=size)
|
||||
sp_b = sparse.COO((b - origin).T, data=np.ones(len(b), dtype=bool), shape=size)
|
||||
|
||||
# apply the logical operation
|
||||
# get a sparse matrix out
|
||||
applied = operation(sp_a, sp_b)
|
||||
# reconstruct the original coordinates
|
||||
coords = np.column_stack(applied.coords) + origin
|
||||
|
||||
return coords
|
||||
|
||||
|
||||
def strip_array(data):
|
||||
shape = data.shape
|
||||
ndims = len(shape)
|
||||
padding = []
|
||||
slices = []
|
||||
for dim in range(len(shape)):
|
||||
axis = tuple(range(dim)) + tuple(range(dim + 1, ndims))
|
||||
filled = np.any(data, axis=axis)
|
||||
(indices,) = np.nonzero(filled)
|
||||
pad_left = indices[0]
|
||||
pad_right = indices[-1]
|
||||
padding.append([pad_left, pad_right])
|
||||
slices.append(slice(pad_left, pad_right))
|
||||
return data[tuple(slices)], np.array(padding, int)
|
||||
|
||||
|
||||
def indices_to_points(indices, pitch=None, origin=None):
|
||||
"""
|
||||
Convert indices of an (n,m,p) matrix into a set of voxel center points.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices: (q, 3) int, index of voxel matrix (n,m,p)
|
||||
pitch: float, what pitch was the voxel matrix computed with
|
||||
origin: (3,) float, what is the origin of the voxel matrix
|
||||
|
||||
Returns
|
||||
----------
|
||||
points: (q, 3) float, list of points
|
||||
"""
|
||||
indices = np.asanyarray(indices)
|
||||
if indices.shape[1:] != (3,):
|
||||
raise ValueError("shape of indices must be (q, 3)")
|
||||
|
||||
points = np.array(indices, dtype=np.float64)
|
||||
if pitch is not None:
|
||||
points *= float(pitch)
|
||||
if origin is not None:
|
||||
origin = np.asanyarray(origin)
|
||||
if origin.shape != (3,):
|
||||
raise ValueError("shape of origin must be (3,)")
|
||||
points += origin
|
||||
|
||||
return points
|
||||
|
||||
|
||||
def matrix_to_points(matrix, pitch=None, origin=None):
|
||||
"""
|
||||
Convert an (n,m,p) matrix into a set of points for each voxel center.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
matrix: (n,m,p) bool, voxel matrix
|
||||
pitch: float, what pitch was the voxel matrix computed with
|
||||
origin: (3,) float, what is the origin of the voxel matrix
|
||||
|
||||
Returns
|
||||
----------
|
||||
points: (q, 3) list of points
|
||||
"""
|
||||
indices = np.column_stack(np.nonzero(matrix))
|
||||
points = indices_to_points(indices=indices, pitch=pitch, origin=origin)
|
||||
return points
|
||||
|
||||
|
||||
def points_to_indices(points, pitch=None, origin=None):
|
||||
"""
|
||||
Convert center points of an (n,m,p) matrix into its indices.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points : (q, 3) float
|
||||
Center points of voxel matrix (n,m,p)
|
||||
pitch : float
|
||||
What pitch was the voxel matrix computed with
|
||||
origin : (3,) float
|
||||
What is the origin of the voxel matrix
|
||||
|
||||
Returns
|
||||
----------
|
||||
indices : (q, 3) int
|
||||
List of indices
|
||||
"""
|
||||
points = np.array(points, dtype=np.float64)
|
||||
if points.shape != (points.shape[0], 3):
|
||||
raise ValueError("shape of points must be (q, 3)")
|
||||
|
||||
if origin is not None:
|
||||
origin = np.asanyarray(origin)
|
||||
if origin.shape != (3,):
|
||||
raise ValueError("shape of origin must be (3,)")
|
||||
points -= origin
|
||||
if pitch is not None:
|
||||
points /= pitch
|
||||
|
||||
origin = np.asanyarray(origin, dtype=np.float64)
|
||||
pitch = float(pitch)
|
||||
|
||||
indices = np.round(points).astype(int)
|
||||
return indices
|
||||
@@ -0,0 +1,718 @@
|
||||
"""
|
||||
Numpy encode/decode/utility implementations for run length encodings.
|
||||
|
||||
# Run Length Encoded Features
|
||||
|
||||
Encoding/decoding functions for run length encoded data.
|
||||
|
||||
We include code for two variations:
|
||||
|
||||
* run length encoding (RLE)
|
||||
* binary run length encdoing (BRLE)
|
||||
|
||||
RLE stores sequences of repeated values as the value followed by its count, e.g.
|
||||
|
||||
```python
|
||||
dense_to_rle([5, 5, 3, 2, 2, 2, 2, 6]) == [5, 2, 3, 1, 2, 4, 6, 1]
|
||||
```
|
||||
|
||||
i.e. the value `5` is repeated `2` times, then `3` is repeated `1` time, `2` is
|
||||
repeated `4` times and `6` is repeated `1` time.
|
||||
|
||||
BRLE is an optimized form for when the stored values can only be `0` or `1`.
|
||||
This means we only need to save the counts, and assume the values alternate
|
||||
(starting at `0`).
|
||||
|
||||
```python
|
||||
dense_to_brle([1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 0]) == \
|
||||
[0, 2, 4, 7, 2]
|
||||
```
|
||||
|
||||
i.e. the value zero occurs `0` times, followed by `2` ones, `4` zeros, `7` ones
|
||||
and `2` zeros.
|
||||
|
||||
Sequences with counts exceeding the data type's maximum value have to be
|
||||
handled carefully. For example, the `uint8` encoding of 300 zeros
|
||||
(`uint8` has a max value of 255) is:
|
||||
|
||||
* RLE: `[0, 255, 0, 45]` (`0` repeated `255` times + `0` repeated `45` times)
|
||||
* BRLE: `[255, 0, 45, 0]` (`255` zeros + `0` ones + `45` zeros + `0` ones)
|
||||
|
||||
This module contains implementations of various RLE/BRLE operations.
|
||||
"""
|
||||
|
||||
import functools
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def brle_length(brle):
|
||||
"""Optimized implementation of `len(brle_to_dense(brle))`"""
|
||||
return np.sum(brle)
|
||||
|
||||
|
||||
def rle_length(rle):
|
||||
"""Optimized implementation of `len(rle_to_dense(rle_to_brle(rle)))`"""
|
||||
return np.sum(rle[1::2])
|
||||
|
||||
|
||||
def rle_to_brle(rle, dtype=None):
|
||||
"""
|
||||
Convert run length encoded (RLE) value/counts to BRLE.
|
||||
|
||||
RLE data is stored in a rank 1 array with each pair giving:
|
||||
(value, count)
|
||||
|
||||
e.g. the RLE encoding of [4, 4, 4, 1, 1, 6] is [4, 3, 1, 2, 6, 1].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rle : (n,) int
|
||||
Run length encoded data
|
||||
|
||||
Returns
|
||||
----------
|
||||
equivalent binary run length encoding. a list if dtype is None,
|
||||
otherwise brle_to_brle is called on that list before returning.
|
||||
|
||||
Raises
|
||||
----------
|
||||
ValueError
|
||||
If any of the even counts of `rle` are not zero or 1.
|
||||
"""
|
||||
curr_val = 0
|
||||
out = [0]
|
||||
acc = 0
|
||||
for value, count in np.reshape(rle, (-1, 2)):
|
||||
acc += count
|
||||
if value not in (0, 1):
|
||||
raise ValueError("Invalid run length encoding for conversion to BRLE")
|
||||
if value == curr_val:
|
||||
out[-1] += count
|
||||
else:
|
||||
out.append(int(count))
|
||||
curr_val = value
|
||||
if len(out) % 2:
|
||||
out.append(0)
|
||||
if dtype is not None:
|
||||
out = brle_to_brle(out, dtype=dtype)
|
||||
return out
|
||||
|
||||
|
||||
def brle_logical_not(brle):
|
||||
"""
|
||||
Get the BRLE encoding of the `logical_not`ed dense form of `brle`.
|
||||
|
||||
Equivalent to `dense_to_brle(np.logical_not(brle_to_dense(brle)))` but
|
||||
highly optimized - just pads brle with a 0 on each end (or strips is
|
||||
existing endpoints are both zero).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
brle: rank 1 int array of binary run length encoded data
|
||||
|
||||
Returns
|
||||
----------
|
||||
rank 1 int array of binary run length encoded data corresponding to
|
||||
element-wise not of the input.
|
||||
"""
|
||||
if brle[0] or brle[-1]:
|
||||
return np.pad(brle, [1, 1], mode="constant")
|
||||
else:
|
||||
return brle[1:-1]
|
||||
|
||||
|
||||
def merge_brle_lengths(lengths):
|
||||
"""Inverse of split_long_brle_lengths."""
|
||||
if len(lengths) == 0:
|
||||
return []
|
||||
|
||||
out = [int(lengths[0])]
|
||||
accumulating = False
|
||||
for length in lengths[1:]:
|
||||
if accumulating:
|
||||
out[-1] += length
|
||||
accumulating = False
|
||||
else:
|
||||
if length == 0:
|
||||
accumulating = True
|
||||
else:
|
||||
out.append(int(length))
|
||||
return out
|
||||
|
||||
|
||||
def split_long_brle_lengths(lengths, dtype=np.int64):
|
||||
"""
|
||||
Split lengths that exceed max dtype value.
|
||||
|
||||
Lengths `l` are converted into [max_val, 0] * l // max_val + [l % max_val]
|
||||
|
||||
e.g. for dtype=np.uint8 (max_value == 255)
|
||||
```
|
||||
split_long_brle_lengths([600, 300, 2, 6], np.uint8) == \
|
||||
[255, 0, 255, 0, 90, 255, 0, 45, 2, 6]
|
||||
```
|
||||
"""
|
||||
lengths = np.asarray(lengths)
|
||||
max_val = np.iinfo(dtype).max
|
||||
bad_length_mask = lengths > max_val
|
||||
if np.any(bad_length_mask):
|
||||
# there are some bad lengths
|
||||
nl = len(lengths)
|
||||
repeats = np.asarray(lengths) // max_val
|
||||
remainders = (lengths % max_val).astype(dtype)
|
||||
|
||||
lengths = np.concatenate(
|
||||
[
|
||||
np.array([max_val, 0] * repeat + [remainder], dtype=dtype)
|
||||
for repeat, remainder in zip(repeats, remainders)
|
||||
]
|
||||
)
|
||||
lengths = lengths.reshape((np.sum(repeats) * 2 + nl,)).astype(dtype)
|
||||
return lengths
|
||||
elif lengths.dtype != dtype:
|
||||
return lengths.astype(dtype)
|
||||
else:
|
||||
return lengths
|
||||
|
||||
|
||||
def dense_to_brle(dense_data, dtype=np.int64):
|
||||
"""
|
||||
Get the binary run length encoding of `dense_data`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dense_data: rank 1 bool array of data to encode.
|
||||
dtype: numpy int type.
|
||||
|
||||
Returns
|
||||
----------
|
||||
Binary run length encoded rank 1 array of dtype `dtype`.
|
||||
|
||||
Raises
|
||||
----------
|
||||
ValuError if dense_data is not a rank 1 bool array.
|
||||
"""
|
||||
if dense_data.dtype != bool:
|
||||
raise ValueError("`dense_data` must be bool")
|
||||
if len(dense_data.shape) != 1:
|
||||
raise ValueError("`dense_data` must be rank 1.")
|
||||
n = len(dense_data)
|
||||
starts = np.r_[0, np.flatnonzero(dense_data[1:] != dense_data[:-1]) + 1]
|
||||
lengths = np.diff(np.r_[starts, n])
|
||||
lengths = split_long_brle_lengths(lengths, dtype=dtype)
|
||||
if dense_data[0]:
|
||||
lengths = np.pad(lengths, [1, 0], mode="constant")
|
||||
return lengths
|
||||
|
||||
|
||||
_ft = np.array([False, True], dtype=bool)
|
||||
|
||||
|
||||
def brle_to_dense(brle_data, vals=None):
|
||||
"""Decode binary run length encoded data to dense.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
brle_data: BRLE counts of False/True values
|
||||
vals: if not `None`, a length 2 array/list/tuple with False/True substitute
|
||||
values, e.g. brle_to_dense([2, 3, 1, 0], [7, 9]) == [7, 7, 9, 9, 9, 7]
|
||||
|
||||
Returns
|
||||
----------
|
||||
rank 1 dense data of dtype `bool if vals is None else vals.dtype`
|
||||
|
||||
Raises
|
||||
----------
|
||||
ValueError if vals it not None and shape is not (2,)
|
||||
"""
|
||||
if vals is None:
|
||||
vals = _ft
|
||||
else:
|
||||
vals = np.asarray(vals)
|
||||
if vals.shape != (2,):
|
||||
raise ValueError(f"vals.shape must be (2,), got {vals.shape}")
|
||||
ft = np.repeat(_ft[np.newaxis, :], (len(brle_data) + 1) // 2, axis=0).flatten()
|
||||
return np.repeat(ft[: len(brle_data)], brle_data).flatten()
|
||||
|
||||
|
||||
def rle_to_dense(rle_data, dtype=np.int64):
|
||||
"""Get the dense decoding of the associated run length encoded data."""
|
||||
values, counts = np.split(np.reshape(rle_data, (-1, 2)), 2, axis=-1)
|
||||
if dtype is not None:
|
||||
values = np.asanyarray(values, dtype=dtype)
|
||||
try:
|
||||
result = np.repeat(np.squeeze(values, axis=-1), np.squeeze(counts, axis=-1))
|
||||
except TypeError:
|
||||
# on windows it sometimes fails to cast data type
|
||||
result = np.repeat(
|
||||
np.squeeze(values.astype(np.int64), axis=-1),
|
||||
np.squeeze(counts.astype(np.int64), axis=-1),
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def dense_to_rle(dense_data, dtype=np.int64):
|
||||
"""Get run length encoding of the provided dense data."""
|
||||
n = len(dense_data)
|
||||
starts = np.r_[0, np.flatnonzero(dense_data[1:] != dense_data[:-1]) + 1]
|
||||
lengths = np.diff(np.r_[starts, n])
|
||||
values = dense_data[starts]
|
||||
values, lengths = split_long_rle_lengths(values, lengths, dtype=dtype)
|
||||
out = np.stack((values, lengths), axis=1)
|
||||
return out.flatten()
|
||||
|
||||
|
||||
def split_long_rle_lengths(values, lengths, dtype=np.int64):
|
||||
"""
|
||||
Split long lengths in the associated run length encoding.
|
||||
|
||||
e.g.
|
||||
```python
|
||||
split_long_rle_lengths([5, 300, 2, 12], np.uint8) == [5, 255, 5, 45, 2, 12]
|
||||
```
|
||||
|
||||
Parameters
|
||||
----------
|
||||
values: values column of run length encoding, or `rle[::2]`
|
||||
lengths: counts in run length encoding, or `rle[1::2]`
|
||||
dtype: numpy data type indicating the maximum value.
|
||||
|
||||
Returns
|
||||
----------
|
||||
values, lengths associated with the appropriate splits. `lengths` will be
|
||||
of type `dtype`, while `values` will be the same as the value passed in.
|
||||
"""
|
||||
max_length = np.iinfo(dtype).max
|
||||
lengths = np.asarray(lengths)
|
||||
repeats = lengths // max_length
|
||||
if np.any(repeats):
|
||||
repeats += 1
|
||||
remainder = lengths % max_length
|
||||
values = np.repeat(values, repeats)
|
||||
lengths = np.zeros(len(repeats), dtype=dtype)
|
||||
lengths.fill(max_length)
|
||||
lengths = np.repeat(lengths, repeats)
|
||||
lengths[np.cumsum(repeats) - 1] = remainder
|
||||
elif lengths.dtype != dtype:
|
||||
lengths = lengths.astype(dtype)
|
||||
return values, lengths
|
||||
|
||||
|
||||
def merge_rle_lengths(values, lengths):
|
||||
"""Inverse of split_long_rle_lengths except returns normal python lists."""
|
||||
ret_values = []
|
||||
ret_lengths = []
|
||||
curr = None
|
||||
for value, length in zip(values, lengths):
|
||||
if length == 0:
|
||||
continue
|
||||
if value == curr:
|
||||
ret_lengths[-1] += length
|
||||
else:
|
||||
curr = value
|
||||
ret_lengths.append(int(length))
|
||||
ret_values.append(value)
|
||||
return ret_values, ret_lengths
|
||||
|
||||
|
||||
def brle_to_rle(brle, dtype=np.int64):
|
||||
if len(brle) % 2 == 1:
|
||||
brle = np.concatenate([brle, [0]])
|
||||
lengths = brle
|
||||
values = np.tile(_ft, len(brle) // 2)
|
||||
return rle_to_rle(np.stack((values, lengths), axis=1).flatten(), dtype=dtype)
|
||||
|
||||
|
||||
def brle_to_brle(brle, dtype=np.int64):
|
||||
"""
|
||||
Almost the identity function.
|
||||
|
||||
Checks for possible merges and required splits.
|
||||
"""
|
||||
return split_long_brle_lengths(merge_brle_lengths(brle), dtype=dtype)
|
||||
|
||||
|
||||
def rle_to_rle(rle, dtype=np.int64):
|
||||
"""
|
||||
Almost the identity function.
|
||||
|
||||
Checks for possible merges and required splits.
|
||||
"""
|
||||
values, lengths = np.reshape(rle, (-1, 2)).T
|
||||
values, lengths = merge_rle_lengths(values, lengths)
|
||||
values, lengths = split_long_rle_lengths(values, lengths, dtype=dtype)
|
||||
return np.stack((values, lengths), axis=1).flatten()
|
||||
|
||||
|
||||
def _unsorted_gatherer(indices, sorted_gather_fn):
|
||||
if not isinstance(indices, np.ndarray):
|
||||
indices = np.array(indices, copy=False)
|
||||
order = np.argsort(indices)
|
||||
ordered_indices = indices[order]
|
||||
|
||||
def f(data, dtype=None):
|
||||
result = np.zeros(len(order), dtype=dtype or getattr(data, "dtype", None))
|
||||
result[order] = tuple(sorted_gather_fn(data, ordered_indices))
|
||||
return result
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def sorted_rle_gather_1d(rle_data, ordered_indices):
|
||||
"""
|
||||
Gather brle_data at ordered_indices.
|
||||
|
||||
This is equivalent to `rle_to_dense(brle_data)[ordered_indices]` but avoids
|
||||
the decoding.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
brle_data: iterable of run-length-encoded data.
|
||||
ordered_indices: iterable of ints in ascending order.
|
||||
|
||||
Returns
|
||||
----------
|
||||
`brle_data` iterable of values at the dense indices, same length as
|
||||
ordered indices.
|
||||
"""
|
||||
data_iter = iter(rle_data)
|
||||
index_iter = iter(ordered_indices)
|
||||
try:
|
||||
index = next(index_iter)
|
||||
except StopIteration:
|
||||
return
|
||||
start = 0
|
||||
while True:
|
||||
while start <= index:
|
||||
try:
|
||||
value = next(data_iter)
|
||||
start += next(data_iter)
|
||||
except StopIteration:
|
||||
raise IndexError(
|
||||
"Index %d out of range of raw_values length %d", index, start
|
||||
)
|
||||
|
||||
try:
|
||||
while index < start:
|
||||
yield value
|
||||
index = next(index_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
|
||||
|
||||
def rle_mask(rle_data, mask):
|
||||
"""
|
||||
Perform masking of the input run-length data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rle_data: iterable of run length encoded data
|
||||
mask: iterable of bools corresponding to the dense mask.
|
||||
|
||||
Returns
|
||||
----------
|
||||
iterable of dense values of rle_data wherever mask is True.
|
||||
"""
|
||||
data_iter = iter(rle_data)
|
||||
mask_iter = iter(mask)
|
||||
while True:
|
||||
try:
|
||||
value = next(data_iter)
|
||||
count = next(data_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
for _ in range(count):
|
||||
m = next(mask_iter)
|
||||
if m:
|
||||
yield value
|
||||
|
||||
|
||||
def brle_mask(rle_data, mask):
|
||||
"""
|
||||
Perform masking of the input binary run-length data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
brle_data: iterable of binary run length encoded data
|
||||
mask: iterable of bools corresponding to the dense mask.
|
||||
|
||||
Returns
|
||||
----------
|
||||
iterable dense values of brle_data wherever mask is True.
|
||||
"""
|
||||
data_iter = iter(rle_data)
|
||||
mask_iter = iter(mask)
|
||||
value = True
|
||||
while True:
|
||||
try:
|
||||
value = not value
|
||||
count = next(data_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
for _ in range(count):
|
||||
m = next(mask_iter)
|
||||
if m:
|
||||
yield value
|
||||
|
||||
|
||||
def rle_gatherer_1d(indices):
|
||||
"""
|
||||
Get a gather function at the given indices.
|
||||
|
||||
Because gathering on RLE data requires sorting, for instances where
|
||||
gathering at the same indices on different RLE data this can save the
|
||||
sorting process.
|
||||
|
||||
If only gathering on a single RLE iterable, use `rle_gather_1d`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices: iterable of integers
|
||||
|
||||
Returns
|
||||
----------
|
||||
gather function, mapping `(rle_data, dtype=None) -> values`.
|
||||
`values` will have the same length as `indices` and dtype provided,
|
||||
or rle_data.dtype if no dtype is provided.
|
||||
"""
|
||||
return _unsorted_gatherer(indices, sorted_rle_gather_1d)
|
||||
|
||||
|
||||
def rle_gather_1d(rle_data, indices, dtype=None):
|
||||
"""
|
||||
Gather RLE data values at the provided dense indices.
|
||||
|
||||
This is equivalent to `rle_to_dense(rle_data)[indices]` but the
|
||||
implementation does not require the construction of the dense array.
|
||||
|
||||
If indices is known to be in order, use `sorted_gather_1d`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rle_data: run length encoded data
|
||||
indices: dense indices
|
||||
dtype: numpy dtype. If not provided, uses rle_data.dtype
|
||||
|
||||
Returns
|
||||
----------
|
||||
numpy array, dense data at indices, same length as indices and dtype as
|
||||
rle_data
|
||||
"""
|
||||
return rle_gatherer_1d(indices)(rle_data, dtype=dtype)
|
||||
|
||||
|
||||
def sorted_brle_gather_1d(brle_data, ordered_indices):
|
||||
"""
|
||||
Gather brle_data at ordered_indices.
|
||||
|
||||
This is equivalent to `brle_to_dense(brle_data)[ordered_indices]` but
|
||||
avoids the decoding.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
raw_data: iterable of run-length-encoded data.
|
||||
ordered_indices: iterable of ints in ascending order.
|
||||
|
||||
Returns
|
||||
----------
|
||||
`raw_data` iterable of values at the dense indices, same length as
|
||||
ordered indices.
|
||||
"""
|
||||
data_iter = iter(brle_data)
|
||||
index_iter = iter(ordered_indices)
|
||||
try:
|
||||
index = next(index_iter)
|
||||
except StopIteration:
|
||||
return
|
||||
start = 0
|
||||
value = True
|
||||
while True:
|
||||
while start <= index:
|
||||
try:
|
||||
value = not value
|
||||
start += next(data_iter)
|
||||
except StopIteration:
|
||||
raise IndexError(
|
||||
"Index %d out of range of raw_values length %d", index, start
|
||||
)
|
||||
|
||||
try:
|
||||
while index < start:
|
||||
yield value
|
||||
index = next(index_iter)
|
||||
except StopIteration:
|
||||
break
|
||||
|
||||
|
||||
def brle_gatherer_1d(indices):
|
||||
"""
|
||||
Get a gather function at the given indices.
|
||||
|
||||
Because gathering on BRLE data requires sorting, for instances where
|
||||
gathering at the same indices on different RLE data this can save the
|
||||
sorting process.
|
||||
|
||||
If only gathering on a single RLE iterable, use `brle_gather_1d`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
indices: iterable of integers
|
||||
|
||||
Returns
|
||||
----------
|
||||
gather function, mapping `(rle_data, dtype=None) -> values`.
|
||||
`values` will have the same length as `indices` and dtype provided,
|
||||
or rle_data.dtype if no dtype is provided.
|
||||
"""
|
||||
return functools.partial(
|
||||
_unsorted_gatherer(indices, sorted_brle_gather_1d), dtype=bool
|
||||
)
|
||||
|
||||
|
||||
def brle_gather_1d(brle_data, indices):
|
||||
"""
|
||||
Gather BRLE data values at the provided dense indices.
|
||||
|
||||
This is equivalent to `rle_to_dense(rle_data)[indices]` but the
|
||||
implementation does not require the construction of the dense array.
|
||||
|
||||
If indices is known to be in order, use `sorted_brle_gather_1d`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rle_data: run length encoded data
|
||||
indices: dense indices
|
||||
|
||||
Returns
|
||||
----------
|
||||
numpy array, dense data at indices, same length as indices and dtype as
|
||||
rle_data
|
||||
"""
|
||||
return brle_gatherer_1d(indices)(brle_data)
|
||||
|
||||
|
||||
def brle_reverse(brle_data):
|
||||
"""Equivalent to dense_to_brle(brle_to_dense(brle_data)[-1::-1])."""
|
||||
if len(brle_data) % 2 == 0:
|
||||
brle_data = np.concatenate([brle_data, [0]], axis=0)
|
||||
end = -1 if brle_data[-1] == 0 else None
|
||||
return brle_data[-1:end:-1]
|
||||
|
||||
|
||||
def rle_reverse(rle_data):
|
||||
"""Get the rle encoding of the reversed dense array."""
|
||||
if not isinstance(rle_data, np.ndarray):
|
||||
rle_data = np.array(rle_data, copy=False)
|
||||
rle_data = np.reshape(rle_data, (-1, 2))
|
||||
rle_data = rle_data[-1::-1]
|
||||
return np.reshape(rle_data, (-1,))
|
||||
|
||||
|
||||
def rle_to_sparse(rle_data):
|
||||
"""Get dense indices associated with non-zeros."""
|
||||
indices = []
|
||||
values = []
|
||||
it = iter(rle_data)
|
||||
index = 0
|
||||
try:
|
||||
while True:
|
||||
value = next(it)
|
||||
counts = int(next(it))
|
||||
end = index + counts
|
||||
if value:
|
||||
indices.append(np.arange(index, end, dtype=np.int64))
|
||||
values.append(np.repeat(value, counts))
|
||||
index = end
|
||||
except StopIteration:
|
||||
pass
|
||||
if len(indices) == 0:
|
||||
assert len(values) == 0
|
||||
return indices, values
|
||||
|
||||
indices = np.concatenate(indices)
|
||||
values = np.concatenate(values, dtype=rle_data.dtype)
|
||||
return indices, values
|
||||
|
||||
|
||||
def brle_to_sparse(brle_data, dtype=np.int64):
|
||||
ends = np.cumsum(brle_data)
|
||||
indices = [np.arange(s, e, dtype=dtype) for s, e in zip(ends[::2], ends[1::2])]
|
||||
return np.concatenate(indices)
|
||||
|
||||
|
||||
def rle_strip(rle_data):
|
||||
"""
|
||||
Remove leading and trailing zeros.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rle_data: run length encoded data
|
||||
|
||||
Returns
|
||||
----------
|
||||
(stripped_rle_data, padding)
|
||||
stripped_rle_data: rle data without any leading or trailing zeros
|
||||
padding: 2-element dense padding
|
||||
"""
|
||||
rle_data = np.reshape(rle_data, (-1, 2))
|
||||
start = 0
|
||||
final_i = len(rle_data)
|
||||
for i, (val, count) in enumerate(rle_data):
|
||||
if val and count > 0:
|
||||
final_i = i
|
||||
break
|
||||
else:
|
||||
start += count
|
||||
|
||||
end = 0
|
||||
final_j = len(rle_data)
|
||||
for j, (val, count) in enumerate(rle_data[::-1]):
|
||||
if val and count > 0:
|
||||
final_j = j
|
||||
break
|
||||
else:
|
||||
end += count
|
||||
|
||||
rle_data = rle_data[final_i : None if final_j == 0 else -final_j].reshape((-1,))
|
||||
return rle_data, (start, end)
|
||||
|
||||
|
||||
def brle_strip(brle_data):
|
||||
"""
|
||||
Remove leading and trailing zeros.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
brle_data: binary run length encoded data.
|
||||
|
||||
Returns
|
||||
----------
|
||||
(stripped_brle_data, padding)
|
||||
stripped_brle_data: rle data without any leading or trailing zeros
|
||||
padding: 2-element dense padding
|
||||
"""
|
||||
start = 0
|
||||
val = True
|
||||
final_i = len(brle_data)
|
||||
for i, count in enumerate(brle_data):
|
||||
val = not val
|
||||
if val and count > 0:
|
||||
final_i = i
|
||||
break
|
||||
else:
|
||||
start += count
|
||||
end = 0
|
||||
final_j = len(brle_data)
|
||||
val = bool(len(brle_data) % 2)
|
||||
for j, count in enumerate(brle_data[::-1]):
|
||||
val = not val
|
||||
if val and count > 0:
|
||||
final_j = j
|
||||
break
|
||||
else:
|
||||
end += count
|
||||
|
||||
brle_data = brle_data[final_i : None if final_j == 0 else -final_j]
|
||||
brle_data = np.concatenate([[0], brle_data])
|
||||
return brle_data, (start, end)
|
||||
@@ -0,0 +1,181 @@
|
||||
import numpy as np
|
||||
|
||||
from .. import caching, util
|
||||
from .. import transformations as tr
|
||||
from ..typed import Optional
|
||||
|
||||
|
||||
class Transform:
|
||||
"""
|
||||
Class for caching metadata associated with 4x4 transformations.
|
||||
|
||||
The transformation matrix is used to define relevant properties
|
||||
for the voxels, including pitch and origin.
|
||||
"""
|
||||
|
||||
def __init__(self, matrix, datastore: Optional[caching.DataStore] = None):
|
||||
"""
|
||||
Initialize with a transform.
|
||||
|
||||
Parameters
|
||||
-----------
|
||||
matrix : (4, 4) float
|
||||
Homogeneous transformation matrix
|
||||
datastore
|
||||
If passed store the actual values in a reference to
|
||||
another datastore.
|
||||
"""
|
||||
matrix = np.asanyarray(matrix, dtype=np.float64)
|
||||
if matrix.shape != (4, 4) or not np.allclose(matrix[3, :], [0, 0, 0, 1]):
|
||||
raise ValueError("matrix is invalid!")
|
||||
|
||||
# store matrix as data
|
||||
if datastore is None:
|
||||
self._data = caching.DataStore()
|
||||
elif isinstance(datastore, caching.DataStore):
|
||||
self._data = datastore
|
||||
else:
|
||||
raise ValueError(f"{type(datastore)} != caching.DataStore")
|
||||
|
||||
self._data["transform_matrix"] = matrix
|
||||
# dump cache when matrix changes
|
||||
self._cache = caching.Cache(id_function=self._data.__hash__)
|
||||
|
||||
def __hash__(self):
|
||||
"""
|
||||
Get the hash of the current transformation matrix.
|
||||
|
||||
Returns
|
||||
------------
|
||||
hash : str
|
||||
Hash of transformation matrix
|
||||
"""
|
||||
return self._data.__hash__()
|
||||
|
||||
@property
|
||||
def translation(self):
|
||||
"""
|
||||
Get the translation component of the matrix
|
||||
|
||||
Returns
|
||||
------------
|
||||
translation : (3,) float
|
||||
Cartesian translation
|
||||
"""
|
||||
return self._data["transform_matrix"][:3, 3]
|
||||
|
||||
@property
|
||||
def matrix(self):
|
||||
"""
|
||||
Get the homogeneous transformation matrix.
|
||||
|
||||
Returns
|
||||
-------------
|
||||
matrix : (4, 4) float
|
||||
Transformation matrix
|
||||
"""
|
||||
return self._data["transform_matrix"]
|
||||
|
||||
@matrix.setter
|
||||
def matrix(self, values):
|
||||
"""
|
||||
Set the homogeneous transformation matrix.
|
||||
|
||||
Parameters
|
||||
-------------
|
||||
matrix : (4, 4) float
|
||||
Transformation matrix
|
||||
"""
|
||||
values = np.asanyarray(values, dtype=np.float64)
|
||||
if values.shape != (4, 4):
|
||||
raise ValueError("matrix must be (4, 4)!")
|
||||
self._data["transform_matrix"] = values
|
||||
|
||||
@caching.cache_decorator
|
||||
def scale(self):
|
||||
"""
|
||||
Get the scale factor of the current transformation.
|
||||
|
||||
Returns
|
||||
-------------
|
||||
scale : (3,) float
|
||||
Scale factor from the matrix
|
||||
"""
|
||||
# get the current transformation
|
||||
matrix = self.matrix
|
||||
# get the (3,) diagonal of the rotation component
|
||||
scale = np.diag(matrix[:3, :3])
|
||||
if not np.allclose(matrix[:3, :3], scale * np.eye(3), scale * 1e-6 + 1e-8):
|
||||
raise RuntimeError("transform features a shear or rotation")
|
||||
return scale
|
||||
|
||||
@caching.cache_decorator
|
||||
def pitch(self):
|
||||
scale = self.scale
|
||||
if not util.allclose(scale[0], scale[1:], np.max(np.abs(scale)) * 1e-6 + 1e-8):
|
||||
raise RuntimeError("transform features non-uniform scaling")
|
||||
return scale
|
||||
|
||||
@caching.cache_decorator
|
||||
def unit_volume(self):
|
||||
"""Volume of a transformed unit cube."""
|
||||
return np.linalg.det(self._data["transform_matrix"][:3, :3])
|
||||
|
||||
def apply_transform(self, matrix):
|
||||
"""Mutate the transform in-place and return self."""
|
||||
self.matrix = np.matmul(matrix, self.matrix)
|
||||
return self
|
||||
|
||||
def apply_translation(self, translation):
|
||||
"""Mutate the transform in-place and return self."""
|
||||
self.matrix[:3, 3] += translation
|
||||
return self
|
||||
|
||||
def apply_scale(self, scale):
|
||||
"""Mutate the transform in-place and return self."""
|
||||
self.matrix[:3] *= scale
|
||||
return self
|
||||
|
||||
def transform_points(self, points):
|
||||
"""
|
||||
Apply the transformation to points (not in-place).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points: (n, 3) float
|
||||
Points in cartesian space
|
||||
|
||||
Returns
|
||||
----------
|
||||
transformed : (n, 3) float
|
||||
Points transformed by matrix
|
||||
"""
|
||||
if self.is_identity:
|
||||
return points.copy()
|
||||
return tr.transform_points(points.reshape(-1, 3), self.matrix).reshape(
|
||||
points.shape
|
||||
)
|
||||
|
||||
def inverse_transform_points(self, points):
|
||||
"""Apply the inverse transformation to points (not in-place)."""
|
||||
if self.is_identity:
|
||||
return points
|
||||
return tr.transform_points(points.reshape(-1, 3), self.inverse_matrix).reshape(
|
||||
points.shape
|
||||
)
|
||||
|
||||
@caching.cache_decorator
|
||||
def inverse_matrix(self):
|
||||
inv = np.linalg.inv(self.matrix)
|
||||
inv.flags.writeable = False
|
||||
return inv
|
||||
|
||||
def copy(self):
|
||||
return Transform(matrix=self.matrix)
|
||||
|
||||
@caching.cache_decorator
|
||||
def is_identity(self):
|
||||
"""
|
||||
Flags this transformation being sufficiently close to eye(4).
|
||||
"""
|
||||
return util.allclose(self.matrix, np.eye(4), 1e-8)
|
||||
Reference in New Issue
Block a user