Avoid tunnels from the mitochondria closing ("holding hands")

synapse_net/inference/mitochondria.py

@@ -47,7 +47,12 @@ def _run_segmentation(
         print("Compute watershed in", time.time() - t0, "s")
 
     seg = apply_size_filter(seg, min_size, verbose=verbose, block_shape=block_shape)
-    seg = _postprocess_seg_3d(seg, area_threshold=area_threshold)
+    t0 = time.time()
+    seg = _postprocess_seg_3d(
+        seg, area_threshold=area_threshold, iterations=2, iterations_3d=3, fill_holes=True, remove_tunnels=True
+    )
+    if verbose:
+        print("Post-process segmentation in", time.time() - t0, "s")
     return seg
 
 

synapse_net/inference/util.py

@@ -19,8 +19,9 @@
 
 from elf.io import open_file
 from numpy.typing import ArrayLike
+from scipy import ndimage
 from scipy.ndimage import binary_closing
-from skimage.measure import regionprops
+from skimage.measure import euler_number, regionprops
 from skimage.morphology import remove_small_holes
 from skimage.transform import rescale, resize
 from torch_em.model.unet import Decoder
@@ -691,7 +692,79 @@ def apply_size_filter(
     return segmentation
 
 
-def _postprocess_seg_3d(seg, area_threshold=1000, iterations=4, iterations_3d=8):
+def _fill_holes_per_slice(mask, axes=(0, 1, 2)):
+    """Fill the holes in each 2D slice of the mask, along each of the axes in turn."""
+    for axis in axes:
+        moved = np.moveaxis(mask, axis, 0)
+        mask = np.moveaxis(np.stack([ndimage.binary_fill_holes(plane) for plane in moved]), 0, axis)
+    return mask
+
+
+def _count_tunnels(mask):
+    """Count the tunnels (handles) of a 3D mask, with the object 26-connected and the background 6-connected.
+
+    It follows from the Euler number = components - tunnels + cavities.
+    """
+    mask = np.pad(mask, 1)
+    n_components = ndimage.label(mask, structure=np.ones((3, 3, 3)))[1]
+    n_cavities = ndimage.label(~mask)[1] - 1
+    return int(n_components + n_cavities - euler_number(mask, connectivity=3))
+
+
+def _remove_tunnel_bridges(mask, closed):
+    """Remove the parts that the post-processing added to an object and that give it a tunnel.
+
+    Each connected component of closed & ~mask that touches mask in two or more places is a candidate,
+    because only such a component can close a ring. From the largest to the smallest candidate, each is removed
+    if the object has fewer tunnels without it. This undoes bridges between two parts of a mitochondrion that only
+    come close, e.g. the two ends of a curled one ("holding hands").
+    """
+    full = np.ones((3, 3, 3), dtype=bool)
+    added, _ = ndimage.label(closed & ~mask, structure=full)
+    candidates = []
+    for component_id, bb in enumerate(ndimage.find_objects(added), start=1):
+        bb = tuple(slice(max(s.start - 1, 0), min(s.stop + 1, n)) for s, n in zip(bb, mask.shape))
+        component = added[bb] == component_id
+        contact = ndimage.binary_dilation(component, structure=full) & mask[bb]
+        if ndimage.label(contact, structure=full)[1] >= 2:
+            candidates.append((bb, component))
+    if not candidates:
+        return closed
+
+    closed = closed.copy()
+    n_tunnels = _count_tunnels(closed)
+    for bb, component in sorted(candidates, key=lambda candidate: -candidate[1].sum()):
+        closed[bb][component] = False
+        n_tunnels_without = _count_tunnels(closed)
+        if n_tunnels_without < n_tunnels:
+            n_tunnels = n_tunnels_without
+        else:
+            closed[bb][component] = True
+    return closed
+
+
+def _postprocess_seg_3d(
+    seg, area_threshold=1000, iterations=4, iterations_3d=8, fill_holes=False, remove_tunnels=False
+):
+    """Smooth each object of a 3D segmentation in place.
+
+    Each object is processed within its bounding box: small holes are removed, then it is closed with
+    a 3D cross structure (`iterations` times) and a 3x3 structure in the XY plane (`iterations_3d` times).
+    The closing may overwrite other objects.
+
+    Args:
+        seg: The segmentation.
+        area_threshold: The maximal size of the holes that are removed.
+        iterations: The number of iterations for the closing with the 3D cross structure.
+        iterations_3d: The number of iterations for the closing with the 3x3 structure in the XY plane.
+        fill_holes: Whether to fill the holes in each 2D slice of the closed object along z, y and x.
+            Unlike the closing, the hole filling only writes into the background.
+        remove_tunnels: Whether to remove the parts added by the closing and the hole filling that give the object
+            a tunnel, see `_remove_tunnel_bridges`.
+
+    Returns:
+        The post-processed segmentation.
+    """
     # Structure lement for 2d dilation in 3d.
     structure_element = np.ones((3, 3))  # 3x3 structure for XY plane
     structure_3d = np.zeros((1, 3, 3))  # Only applied in the XY plane
@@ -705,9 +778,15 @@ def _postprocess_seg_3d(seg, area_threshold=1000, iterations=4, iterations_3d=8)
 
         # Fill small holes and apply closing.
         mask = remove_small_holes(mask, area_threshold=area_threshold)
-        mask = np.logical_or(binary_closing(mask, iterations=iterations), mask)
-        mask = np.logical_or(binary_closing(mask, iterations=iterations_3d, structure=structure_3d), mask)
-        seg[bb][mask] = prop.label
+        closed = np.logical_or(binary_closing(mask, iterations=iterations), mask)
+        closed = np.logical_or(binary_closing(closed, iterations=iterations_3d, structure=structure_3d), closed)
+
+        if fill_holes:
+            free = np.isin(seg[bb], [0, prop.label]) | closed
+            closed = closed | (_fill_holes_per_slice(closed) & free)
+        if remove_tunnels:
+            closed = _remove_tunnel_bridges(mask, closed)
+        seg[bb][closed] = prop.label
 
     return seg
 

test/test_mitochondria_postprocessing.py

@@ -0,0 +1,83 @@
+"""Offline tests for the post-processing of the mitochondria segmentation (_postprocess_seg_3d)."""
+import unittest
+
+import numpy as np
+
+MITO_SETTINGS = dict(area_threshold=5000, iterations=2, iterations_3d=3, fill_holes=True, remove_tunnels=True)
+
+
+def _curled_object(gap, length=16, width=3, overhang=4, offset=4):
+    """A rod along six edges of a cube that almost closes into a ring: its end stops `gap` voxels before its start.
+
+    The ring does not lie in a plane, so no hole filling per slice can close it. The rods overhang the corners,
+    so that the gap is inside the bounding box of the object.
+    """
+    shift = offset + overhang
+    shape = (2 * shift + length + width,) * 3
+    seg = np.zeros(shape, dtype="uint32")
+    corners = [(0, 0, 0), (length, 0, 0), (length, length, 0), (length, length, length),
+               (0, length, length), (0, 0, length), (0, 0, 0)]
+    for start, stop in zip(corners[:-1], corners[1:]):
+        lower = [shift + min(a, b) - (overhang if a != b else 0) for a, b in zip(start, stop)]
+        upper = [shift + max(a, b) + width + (overhang if a != b else 0) for a, b in zip(start, stop)]
+        if stop == corners[0]:
+            lower[2] = shift + width + gap
+        seg[lower[0]:upper[0], lower[1]:upper[1], lower[2]:upper[2]] = 1
+    return seg
+
+
+class TestMitochondriaPostprocessing(unittest.TestCase):
+    def test_count_tunnels(self):
+        from synapse_net.inference.util import _count_tunnels
+
+        z, y, x = np.mgrid[:30, :30, :30] - 14.5
+        ball = z ** 2 + y ** 2 + x ** 2 < 100
+        torus = (np.sqrt(y ** 2 + x ** 2) - 9) ** 2 + z ** 2 < 9
+        self.assertEqual(_count_tunnels(ball), 0)
+        self.assertEqual(_count_tunnels(torus), 1)
+        self.assertEqual(_count_tunnels(np.concatenate([torus, torus])), 2)
+
+    def test_curled_object(self):
+        from synapse_net.inference.util import _count_tunnels, _postprocess_seg_3d
+
+        seg = _curled_object(gap=2)
+        self.assertEqual(_count_tunnels(seg == 1), 0)
+
+        # The default closing joins the two ends and gives the object a tunnel.
+        default = _postprocess_seg_3d(seg.copy(), area_threshold=5000)
+        self.assertEqual(_count_tunnels(default == 1), 1)
+
+        # The weaker closing joins them too, but the join is removed again.
+        bridged = _postprocess_seg_3d(seg.copy(), **{**MITO_SETTINGS, "remove_tunnels": False})
+        self.assertEqual(_count_tunnels(bridged == 1), 1)
+        result = _postprocess_seg_3d(seg.copy(), **MITO_SETTINGS)
+        self.assertEqual(_count_tunnels(result == 1), 0)
+
+        # The other added voxels, which fill the corners, are kept.
+        self.assertTrue(np.all(result[seg == 1] == 1))
+        self.assertGreater(np.count_nonzero(result), np.count_nonzero(seg))
+
+    def test_fill_holes(self):
+        from synapse_net.inference.util import _count_tunnels, _postprocess_seg_3d
+
+        # A tube along z with a second object inside it.
+        seg = np.zeros((12, 40, 40), dtype="uint32")
+        seg[1:11, 5:35, 5:35] = 1
+        seg[1:11, 8:32, 8:32] = 0
+        seg[4:8, 18:22, 18:22] = 2
+        lumen = np.zeros(seg.shape, dtype=bool)
+        lumen[1:11, 8:32, 8:32] = True
+        lumen[seg == 2] = False
+
+        result = _postprocess_seg_3d(seg.copy(), **MITO_SETTINGS)
+        self.assertTrue(np.all(result[lumen] == 1))
+        self.assertTrue(np.all(result[seg == 2] == 2))
+        self.assertEqual(_count_tunnels(result == 1), 0)
+
+        # Without the hole filling, the lumen stays open.
+        result = _postprocess_seg_3d(seg.copy(), **{**MITO_SETTINGS, "fill_holes": False})
+        self.assertEqual(_count_tunnels(result == 1), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()