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)