elf.mesh.mesh

 1from typing import Optional, Tuple
 2
 3import nifty
 4import numpy as np
 5from skimage.measure import marching_cubes as marching_cubes_impl
 6
 7
 8def marching_cubes(
 9    obj: np.ndarray,
10    smoothing_iterations: int = 0,
11    resolution: Optional[Tuple[float, float, float]] = None,
12) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
13    """Compute mesh via marching cubes.
14
15    This is a wrapper around the skimage marching cubes implementation that provides
16    additional mesh smoothing.
17
18    Args:
19        obj: Volume containing the object to be meshed.
20        smoothing_iterations: Number of mesh smoothing iterations.
21        resolution: Resolution of the data.
22
23    Returns:
24        The vertices of the mesh.
25        The faces of the mesh.
26        The normals of the mesh.
27    """
28    resolution = (1.0, 1.0, 1.0) if resolution is None else resolution
29    if len(resolution) != 3:
30        raise ValueError(f"Invalid resolution argument: {resolution}")
31    resolution = tuple(resolution)
32
33    verts, faces, normals, _ = marching_cubes_impl(obj, spacing=resolution)
34    if smoothing_iterations > 0:
35        verts, normals = smooth_mesh(verts, normals, faces, smoothing_iterations)
36
37    return verts, faces, normals
38
39
40def smooth_mesh(
41    verts: np.ndarray, normals: np.ndarray, faces: np.ndarray, iterations: int
42) -> Tuple[np.ndarray, np.ndarray]:
43    """Smooth mesh surface via laplacian smoothing.
44
45    Args:
46        verts: The mesh vertices.
47        normals: The mesh normals.
48        faces: The mesh faces.
49        iterations: The number of smoothing iterations.
50
51    Returns:
52        The vertices after smoothing.
53        The normals after smoothing.
54    """
55    n_verts = len(verts)
56    g = nifty.graph.undirectedGraph(n_verts)
57
58    edges = np.concatenate([faces[:, :2], faces[:, 1:], faces[:, ::2]], axis=0)
59    g.insertEdges(edges)
60
61    current_verts = verts
62    current_normals = normals
63    new_verts = np.zeros_like(verts, dtype=verts.dtype)
64    new_normals = np.zeros_like(normals, dtype=normals.dtype)
65
66    # Implement this directly in nifty for speed up?
67    for it in range(iterations):
68        for vert in range(n_verts):
69            nbrs = np.array([vert] + [nbr[0] for nbr in g.nodeAdjacency(vert)], dtype="int")
70            new_verts[vert] = np.mean(current_verts[nbrs], axis=0)
71            new_normals[vert] = np.mean(current_normals[nbrs], axis=0)
72        current_verts = new_verts
73        current_normals = new_normals
74
75    return new_verts, new_normals