fix: seg2mesh writes closed x/y/z meshes, and mesh2seg fills them back exactly (v2.0)

src/copick_utils/cli/seg2mesh.py

@@ -129,6 +129,7 @@ def seg2mesh(
         workers=workers,
         level=level,
         step_size=step_size,
+        individual_meshes=individual_meshes,
     )
 
     successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0)

src/copick_utils/converters/mesh_from_segmentation.py

@@ -21,31 +21,80 @@ def volume_to_mesh(volume: np.ndarray, voxel_spacing: float, level: float = 0.5,
     """
     Convert a binary volume to a mesh using marching cubes.
 
+    Vertices are (x, y, z) in physical units, with voxel index i at i * voxel_spacing, like copick's other meshes
+    and picks. The volume is padded with a voxel of background first, so the surface is closed where the
+    segmentation touches the edge of the volume.
+
     Args:
         volume: Binary volume array with shape (z, y, x)
         voxel_spacing: Spacing between voxels in physical units
         level: Isosurface level for marching cubes
         step_size: Step size for marching cubes (higher = coarser mesh)
 
     Returns:
-        Trimesh object representing the mesh
+        Trimesh object representing the mesh (outward-facing triangles)
     """
     from skimage import measure
 
     # Generate mesh using marching cubes
     verts, faces, normals, values = measure.marching_cubes(
-        volume.astype(float),
+        np.pad(volume.astype(float), 1),
         level=level,
         step_size=step_size,
         spacing=(voxel_spacing, voxel_spacing, voxel_spacing),
     )
 
-    # Create trimesh object
-    mesh = tm.Trimesh(vertices=verts, faces=faces, vertex_normals=normals)
+    # marching_cubes works in the array's (z, y, x) order: reverse to (x, y, z) and undo the padding. Reversing the
+    # axes mirrors the mesh, which also turns its triangles to face outward, so the faces stay as they are.
+    mesh = tm.Trimesh(vertices=verts[:, ::-1] - voxel_spacing, faces=faces, vertex_normals=normals[:, ::-1])
 
     return mesh
 
 
+def label_meshes(
+    volume: np.ndarray,
+    voxel_spacing: float,
+    level: float = 0.5,
+    step_size: int = 1,
+) -> Dict[int, tm.Trimesh]:
+    """
+    Convert each non-zero label of a volume to its own mesh (see :func:`volume_to_mesh`).
+
+    Each label is meshed from its own mask, cropped to the label's bounding box.
+
+    Args:
+        volume: Label volume with shape (z, y, x)
+        voxel_spacing: Spacing between voxels in physical units
+        level: Isosurface level for marching cubes (of each label's binary mask)
+        step_size: Step size for marching cubes
+
+    Returns:
+        Meshes by label value, in increasing label order
+    """
+    from scipy import ndimage
+
+    if np.issubdtype(volume.dtype, np.integer) and volume.size and volume.min() >= 0:
+        labels = volume
+        values = None
+    else:
+        # Map the values to consecutive indices, 0 staying background.
+        values, labels = np.unique(volume, return_inverse=True)
+        labels = labels.reshape(volume.shape)
+        if values[0] != 0:
+            values = np.concatenate([[0], values])
+            labels = labels + 1
+
+    meshes = {}
+    for index, box in enumerate(ndimage.find_objects(labels), start=1):
+        if box is None:
+            continue
+        mesh = volume_to_mesh(labels[box] == index, voxel_spacing, level, step_size)
+        # The crop starts at (z, y, x) = box starts: move the mesh there.
+        mesh.apply_translation(np.array([s.start for s in box][::-1]) * voxel_spacing)
+        meshes[int(index if values is None else values[index])] = mesh
+    return meshes
+
+
 def mesh_from_segmentation(
     segmentation: "CopickSegmentation",
     run: "CopickRun",
@@ -55,6 +104,7 @@ def mesh_from_segmentation(
     level: float = 0.5,
     step_size: int = 1,
     voxel_spacing: Optional[float] = None,
+    individual_meshes: bool = False,
     **kwargs,
 ) -> Optional[Tuple["CopickMesh", Dict[str, int]]]:
     """
@@ -64,14 +114,16 @@ def mesh_from_segmentation(
         segmentation: CopickSegmentation object to convert
         run: CopickRun object
         object_name: Name for the output mesh object
-        session_id: Session ID for the output mesh
+        session_id: Session ID for the output mesh; with ``individual_meshes``, a template whose ``{instance_id}``
+            is replaced by each label (without it, ``-{label}`` is appended)
         user_id: User ID for the output mesh
         level: Isosurface level for marching cubes
         step_size: Step size for marching cubes
+        individual_meshes: Store one mesh per non-zero label instead of one mesh of the whole segmentation
 
     Returns:
-        Tuple of (CopickMesh object, stats dict) or None if creation failed.
-        Stats dict contains 'vertices_created' and 'faces_created'.
+        Tuple of (CopickMesh object, stats dict) or None if creation failed. With ``individual_meshes``, the first
+        label's mesh and the totals. Stats dict contains 'vertices_created' and 'faces_created'.
     """
     try:
         # Load the volume from the segmentation
@@ -84,6 +136,32 @@ def mesh_from_segmentation(
         # Get voxel spacing from segmentation
         voxel_spacing = segmentation.voxel_size
 
+        if individual_meshes:
+            meshes = label_meshes(volume, voxel_spacing, level, step_size)
+            if not meshes:
+                logger.error("No labels to mesh")
+                return None
+            first = None
+            stats = {"vertices_created": 0, "faces_created": 0}
+            for label, mesh in meshes.items():
+                if "{instance_id}" in session_id:
+                    label_session_id = session_id.replace("{instance_id}", str(label))
+                else:
+                    label_session_id = f"{session_id}-{label}"
+                copick_mesh, mesh_stats = store_mesh_with_stats(
+                    run,
+                    mesh,
+                    object_name,
+                    label_session_id,
+                    user_id,
+                    "mesh",
+                )
+                if first is None:
+                    first = copick_mesh
+                for key in stats:
+                    stats[key] += mesh_stats[key]
+            return first, stats
+
         # Convert volume to mesh
         mesh = volume_to_mesh(volume, voxel_spacing, level, step_size)
 

src/copick_utils/converters/segmentation_from_mesh.py

@@ -41,62 +41,93 @@ def ensure_mesh(trimesh_object) -> Optional[tm.Trimesh]:
         raise ValueError("Input must be a Trimesh or Scene object")
 
 
+#: Ray origins sit this far (in voxels) off the voxel centres, across the ray. Meshes made by marching cubes have
+#: their vertices and edges on the voxel grid, where rays through the voxel centres would hit edges and vertices
+#: exactly and count the crossing twice or not at all. Rounding back to voxel indices removes the offset; a voxel
+#: centre within about this distance of a sloped surface may fall on either side of it.
+_RAY_OFFSET = (0.01 * np.sqrt(2), 0.01 * np.sqrt(3))
+
+#: A crossing this close (in voxels) to a voxel centre counts as before it, so a face through the centres (e.g. a
+#: box at the volume's edge) includes them where the ray enters and excludes them where it leaves.
+_CENTRE_TOLERANCE = 1e-6
+
+
+def _crossing_voxels(loc: np.ndarray, voxel_spacing: float, axis: int) -> np.ndarray:
+    """Voxel indices of ray crossings: the ray's own voxel across it, and along it the first voxel centre after the
+    crossing, so a voxel is inside when its centre is. Crossings before the volume move to its first voxel."""
+    pos = loc / voxel_spacing
+    idx = np.round(pos).astype("int")
+    idx[:, axis] = np.maximum(np.ceil(pos[:, axis] - _CENTRE_TOLERANCE).astype("int"), 0)
+    return idx
+
+
+def _fill_from_crossings(
+    int_loc: np.ndarray,
+    voxel_dims: Tuple[int, int, int],
+    axis: int,
+) -> np.ndarray:
+    """Inside/outside of every voxel along rays on ``axis``, from the voxels of their crossings: a voxel is inside
+    after an odd number of crossings. Crossings in the same voxel cancel in pairs. Returns a (z, y, x) volume."""
+    toggles = np.zeros((voxel_dims[2], voxel_dims[1], voxel_dims[0]), dtype=np.uint8)
+    zyx = int_loc[int_loc[:, axis] < voxel_dims[axis]][:, ::-1]
+    np.add.at(toggles, (zyx[:, 0], zyx[:, 1], zyx[:, 2]), 1)
+    toggles &= 1
+    return np.bitwise_xor.accumulate(toggles, axis=2 - axis).astype(bool)
+
+
+def _ray_start(mesh: tm.Trimesh, voxel_spacing: float, axis: int) -> float:
+    """Where rays along an axis start (in voxels): before the volume and before the mesh."""
+    return min(-1.0, np.floor(mesh.bounds[0, axis] / voxel_spacing) - 1)
+
+
 def _onesmask_z(mesh: tm.Trimesh, voxel_dims: Tuple[int, int, int], voxel_spacing: float) -> np.ndarray:
     """Create mask by ray casting in Z direction."""
     intersector = RayMeshIntersector(mesh)
 
-    # Create a grid of rays in XY plane, shooting in Z direction
+    # Create a grid of rays in XY plane (off the voxel centres, see _RAY_OFFSET), shooting in Z direction
     grid_x, grid_y = np.mgrid[0 : voxel_dims[0], 0 : voxel_dims[1]]
-    ray_grid = np.vstack([grid_x.ravel(), grid_y.ravel(), -np.ones((grid_x.size,))]).T * voxel_spacing
+    ray_grid = (
+        np.vstack(
+            [
+                grid_x.ravel() + _RAY_OFFSET[0],
+                grid_y.ravel() + _RAY_OFFSET[1],
+                np.full(grid_x.size, _ray_start(mesh, voxel_spacing, 2)),
+            ],
+        ).T
+        * voxel_spacing
+    )
     ray_dir = np.zeros((ray_grid.shape[0], 3))
     ray_dir[:, 2] = 1
 
     loc, _, _ = intersector.intersects_location(ray_grid, ray_dir)
 
-    # Convert to voxel coordinates and sort by z
-    int_loc = np.round(loc / voxel_spacing).astype("int")
-    sort_idx = int_loc[:, 2].argsort()
-    int_loc = int_loc[sort_idx, :]
-
     # Build volume by tracking crossings
-    img = np.zeros((voxel_dims[1], voxel_dims[0]), dtype="bool")
-    vol = np.zeros((voxel_dims[2], voxel_dims[1], voxel_dims[0]), dtype="bool")
-
-    for z in range(voxel_dims[2]):
-        idx = int_loc[:, 2] == z
-        img[int_loc[idx, 1], int_loc[idx, 0]] = np.logical_not(img[int_loc[idx, 1], int_loc[idx, 0]])
-        vol[z, :, :] = img
-
-    return vol
+    return _fill_from_crossings(_crossing_voxels(loc, voxel_spacing, 2), voxel_dims, 2)
 
 
 def _onesmask_x(mesh: tm.Trimesh, voxel_dims: Tuple[int, int, int], voxel_spacing: float) -> np.ndarray:
     """Create mask by ray casting in X direction."""
     intersector = RayMeshIntersector(mesh)
 
-    # Create a grid of rays in YZ plane, shooting in X direction
+    # Create a grid of rays in YZ plane (off the voxel centres, see _RAY_OFFSET), shooting in X direction
     grid_y, grid_z = np.mgrid[0 : voxel_dims[1], 0 : voxel_dims[2]]
-    ray_grid = np.vstack([-np.ones((grid_y.size,)), grid_y.ravel(), grid_z.ravel()]).T * voxel_spacing
+    ray_grid = (
+        np.vstack(
+            [
+                np.full(grid_y.size, _ray_start(mesh, voxel_spacing, 0)),
+                grid_y.ravel() + _RAY_OFFSET[0],
+                grid_z.ravel() + _RAY_OFFSET[1],
+            ],
+        ).T
+        * voxel_spacing
+    )
     ray_dir = np.zeros((ray_grid.shape[0], 3))
     ray_dir[:, 0] = 1
 
     loc, _, _ = intersector.intersects_location(ray_grid, ray_dir)
 
-    # Convert to voxel coordinates and sort by x
-    int_loc = np.round(loc / voxel_spacing).astype("int")
-    sort_idx = int_loc[:, 0].argsort()
-    int_loc = int_loc[sort_idx, :]
-
     # Build volume by tracking crossings
-    img = np.zeros((voxel_dims[2], voxel_dims[1]), dtype="bool")
-    vol = np.zeros((voxel_dims[2], voxel_dims[1], voxel_dims[0]), dtype="bool")
-
-    for x in range(voxel_dims[0]):
-        idx = int_loc[:, 0] == x
-        img[int_loc[idx, 2], int_loc[idx, 1]] = np.logical_not(img[int_loc[idx, 2], int_loc[idx, 1]])
-        vol[:, :, x] = img
-
-    return vol
+    return _fill_from_crossings(_crossing_voxels(loc, voxel_spacing, 0), voxel_dims, 0)
 
 
 def mesh_to_volume(mesh: tm.Trimesh, voxel_dims: Tuple[int, int, int], voxel_spacing: float) -> np.ndarray: