Migrate mesh module

Author: constantinpape

elf/mesh/mesh.py

@@ -1,6 +1,6 @@
 from typing import Optional, Tuple
 
-import nifty
+import bioimage_cpp as bic
 import numpy as np
 from skimage.measure import marching_cubes as marching_cubes_impl
 
@@ -52,24 +52,4 @@ def smooth_mesh(
         The vertices after smoothing.
         The normals after smoothing.
     """
-    n_verts = len(verts)
-    g = nifty.graph.undirectedGraph(n_verts)
-
-    edges = np.concatenate([faces[:, :2], faces[:, 1:], faces[:, ::2]], axis=0)
-    g.insertEdges(edges)
-
-    current_verts = verts
-    current_normals = normals
-    new_verts = np.zeros_like(verts, dtype=verts.dtype)
-    new_normals = np.zeros_like(normals, dtype=normals.dtype)
-
-    # Implement this directly in nifty for speed up?
-    for it in range(iterations):
-        for vert in range(n_verts):
-            nbrs = np.array([vert] + [nbr[0] for nbr in g.nodeAdjacency(vert)], dtype="int")
-            new_verts[vert] = np.mean(current_verts[nbrs], axis=0)
-            new_normals[vert] = np.mean(current_normals[nbrs], axis=0)
-        current_verts = new_verts
-        current_normals = new_normals
-
-    return new_verts, new_normals
+    return bic.utils.smooth_mesh(verts, normals, faces, iterations)

elf/mesh/mesh_to_segmentation.py

@@ -3,7 +3,7 @@
 from typing import Optional, Tuple
 
 import numpy as np
-import vigra
+from bioimage_cpp.segmentation import label as cc_label
 from tqdm import tqdm
 
 from elf.parallel import label
@@ -86,7 +86,7 @@ def mesh_to_segmentation(
     seg[coords] = 0
 
     if block_shape is None:
-        seg_out = vigra.analysis.labelVolumeWithBackground(seg) == 2
+        seg_out = cc_label(seg, background=0, connectivity=1) == 2
     else:
         seg_out = np.zeros_like(seg)
         seg_out = label(seg, seg_out, with_background=True, block_shape=block_shape) == 2

test/mesh/test_mesh.py

@@ -22,6 +22,69 @@ def test_marching_cubes_smoothing(self):
             self.assertGreater(len(faces), 0)
             self.assertGreater(len(normals), 0)
 
+    def test_marching_cubes_resolution(self):
+        from elf.mesh import marching_cubes
+        shape = (32,) * 3
+        rng = np.random.default_rng(0)
+        seg = (rng.random(shape) > 0.8).astype("uint32")
+
+        verts_unit, _, _ = marching_cubes(seg, resolution=(1.0, 1.0, 1.0))
+        resolution = (2.0, 1.0, 0.5)
+        verts_scaled, _, _ = marching_cubes(seg, resolution=resolution)
+
+        self.assertEqual(verts_unit.shape, verts_scaled.shape)
+        expected = verts_unit * np.array(resolution)
+        self.assertTrue(np.allclose(verts_scaled, expected))
+
+    def test_smooth_mesh_tetrahedron(self):
+        from elf.mesh.mesh import smooth_mesh
+
+        verts = np.array(
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+            dtype=np.float64,
+        )
+        normals = verts + 0.5
+        faces = np.array(
+            [[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]], dtype=np.int64
+        )
+
+        out_verts, out_normals = smooth_mesh(verts, normals, faces, iterations=1)
+
+        self.assertEqual(out_verts.shape, verts.shape)
+        self.assertEqual(out_normals.shape, normals.shape)
+        self.assertEqual(out_verts.dtype, verts.dtype)
+        self.assertEqual(out_normals.dtype, normals.dtype)
+
+        # Every vertex of a tetrahedron is connected to every other vertex,
+        # so one Jacobi-Laplacian iteration collapses each vertex onto the
+        # mean of all four vertices (the centroid).
+        centroid_verts = verts.mean(axis=0)
+        centroid_normals = normals.mean(axis=0)
+        for i in range(len(verts)):
+            self.assertTrue(np.allclose(out_verts[i], centroid_verts))
+            self.assertTrue(np.allclose(out_normals[i], centroid_normals))
+
+    def test_smooth_mesh_no_op(self):
+        from elf.mesh.mesh import smooth_mesh
+
+        verts = np.array(
+            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+            dtype=np.float64,
+        )
+        normals = verts + 0.5
+        faces = np.array(
+            [[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]], dtype=np.int64
+        )
+
+        out_verts, out_normals = smooth_mesh(verts, normals, faces, iterations=0)
+
+        self.assertEqual(out_verts.shape, verts.shape)
+        self.assertEqual(out_normals.shape, normals.shape)
+        self.assertEqual(out_verts.dtype, verts.dtype)
+        self.assertEqual(out_normals.dtype, normals.dtype)
+        self.assertTrue(np.array_equal(out_verts, verts))
+        self.assertTrue(np.array_equal(out_normals, normals))
+
 
 if __name__ == '__main__':
     unittest.main()

test/mesh/test_mesh_to_segmentation.py

@@ -0,0 +1,58 @@
+import os
+import tempfile
+import unittest
+
+import numpy as np
+
+try:
+    import madcad  # noqa: F401
+except ImportError:
+    madcad = None
+
+
+def _make_sphere(shape, radius):
+    z, y, x = np.indices(shape)
+    center = np.array(shape) // 2
+    dist = np.sqrt(
+        (z - center[0]) ** 2 + (y - center[1]) ** 2 + (x - center[2]) ** 2
+    )
+    return (dist <= radius).astype("uint8")
+
+
+@unittest.skipIf(madcad is None, "Needs madcad")
+class TestMeshToSegmentation(unittest.TestCase):
+    def _roundtrip(self, block_shape):
+        from elf.mesh import marching_cubes
+        from elf.mesh.io import write_obj
+        from elf.mesh.mesh_to_segmentation import mesh_to_segmentation
+
+        shape = (32, 32, 32)
+        seg = _make_sphere(shape, radius=10)
+        verts, faces, _ = marching_cubes(seg)
+
+        with tempfile.NamedTemporaryFile(suffix=".obj", delete=False) as f:
+            obj_path = f.name
+        try:
+            write_obj(obj_path, verts, faces)
+            recovered = mesh_to_segmentation(
+                obj_path, shape=shape, block_shape=block_shape
+            )
+        finally:
+            if os.path.exists(obj_path):
+                os.remove(obj_path)
+
+        # The recovered mask should agree with the original sphere on the
+        # interior up to a thin surface tolerance.
+        diff = recovered.astype(bool) ^ seg.astype(bool)
+        # Number of disagreements should be small relative to the sphere volume.
+        self.assertLess(diff.sum(), seg.sum() * 0.5)
+
+    def test_mesh_to_segmentation_serial(self):
+        self._roundtrip(block_shape=None)
+
+    def test_mesh_to_segmentation_blockwise(self):
+        self._roundtrip(block_shape=(16, 16, 16))
+
+
+if __name__ == "__main__":
+    unittest.main()