skeleton postprocessing - edits to code style and documentation

Author: GOESTERN-1107078

examples/skeleton/instance_segmentation.py

@@ -17,37 +17,35 @@
 def parse_args():
     parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
     parser.add_argument("--mask_path", required=True,
-                        help="Binary mask to skeletonize; an .h5/.hdf5 or MRC volume.")
+                        help="Binary mask to skeletonize; an .h5 or MRC volume.")
     parser.add_argument("--seg_key", default=None,
-                        help="HDF5 dataset holding the mask; required for .h5/.hdf5 inputs.")
+                        help="HDF5 dataset holding the mask; required for .h5 inputs.")
     parser.add_argument("--output_dir", default=".",
                         help="Directory for the skeleton and instance .npz files.")
-    parser.add_argument("--tag", default="",
-                        help="Suffix appended to the output filenames.")
     parser.add_argument("--pixel_size", type=float, default=10.0,
                         help="Isotropic voxel spacing for teasar; graph coordinates use these units.")
-    parser.add_argument("--scale", type=float, default=0.0,
+    parser.add_argument("--teasar_scale", type=float, default=0.0,
                         help="teasar invalidation-radius scale.")
-    parser.add_argument("--constant", type=float, default=70.0,
+    parser.add_argument("--teasar_const", type=float, default=70.0,
                         help="teasar invalidation-radius constant.")
     parser.add_argument("--number_of_threads", type=int, default=8,
                         help="Threads for teasar.")
     parser.add_argument("--direction_span", type=int, default=10,
                         help="Nodes walked along each arm to estimate its direction at a junction.")
     parser.add_argument("--min_through_angle", type=float, default=170.0,
-                        help="Min through-pair angle (deg) for a degree-4 crossing to split.")
+                        help="Min through-pair angle (degrees) for a degree-4 crossing to split.")
     parser.add_argument("--min_branch_angle", type=float, default=30.0,
-                        help="Min branch angle (deg) for a degree-3 odd arm to be separated.")
+                        help="Min branch angle (degrees) for a degree-3 odd arm to be separated.")
     parser.add_argument("--tick_length", type=float, default=50.0,
-                        help="Prune dead-end branches shorter than this; 0 disables.")
-    parser.add_argument("--join_radius", type=float, default=50.0,
+                        help="Prune dead-end branches shorter than this distance; 0 disables.")
+    parser.add_argument("--join_dist", type=float, default=50.0,
                         help="Join collinear endpoints across gaps up to this distance; 0 disables.")
     parser.add_argument("--min_join_angle", type=float, default=175.0,
-                        help="Min straightness (deg) for a join; 180 is collinear.")
-    parser.add_argument("--circle_size", type=float, default=40.0,
+                        help="Min straightness (degrees) for a join; 180 is collinear.")
+    parser.add_argument("--circle_size", type=float, default=70.0,
                         help="Instance tube diameter for the --view render, in --pixel_size units.")
     parser.add_argument("--view", action="store_true",
-                        help="Open the pipeline stages in napari.")
+                        help="Open the binary mask and instances in napari.")
     return parser.parse_args()
 
 
@@ -69,43 +67,26 @@ def load_mask(mask_path, seg_key):
         return np.asarray(mrc.data)
 
 
-def view_stages(binary, stages, pixel_size, circle_size):
-    import napari
-
-    radius = (circle_size / 2) / pixel_size
-    viewer = napari.Viewer(ndisplay=3)
-    viewer.add_labels(binary, name="mask", opacity=0.4)
-    for index, (name, vertices, edges, _) in enumerate(stages):
-        labels = connected_components(skeleton_to_graph(vertices, edges))
-        volume = draw_instances(vertices / pixel_size, edges, labels, binary.shape, radius)
-        viewer.add_labels(volume, name=f"instances ({name})", visible=index == len(stages) - 1)
-    napari.run()
-
-
 def main():
     args = parse_args()
     mask = load_mask(args.mask_path, args.seg_key)
     binary = (mask > 0).astype(np.uint8)
 
     spacing = (args.pixel_size,) * 3
     raw_vertices, raw_edges, raw_radii = teasar(
-        binary, spacing=spacing, scale=args.scale, constant=args.constant,
+        binary, spacing=spacing, scale=args.teasar_scale, constant=args.teasar_const,
         number_of_threads=args.number_of_threads,
     )
 
-    stages = []
     vertices, edges, radii = clean_graph(
         raw_vertices, raw_edges, radii=raw_radii,
         direction_span=args.direction_span, min_through_angle=args.min_through_angle,
         min_branch_angle=args.min_branch_angle, tick_length=args.tick_length,
-        join_radius=args.join_radius, min_join_angle=args.min_join_angle,
-        save_intermediates=stages,
+        join_dist=args.join_dist, min_join_angle=args.min_join_angle,
     )
     labels = connected_components(skeleton_to_graph(vertices, edges))
 
     name = Path(args.mask_path).stem
