perf(multires meshes): vastly faster mesh chunking

igneous/tasks/mesh/multires.py

@@ -109,7 +109,11 @@ def process_mesh(
     return (None, None)
 
   lods = [
-    create_octree_level_from_mesh(lods[lod], chunk_shape, lod, len(lods), grid_origin, mesh_shape)
+    create_octree_level_from_mesh(
+      lods[lod], chunk_shape, 
+      lod, len(lods), 
+      grid_origin, mesh_shape
+    )
     for lod in range(len(lods)) 
   ]
   fragment_positions = [ nodes for submeshes, nodes in lods ]
@@ -132,6 +136,9 @@ def process_mesh(
   mesh_binaries = []
   for lod, submeshes in enumerate(lods):
     for frag_no, submesh in enumerate(submeshes):
+      if submesh.empty():
+        continue
+
       submesh.vertices = to_stored_model_space(
         submesh.vertices, manifest, 
         lod=lod,
@@ -222,7 +229,8 @@ def MultiResShardedMeshMergeTask(
   # important to iterate this way to avoid
   # creating a copy of meshes vs. { ... for in }
   for label in labels:
-    meshes[label] = Mesh.concatenate(*meshes[label])
+    mesh = Mesh.concatenate(*meshes[label], segid=label)
+    meshes[label] = mesh.consolidate()
   del labels
 
   fname, shard = create_mesh_shard(
@@ -324,7 +332,7 @@ def generate_lods(
     )
 
     lods.append(
-      Mesh(*simplifier.getMesh())
+      Mesh(*simplifier.getMesh(), segid=label)
     )
 
   return lods
@@ -493,7 +501,7 @@ def less_msb(x: int, y: int) -> bool:
   return lhs[msd] - rhs[msd]
 
 def determine_mesh_shape_from_lods(
-    lods: list[trimesh.Trimesh],
+    lods: list[Mesh],
   ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.int_]]:
   mesh_starts = [np.min(lod.vertices, axis=0) for lod in lods]
   mesh_ends = [np.max(lod.vertices, axis=0) for lod in lods]
@@ -503,54 +511,17 @@ def determine_mesh_shape_from_lods(
 
   return grid_origin, mesh_shape
 
-def generate_gridded_submeshes(
-  mesh: trimesh.Trimesh,
-  offset: np.ndarray,
-  grid_size: Vec,
-  scale: Vec,
-) -> Iterator[Tuple[trimesh.Trimesh, Tuple[int, int, int]]]:
-  nx, ny, nz = np.eye(3)
-  ox, oy, oz = offset * np.eye(3)
-
-  for x in range(0, grid_size.x):
-    # list(...) required b/c it doesn't like Vec classes
-    mesh_x = trimesh.intersections.slice_mesh_plane(mesh, plane_normal=nx, plane_origin=list(nx*x*scale.x+ox))
-    mesh_x = trimesh.intersections.slice_mesh_plane(mesh_x, plane_normal=-nx, plane_origin=list(nx*(x+1)*scale.x+ox))
-    for y in range(0, grid_size.y):
-      mesh_y = trimesh.intersections.slice_mesh_plane(mesh_x, plane_normal=ny, plane_origin=list(ny*y*scale.y+oy))
-      mesh_y = trimesh.intersections.slice_mesh_plane(mesh_y, plane_normal=-ny, plane_origin=list(ny*(y+1)*scale.y+oy))
-      for z in range(0, grid_size.z):
-        mesh_z = trimesh.intersections.slice_mesh_plane(mesh_y, plane_normal=nz, plane_origin=list(nz*z*scale.z+oz))
-        mesh_z = trimesh.intersections.slice_mesh_plane(mesh_z, plane_normal=-nz, plane_origin=list(nz*(z+1)*scale.z+oz))
-
-        if len(mesh_z.vertices) == 0:
-          continue
-
-        # test for totally degenerate meshes by checking if 
-        # all of two axes match, meaning the mesh must be a 
-        # point or a line.
-        if np.sum([ np.all(mesh_z.vertices[:,i] == mesh_z.vertices[0,i]) for i in range(3) ]) >= 2:
-          continue
-
-        yield mesh_z, (x, y, z)
-
 def retriangulate_mesh(
-    mesh: trimesh.Trimesh,
+    mesh: Mesh,
     offset: np.ndarray,
-    grid_size: Vec,
     scale: Vec,
-  ) -> trimesh.Trimesh:
+  ) -> Mesh:
   """
   Retriangulate the input mesh to avoid any cases where the boundaries of a triangle are split across the boundaries of the submeshes
   """
-  new_mesh = trimesh.Trimesh()
-
-  for submesh, _ in generate_gridded_submeshes(
-      mesh, offset, grid_size, scale
-  ):
-      new_mesh = trimesh.util.concatenate(new_mesh, submesh)
-
-  return new_mesh
+  chunks = zmesh.chunk_mesh(mesh, scale, offset)
+  new_mesh = zmesh.Mesh.concatenate(*chunks.values(), id=mesh.segid)
+  return new_mesh.merge_close_vertices(radius=1e-5)
 
 def create_octree_level_from_mesh(mesh, chunk_shape, lod, num_lods, offset, grid_length):
   """
@@ -559,7 +530,6 @@ def create_octree_level_from_mesh(mesh, chunk_shape, lod, num_lods, offset, grid
 
   This creates (2^lod)^3 submeshes.
   """
-  mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
   scale = Vec(*(np.array(chunk_shape) * (2**lod)))
   grid_size = Vec(*(np.ceil(grid_length / scale)), dtype=int)
 
@@ -568,26 +538,21 @@ def create_octree_level_from_mesh(mesh, chunk_shape, lod, num_lods, offset, grid
   # at the higher level of the octree
   if lod > 0:
       upper_grid_scale = Vec(*(np.array(chunk_shape) * (2 ** (lod - 1))))
-      upper_grid_shape = Vec(*np.ceil(grid_length / upper_grid_scale), dtype=int)
-      mesh = retriangulate_mesh(mesh, offset, upper_grid_shape, upper_grid_scale)
+      mesh = retriangulate_mesh(mesh, offset, upper_grid_scale)
 
   if lod == num_lods - 1:
       return ([Mesh(mesh.vertices, mesh.faces)], ((0, 0, 0),))
 
-  submeshes = []
-  nodes = []
-  for submesh, node in generate_gridded_submeshes(
-      mesh, offset, grid_size, scale
-  ):
-    submeshes.append(submesh)
-    nodes.append(node)
+  grid = zmesh.chunk_mesh(mesh, scale, offset)
 
   # Sort in Z-curve order
-  submeshes, nodes = zip(
-    *sorted(zip(submeshes, nodes),
-    key=functools.cmp_to_key(lambda x, y: cmp_zorder(x[1], y[1])))
+  nodes, submeshes = zip(
+    *sorted(
+      grid.items(),
+      key=functools.cmp_to_key(lambda x, y: cmp_zorder(x[0], y[0]))
+    )
   )
-  # convert back from trimesh to CV Mesh class
-  submeshes = [ Mesh(m.vertices, m.faces) for m in submeshes ]
+  # convert back from zmesh.Mesh to CV Mesh class
+  submeshes = [ Mesh(m.vertices, m.faces, segid=mesh.segid) for m in submeshes ]
 
   return (submeshes, nodes)

requirements.txt

@@ -25,4 +25,4 @@ task-queue>=2.4.0
 tqdm
 trimesh[easy]
 xs3d>=1.11.0
-zmesh>=1.4,<2.0
+zmesh>=1.9.0,<2.0