1import numpy as np
2
3from .. import util
4from ..constants import log
5
6
7def fill_orthographic(dense):
8    shape = dense.shape
9    indices = np.stack(
10        np.meshgrid(*(np.arange(s) for s in shape), indexing='ij'),
11        axis=-1)
12    empty = np.logical_not(dense)
13
14    def fill_axis(axis):
15        base_local_indices = indices[..., axis]
16        local_indices = base_local_indices.copy()
17        local_indices[empty] = shape[axis]
18        mins = np.min(local_indices, axis=axis, keepdims=True)
19        local_indices = base_local_indices.copy()
20        local_indices[empty] = -1
21        maxs = np.max(local_indices, axis=axis, keepdims=True)
22
23        return np.logical_and(
24            base_local_indices >= mins,
25            base_local_indices <= maxs,
26        )
27
28    filled = fill_axis(axis=0)
29    for axis in range(1, len(shape)):
30        filled = np.logical_and(filled, fill_axis(axis))
31    return filled
32
33
34def fill_base(sparse_indices):
35    """
36    Given a sparse surface voxelization, fill in between columns.
37
38    Parameters
39    --------------
40    sparse_indices: (n, 3) int, location of filled cells
41
42    Returns
43    --------------
44    filled: (m, 3) int, location of filled cells
45    """
46    # validate inputs
47    sparse_indices = np.asanyarray(sparse_indices, dtype=np.int64)
48    if not util.is_shape(sparse_indices, (-1, 3)):
49        raise ValueError('incorrect shape')
50
51    # create grid and mark inner voxels
52    max_value = sparse_indices.max() + 3
53
54    grid = np.zeros((max_value,
55                     max_value,
56                     max_value),
57                    bool)
58    voxels_sparse = np.add(sparse_indices, 1)
59
60    grid[tuple(voxels_sparse.T)] = 1
61
62    for i in range(max_value):
63        check_dir2 = False
64        for j in range(0, max_value - 1):
65            idx = []
66            # find transitions first
67            # transition positions are from 0 to 1 and from 1 to 0
68            eq = np.equal(grid[i, j, :-1], grid[i, j, 1:])
69            idx = np.where(np.logical_not(eq))[0] + 1
70            c = len(idx)
71            check_dir2 = (c % 4) > 0 and c > 4
72            if c < 4:
73                continue
74            for s in range(0, c - c % 4, 4):
75                grid[i, j, idx[s]:idx[s + 3]] = 1
76        if not check_dir2:
77            continue
78
79        # check another direction for robustness
80        for k in range(0, max_value - 1):
81            idx = []
82            # find transitions first
83            eq = np.equal(grid[i, :-1, k], grid[i, 1:, k])
84            idx = np.where(np.logical_not(eq))[0] + 1
85            c = len(idx)
86            if c < 4:
87                continue
88            for s in range(0, c - c % 4, 4):
89                grid[i, idx[s]:idx[s + 3], k] = 1
90
91    # generate new voxels
92    filled = np.column_stack(np.where(grid))
93    filled -= 1
94
95    return filled
96
97
98fill_voxelization = fill_base
99
100
101def matrix_to_marching_cubes(matrix, pitch=1.0):
102    """
103    Convert an (n,m,p) matrix into a mesh, using marching_cubes.
104
105    Parameters
106    -----------
107    matrix : (n, m, p) bool
108      Occupancy array
109
110    Returns
111    ----------
112    mesh : trimesh.Trimesh
113      Mesh generated by meshing voxels using
114      the marching cubes algorithm in skimage
115    """
116    from skimage import measure
117    from ..base import Trimesh
118
119    matrix = np.asanyarray(matrix, dtype=np.bool)
120
121    rev_matrix = np.logical_not(matrix)  # Takes set about 0.
122    # Add in padding so marching cubes can function properly with
123    # voxels on edge of AABB
124    pad_width = 1
125    rev_matrix = np.pad(rev_matrix,
126                        pad_width=(pad_width),
127                        mode='constant',
128                        constant_values=(1))
129
130    # pick between old and new API
131    if hasattr(measure, 'marching_cubes_lewiner'):
132        func = measure.marching_cubes_lewiner
133    else:
134        func = measure.marching_cubes
135
136    # Run marching cubes.
137    pitch = np.asanyarray(pitch)
138    if pitch.size == 1:
139        pitch = (pitch,) * 3
140    meshed = func(volume=rev_matrix,
141                  level=.5,  # it is a boolean voxel grid
142                  spacing=pitch)
143
144    # allow results from either marching cubes function in skimage
145    # binaries available for python 3.3 and 3.4 appear to use the classic
146    # method
147    if len(meshed) == 2:
148        log.warning('using old marching cubes, may not be watertight!')
149        vertices, faces = meshed
150        normals = None
151    elif len(meshed) == 4:
152        vertices, faces, normals, vals = meshed
153
154    # Return to the origin, add in the pad_width
155    vertices = np.subtract(vertices, pad_width * pitch)
156    # create the mesh
157    mesh = Trimesh(vertices=vertices,
158                   faces=faces,
159                   vertex_normals=normals)
160    return mesh
161
162
163def sparse_to_matrix(sparse):
164    """
165    Take a sparse (n,3) list of integer indexes of filled cells,
166    turn it into a dense (m,o,p) matrix.
167
168    Parameters
169    -----------
170    sparse : (n, 3) int
171      Index of filled cells
172
173    Returns
174    ------------
175    dense : (m, o, p) bool
176      Matrix of filled cells
177    """
178
179    sparse = np.asanyarray(sparse, dtype=np.int)
180    if not util.is_shape(sparse, (-1, 3)):
181        raise ValueError('sparse must be (n,3)!')
182
183    shape = sparse.max(axis=0) + 1
184    matrix = np.zeros(np.product(shape), dtype=np.bool)
185    multiplier = np.array([np.product(shape[1:]), shape[2], 1])
186
187    index = (sparse * multiplier).sum(axis=1)
188    matrix[index] = True
189
190    dense = matrix.reshape(shape)
191    return dense
192
193
194def points_to_marching_cubes(points, pitch=1.0):
195    """
196    Mesh points by assuming they fill a voxel box, and then
197    running marching cubes on them
198
199    Parameters
200    ------------
201    points : (n, 3) float
202      Points in 3D space
203
204    Returns
205    -------------
206    mesh : trimesh.Trimesh
207      Points meshed using marching cubes
208    """
209    # make sure inputs are as expected
210    points = np.asanyarray(points, dtype=np.float64)
211    pitch = np.asanyarray(pitch, dtype=float)
212
213    # find the minimum value of points for origin
214    origin = points.min(axis=0)
215    # convert points to occupied voxel cells
216    index = ((points - origin) / pitch).round().astype(np.int64)
217
218    # convert voxel indices to a matrix
219    matrix = sparse_to_matrix(index)
220
221    # run marching cubes on the matrix to generate a mesh
222    mesh = matrix_to_marching_cubes(matrix)
223    mesh.vertices += origin
224
225    return mesh
226
227
228def multibox(centers, pitch=1.0, colors=None):
229    """
230    Return a Trimesh object with a box at every center.
231
232    Doesn't do anything nice or fancy.
233
234    Parameters
235    -----------
236    centers : (n, 3) float
237      Center of boxes that are occupied
238    pitch : float
239      The edge length of a voxel
240    colors : (3,) or (4,) or (n,3) or (n, 4) float
241      Color of boxes
242
243    Returns
244    ---------
245    rough : Trimesh
246      Mesh object representing inputs
247    """
248    from .. import primitives
249    from ..base import Trimesh
250
251    # get centers as numpy array
252    centers = np.asanyarray(
253        centers, dtype=np.float64)
254
255    # get a basic box
256    b = primitives.Box()
257    # apply the pitch
258    b.apply_scale(float(pitch))
259    # tile into one box vertex per center
260    v = np.tile(
261        centers,
262        (1, len(b.vertices))).reshape((-1, 3))
263    # offset to centers
264    v += np.tile(b.vertices, (len(centers), 1))
265
266    f = np.tile(b.faces, (len(centers), 1))
267    f += np.tile(
268        np.arange(len(centers)) * len(b.vertices),
269        (len(b.faces), 1)).T.reshape((-1, 1))
270
271    face_colors = None
272    if colors is not None:
273        colors = np.asarray(colors)
274        if colors.ndim == 1:
275            colors = colors[None].repeat(len(centers), axis=0)
276        if colors.ndim == 2 and len(colors) == len(centers):
277            face_colors = colors.repeat(12, axis=0)
278
279    mesh = Trimesh(vertices=v,
280                   faces=f,
281                   face_colors=face_colors)
282
283    return mesh
284
285
286def boolean_sparse(a, b, operation=np.logical_and):
287    """
288    Find common rows between two arrays very quickly
289    using 3D boolean sparse matrices.
290
291    Parameters
292    -----------
293    a: (n, d)  int, coordinates in space
294    b: (m, d)  int, coordinates in space
295    operation: numpy operation function, ie:
296                  np.logical_and
297                  np.logical_or
298
299    Returns
300    -----------
301    coords: (q, d) int, coordinates in space
302    """
303    # 3D sparse arrays, using wrapped scipy.sparse
304    # pip install sparse
305    import sparse
306
307    # find the bounding box of both arrays
308    extrema = np.array([a.min(axis=0),
309                        a.max(axis=0),
310                        b.min(axis=0),
311                        b.max(axis=0)])
312    origin = extrema.min(axis=0) - 1
313    size = tuple(extrema.ptp(axis=0) + 2)
314
315    # put nearby voxel arrays into same shape sparse array
316    sp_a = sparse.COO((a - origin).T,
317                      data=np.ones(len(a), dtype=np.bool),
318                      shape=size)
319    sp_b = sparse.COO((b - origin).T,
320                      data=np.ones(len(b), dtype=np.bool),
321                      shape=size)
322
323    # apply the logical operation
324    # get a sparse matrix out
325    applied = operation(sp_a, sp_b)
326    # reconstruct the original coordinates
327    coords = np.column_stack(applied.coords) + origin
328
329    return coords
330
331
332def strip_array(data):
333    shape = data.shape
334    ndims = len(shape)
335    padding = []
336    slices = []
337    for dim, size in enumerate(shape):
338        axis = tuple(range(dim)) + tuple(range(dim + 1, ndims))
339        filled = np.any(data, axis=axis)
340        indices, = np.nonzero(filled)
341        pad_left = indices[0]
342        pad_right = indices[-1]
343        padding.append([pad_left, pad_right])
344        slices.append(slice(pad_left, pad_right))
345    return data[tuple(slices)], np.array(padding, int)
346
347
348def indices_to_points(indices, pitch=None, origin=None):
349    """
350    Convert indices of an (n,m,p) matrix into a set of voxel center points.
351
352    Parameters
353    ----------
354    indices: (q, 3) int, index of voxel matrix (n,m,p)
355    pitch: float, what pitch was the voxel matrix computed with
356    origin: (3,) float, what is the origin of the voxel matrix
357
358    Returns
359    ----------
360    points: (q, 3) float, list of points
361    """
362    indices = np.asanyarray(indices)
363    if indices.shape[1:] != (3,):
364        from IPython import embed
365        embed()
366        raise ValueError('shape of indices must be (q, 3)')
367
368    points = np.array(indices, dtype=np.float64)
369    if pitch is not None:
370        points *= float(pitch)
371    if origin is not None:
372        origin = np.asanyarray(origin)
373        if origin.shape != (3,):
374            raise ValueError('shape of origin must be (3,)')
375        points += origin
376
377    return points
378
379
380def matrix_to_points(matrix, pitch=None, origin=None):
381    """
382    Convert an (n,m,p) matrix into a set of points for each voxel center.
383
384    Parameters
385    -----------
386    matrix: (n,m,p) bool, voxel matrix
387    pitch: float, what pitch was the voxel matrix computed with
388    origin: (3,) float, what is the origin of the voxel matrix
389
390    Returns
391    ----------
392    points: (q, 3) list of points
393    """
394    indices = np.column_stack(np.nonzero(matrix))
395    points = indices_to_points(indices=indices,
396                               pitch=pitch,
397                               origin=origin)
398    return points
399
400
401def points_to_indices(points, pitch=None, origin=None):
402    """
403    Convert center points of an (n,m,p) matrix into its indices.
404
405    Parameters
406    ----------
407    points : (q, 3) float
408      Center points of voxel matrix (n,m,p)
409    pitch : float
410      What pitch was the voxel matrix computed with
411    origin : (3,) float
412      What is the origin of the voxel matrix
413
414    Returns
415    ----------
416    indices : (q, 3) int
417      List of indices
418    """
419    points = np.array(points, dtype=np.float64)
420    if points.shape != (points.shape[0], 3):
421        raise ValueError('shape of points must be (q, 3)')
422
423    if origin is not None:
424        origin = np.asanyarray(origin)
425        if origin.shape != (3,):
426            raise ValueError('shape of origin must be (3,)')
427        points -= origin
428    if pitch is not None:
429        points /= pitch
430
431    origin = np.asanyarray(origin, dtype=np.float64)
432    pitch = float(pitch)
433
434    indices = np.round(points).astype(int)
435    return indices
436