-    if args.tag:
-        name = f"{name}_{args.tag}"
     output_dir = Path(args.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -120,7 +101,14 @@ def main():
     print(f"{name}: {len(np.unique(labels))} instances -> {output_dir}")
 
     if args.view:
-        view_stages(binary, stages, args.pixel_size, args.circle_size)
+        import napari
+
+        radius = (args.circle_size / 2) / args.pixel_size
+        volume = draw_instances(vertices / args.pixel_size, edges, labels, binary.shape, radius)
+        viewer = napari.Viewer(ndisplay=3)
+        viewer.add_labels(binary, name="mask", opacity=0.4)
+        viewer.add_labels(volume, name="instances")
+        napari.run()
 
 
 if __name__ == "__main__":

src/bioimage_cpp/skeleton/postprocessing.py

@@ -3,7 +3,7 @@
 Resolve junction nodes so each filament becomes its own connected component.
 Degree-3 and degree-4 nodes are split or pruned based on the angles between their
 incident edges: the straightest pair is kept as the through-going filament and
-the remaining arm(s) are either separated, or for short dead-ends (spurs), removed.
+the remaining arm(s) are either separated, or for short dead-ends (spurs), pruned.
 """
 
 from __future__ import annotations
@@ -13,6 +13,7 @@
 import numpy as np
 
 from ..graph import connected_components
+from ..utils import UnionFind
 from ._graph import skeleton_to_graph
 
 
@@ -33,6 +34,17 @@ def _pair_angle(di, dj):
     return np.degrees(np.arccos(np.clip(di @ dj, -1.0, 1.0)))
 
 
+def _compact(vertices, edges, radii):
+    used = np.zeros(len(vertices), dtype=bool)
+    if len(edges):
+        used[edges.reshape(-1)] = True
+    remap = np.full(len(vertices), -1, dtype=np.int64)
+    remap[used] = np.arange(int(used.sum()))
+    if radii is not None:
+        radii = radii[used]
+    return vertices[used], remap[edges], radii
+
+
 def split_degree3(v, graph, vertices, direction_span=1, min_branch_angle=30.0):
     adj = np.asarray(graph.node_adjacency(int(v)))
     if adj.shape[0] != 3:
@@ -78,7 +90,7 @@ def clean_graph(
     min_through_angle: float = 160.0,
     min_branch_angle: float = 30.0,
     tick_length: float = 0.0,
-    join_radius: float = 0.0,
+    join_dist: float = 0.0,
     min_join_angle: float = 175.0,
     save_intermediates: list | None = None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
@@ -92,8 +104,8 @@ def clean_graph(
        from the through pair by at least ``min_branch_angle``.
     3. Split each degree-4 crossing, separating its two through pairs when they
        are collinear to within ``min_through_angle``.
-    4. If ``join_radius > 0``, reconnect collinear endpoints across gaps less
-       than this distance via :func:`join_close_components`.
+    4. If ``join_dist > 0``, join collinear endpoints across gaps up to
+       this distance via :func:`join_close_components`.
 
     Parameters
     ----------
@@ -112,8 +124,8 @@ def clean_graph(
         through pair for the odd arm to be separated.
     tick_length:
         If > 0, prune dead-end branches shorter than this (physical) distance.
-    join_radius:
-        If > 0, reconnect collinear endpoints across gaps up to this (physical)
+    join_dist:
+        If > 0, join collinear endpoints across gaps up to this (physical)
         distance.
     min_join_angle:
         Minimum straightness (degrees) required for a join; 180 is collinear.
@@ -175,20 +187,12 @@ def _snapshot(name):
         keep[list(prune_edges)] = False
         edges = edges[keep]
 
-    used = np.zeros(len(vertices), dtype=bool)
-    if len(edges):
-        used[edges.reshape(-1)] = True
-    remap = np.full(len(vertices), -1, dtype=np.int64)
-    remap[used] = np.arange(int(used.sum()))
-    vertices = vertices[used]
-    if radii is not None:
-        radii = radii[used]
-    edges = remap[edges]
+    vertices, edges, radii = _compact(vertices, edges, radii)
     _snapshot("split")
 
-    if join_radius > 0:
+    if join_dist > 0:
         vertices, edges, radii = join_close_components(
-            vertices, edges, join_radius,
+            vertices, edges, join_dist,
             min_join_angle=min_join_angle, direction_span=direction_span, radii=radii,
         )
     _snapshot("join")
@@ -208,15 +212,15 @@ def _adjacency(num_nodes, edges):
 
 
 def remove_ticks(vertices, edges, tick_length, radii=None):
-    """Remove short dead-end branches ("ticks") from a skeleton graph.
+    """Prune short dead-end branches ("ticks") from a skeleton graph.
 
     Ports kimimaro's `remove_ticks`. A distance graph is built over the critical
     points (terminals, degree 1; branch points, degree >= 3), whose superedges
     are the paths between them weighted by physical length. The shortest terminal
-    branch below ``tick_length`` is removed repeatedly; when a branch point drops
+    branch below ``tick_length`` is pruned repeatedly; when a branch point drops
     to degree 2 its two superedges are fused into one, so a real filament end is
     re-measured rather than clipped. Standalone paths (both ends terminal) are
-    never removed.
+    never pruned.
 
     Parameters
     ----------
@@ -225,7 +229,7 @@ def remove_ticks(vertices, edges, tick_length, radii=None):
     edges:
         Integer array with shape ``(E, 2)`` indexing ``vertices``.
     tick_length:
-        Maximum branch length (physical) that may be culled.
+        Maximum branch length (physical) that may be pruned.
     radii:
         Optional per-vertex radii, carried through the same remapping.
 
@@ -305,15 +309,7 @@ def remove_ticks(vertices, edges, tick_length, radii=None):
     else:
         edges = edges.copy()
 
-    used = np.zeros(num_nodes, dtype=bool)
-    if len(edges):
-        used[edges.reshape(-1)] = True
-    remap = np.full(num_nodes, -1, dtype=np.int64)
-    remap[used] = np.arange(int(used.sum()))
-    vertices = vertices[used]
-    if radii is not None:
-        radii = radii[used]
-    edges = remap[edges]
+    vertices, edges, radii = _compact(vertices, edges, radii)
     return vertices, edges, radii
 
 
@@ -333,23 +329,23 @@ def _endpoint_tangent(endpoint, indptr, dst, degrees, vertices, span):
     return direction / norm if norm > 0 else None
 
 
-def join_close_components(vertices, edges, radius, *, min_join_angle=175.0,
+def join_close_components(vertices, edges, dist, *, min_join_angle=175.0,
                           direction_span=5, radii=None):
     """Reconnect fragmented filaments by joining collinear endpoints across gaps.
 
     Endpoints (degree = 1) of different connected components are joined with a
-    new edge when they are within ``radius`` and the two fragments are nearly
+    new edge when they are within ``dist`` and the two fragments are nearly
     collinear through the gap: the outward tangent at each endpoint must point
     along the gap to within ``180 - min_join_angle`` degrees, so a straight
     continuation reads ~180 (``min_join_angle``). Joins are made shortest-first
-    with a union-find so each pair of components is linked at most once and each
+    with a union-find so each pair of components is joined at most once and each
     endpoint is used once.
 
     Parameters
     ----------
     vertices, edges:
         Skeleton graph; ``edges`` indexes ``vertices``.
-    radius:
+    dist:
         Maximum gap (physical) across which endpoints may be joined.
     min_join_angle:
         Minimum straightness (degrees) of the joined path; 180 is perfectly
@@ -381,7 +377,7 @@ def join_close_components(vertices, edges, radius, *, min_join_angle=175.0,
     tol = np.deg2rad(180.0 - min_join_angle)
     tree = cKDTree(vertices[endpoints])
     candidates = []
-    for ia, ib in tree.query_pairs(radius):
+    for ia, ib in tree.query_pairs(dist):
         a, b = int(endpoints[ia]), int(endpoints[ib])
         if labels[a] == labels[b]:
             continue
@@ -400,25 +396,15 @@ def join_close_components(vertices, edges, radius, *, min_join_angle=175.0,
             candidates.append((dist, a, b))
 
     candidates.sort()
-    parent = {}
-
-    def find(x):
-        parent.setdefault(x, x)
-        root = x
-        while parent[root] != root:
-            root = parent[root]
-        while parent[x] != root:
-            parent[x], x = root, parent[x]
-        return root
+    uf = UnionFind(int(labels.max()) + 1)
 
     used, new_edges = set(), []
     for _, a, b in candidates:
         if a in used or b in used:
             continue
-        ra, rb = find(int(labels[a])), find(int(labels[b]))
-        if ra == rb:
+        if uf.find(int(labels[a])) == uf.find(int(labels[b])):
             continue
-        parent[ra] = rb
+        uf.merge(int(labels[a]), int(labels[b]))
         new_edges.append([a, b])
         used.add(a)
         used.add(b)
@@ -430,13 +416,6 @@ def find(x):
     return vertices, edges, radii
 
 
-def _line_voxels(start, stop):
-    delta = stop - start
-    steps = int(np.abs(delta).max()) + 1
-    t = np.linspace(0.0, 1.0, steps)
-    return np.rint(start[None, :] + t[:, None] * delta[None, :]).astype(np.int64)
-
-
 def draw_instances(vertices, edges, labels, shape, radius=1):
     """Rasterize a labeled skeleton graph into a dense instance volume.
 
@@ -480,7 +459,10 @@ def draw_instances(vertices, edges, labels, shape, radius=1):
     vi = np.rint(vertices).astype(np.int64)
     centers, center_labels = [], []
     for a, b in edges:
-        pts = _line_voxels(vi[a], vi[b])
+        delta = vi[b] - vi[a]
+        steps = int(np.abs(delta).max()) + 1
+        t = np.linspace(0.0, 1.0, steps)
+        pts = np.rint(vi[a][None, :] + t[:, None] * delta[None, :]).astype(np.int64)
         centers.append(pts)
         center_labels.append(np.full(len(pts), int(labels[a]) + 1))
     centers = np.concatenate(centers)