diff --git a/examples/skeleton/instance_segmentation.py b/examples/skeleton/instance_segmentation.py new file mode 100644 index 0000000..acd8d66 --- /dev/null +++ b/examples/skeleton/instance_segmentation.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python +"""Split a skeletonized mask into filament instances with bioimage_cpp. + +Skeletonize a binary mask with TEASAR, clean the skeleton graph so each filament +becomes its own connected component, and save the labeled instances. +""" + +import argparse +from pathlib import Path + +import numpy as np + +from bioimage_cpp.graph import connected_components +from bioimage_cpp.skeleton import clean_graph, draw_instances, skeleton_to_graph, teasar + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--mask_path", required=True, + 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 inputs.") + parser.add_argument("--output_dir", default=".", + help="Directory for the skeleton and instance .npz files.") + parser.add_argument("--pixel_size", type=float, default=10.0, + help="Isotropic voxel spacing for teasar; graph coordinates use these units.") + parser.add_argument("--teasar_scale", type=float, default=0.0, + help="teasar invalidation-radius scale.") + 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 (degrees) for a degree-4 crossing to split.") + parser.add_argument("--min_branch_angle", type=float, default=30.0, + 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 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 (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 binary mask and instances in napari.") + return parser.parse_args() + + +def load_mask(mask_path, seg_key): + path = Path(mask_path) + if path.suffix in (".h5", ".hdf5"): + import h5py + + if seg_key is None: + raise ValueError("--seg_key is required for HDF5 inputs") + with h5py.File(path, "r") as f: + if seg_key not in f: + raise KeyError(f"{seg_key} not in {path}") + return np.asarray(f[seg_key][:]) + + import mrcfile + + with mrcfile.mmap(path, mode="r", permissive=True) as mrc: + return np.asarray(mrc.data) + + +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.teasar_scale, constant=args.teasar_const, + number_of_threads=args.number_of_threads, + ) + + 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_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 + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + skeleton = dict(vertices=raw_vertices, edges=raw_edges) + if raw_radii is not None: + skeleton["radii"] = raw_radii + instances = dict(vertices=vertices, edges=edges, labels=labels) + if radii is not None: + instances["radii"] = radii + np.savez_compressed(output_dir / f"{name}_skeleton.npz", **skeleton) + np.savez_compressed(output_dir / f"{name}_instances.npz", **instances) + print(f"{name}: {len(np.unique(labels))} instances -> {output_dir}") + + if args.view: + 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__": + main() diff --git a/src/bioimage_cpp/skeleton/__init__.py b/src/bioimage_cpp/skeleton/__init__.py index 5c80fb8..f8de89f 100644 --- a/src/bioimage_cpp/skeleton/__init__.py +++ b/src/bioimage_cpp/skeleton/__init__.py @@ -10,6 +10,7 @@ from .._validation import strict_index from ..distance._distance import _as_binary_input, _normalize_sampling, _normalize_threads from ._graph import skeleton_to_graph +from .postprocessing import clean_graph, draw_instances, join_close_components, remove_ticks _TEASAR_LABELS_BY_DTYPE = { @@ -205,4 +206,4 @@ def teasar_labels( from . import distributed -__all__ = ["distributed", "skeleton_to_graph", "teasar", "teasar_labels"] +__all__ = ["clean_graph", "distributed", "draw_instances", "join_close_components", "remove_ticks", "skeleton_to_graph", "teasar", "teasar_labels"] diff --git a/src/bioimage_cpp/skeleton/postprocessing.py b/src/bioimage_cpp/skeleton/postprocessing.py new file mode 100644 index 0000000..a95d951 --- /dev/null +++ b/src/bioimage_cpp/skeleton/postprocessing.py @@ -0,0 +1,482 @@ +"""Post-processing for skeleton graphs. + +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), pruned. +""" + +from __future__ import annotations + +from collections import defaultdict + +import numpy as np + +from ..graph import connected_components +from ..utils import UnionFind +from ._graph import skeleton_to_graph + + +def _tangent(v, first_neighbor, graph, vertices, span): + prev, cur = v, int(first_neighbor) + for _ in range(span - 1): + nbrs = np.asarray(graph.node_adjacency(cur))[:, 0] + if len(nbrs) != 2: + break + nxt = int(nbrs[0]) if int(nbrs[0]) != prev else int(nbrs[1]) + prev, cur = cur, nxt + d = vertices[cur] - vertices[v] + n = np.linalg.norm(d) + return d / n if n > 0 else d + + +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: + return None + neighbors, edge_ids = adj[:, 0], adj[:, 1] + dirs = np.stack([_tangent(int(v), n, graph, vertices, direction_span) for n in neighbors]) + best, best_angle = None, -1.0 + for a, b in [(0, 1), (0, 2), (1, 2)]: + ang = _pair_angle(dirs[a], dirs[b]) + if ang > best_angle: + best_angle, best = ang, (a, b) + i, j = best + odd = ({0, 1, 2} - {i, j}).pop() + branch_angle = min(_pair_angle(dirs[odd], dirs[i]), _pair_angle(dirs[odd], dirs[j])) + if branch_angle < min_branch_angle: + return None + return [int(edge_ids[odd])] + + +def split_degree4(v, graph, vertices, direction_span=1, min_through_angle=160.0): + adj = np.asarray(graph.node_adjacency(int(v))) + if adj.shape[0] != 4: + return None + neighbors, edge_ids = adj[:, 0], adj[:, 1] + dirs = np.stack([_tangent(int(v), n, graph, vertices, direction_span) for n in neighbors]) + best, best_score, best_min = None, -np.inf, 0.0 + for (a, b), (c, d) in [((0, 1), (2, 3)), ((0, 2), (1, 3)), ((0, 3), (1, 2))]: + ang1, ang2 = _pair_angle(dirs[a], dirs[b]), _pair_angle(dirs[c], dirs[d]) + if ang1 + ang2 > best_score: + best_score, best, best_min = ang1 + ang2, ((a, b), (c, d)), min(ang1, ang2) + if best_min < min_through_angle: + return None + (_, pair_b) = best + return [int(edge_ids[k]) for k in pair_b] + + +def clean_graph( + vertices: np.ndarray, + edges: np.ndarray, + radii: np.ndarray | None = None, + *, + direction_span: int = 5, + min_through_angle: float = 160.0, + min_branch_angle: float = 30.0, + tick_length: 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]: + """Split skeleton junctions so each filament is its own component. + + Postprocessing steps, in order: + + 1. If ``tick_length > 0``, prune dead-end branches shorter than it via + :func:`remove_ticks`. + 2. Split each degree-3 junction, separating the odd arm when it diverges + 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_dist > 0``, join collinear endpoints across gaps up to + this distance via :func:`join_close_components`. + + Parameters + ---------- + vertices: + Float array with shape ``(V, D)`` of skeleton vertex coordinates. + edges: + Integer array with shape ``(E, 2)`` indexing ``vertices``. + radii: + Optional per-vertex radii, carried through the same remapping. + direction_span: + Number of nodes over which each arm's tangent is measured. + min_through_angle: + Minimum through-pair angle (degrees) for a degree-4 crossing to split. + min_branch_angle: + Minimum angle (degrees) between a degree-3 node's odd arm and its + through pair for the odd arm to be separated. + tick_length: + If > 0, prune dead-end branches shorter than this (physical) distance. + 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. + save_intermediates: + If a list is given, ``(name, vertices, edges, radii)`` snapshots are + appended after each step ("raw", "ticks", "split", "join"). + + Returns + ------- + vertices, edges, radii: + The cleaned arrays, reindexed with unused vertices dropped. ``radii`` is + ``None`` when no input radii were given. + """ + vertices = np.asarray(vertices, dtype=np.float64) + edges = np.asarray(edges, dtype=np.int64).copy() + + def _snapshot(name): + if save_intermediates is not None: + save_intermediates.append(( + name, vertices.copy(), edges.copy(), + None if radii is None else np.asarray(radii).copy(), + )) + + _snapshot("raw") + if tick_length > 0: + vertices, edges, radii = remove_ticks(vertices, edges, tick_length, radii=radii) + edges = edges.copy() + _snapshot("ticks") + graph = skeleton_to_graph(vertices, edges) + degrees = np.bincount(edges.reshape(-1), minlength=len(vertices)) + + splits, prune_edges = [], set() + for v in np.where(degrees == 3)[0]: + ids = split_degree3(v, graph, vertices, direction_span, min_branch_angle) + if ids: + prune_edges.update(ids) + for v in np.where(degrees == 4)[0]: + ids = split_degree4(v, graph, vertices, direction_span, min_through_angle) + if ids: + splits.append((int(v), ids)) + + extra_vertices, extra_radii = [], [] + next_id = len(vertices) + for node, edge_ids in splits: + dup = next_id + next_id += 1 + extra_vertices.append(vertices[node]) + if radii is not None: + extra_radii.append(radii[node]) + for e in edge_ids: + edges[e] = np.where(edges[e] == node, dup, edges[e]) + if extra_vertices: + vertices = np.concatenate([vertices, np.asarray(extra_vertices)], axis=0) + if radii is not None: + radii = np.concatenate([radii, np.asarray(extra_radii)]) + + if prune_edges: + keep = np.ones(len(edges), dtype=bool) + keep[list(prune_edges)] = False + edges = edges[keep] + + vertices, edges, radii = _compact(vertices, edges, radii) + _snapshot("split") + + if join_dist > 0: + vertices, edges, radii = join_close_components( + vertices, edges, join_dist, + min_join_angle=min_join_angle, direction_span=direction_span, radii=radii, + ) + _snapshot("join") + return vertices, edges, radii + + +def _adjacency(num_nodes, edges): + src = np.concatenate([edges[:, 0], edges[:, 1]]) + dst = np.concatenate([edges[:, 1], edges[:, 0]]) + eid = np.concatenate([np.arange(len(edges)), np.arange(len(edges))]) + order = np.argsort(src, kind="stable") + dst, eid = dst[order], eid[order] + degrees = np.bincount(src, minlength=num_nodes) + indptr = np.zeros(num_nodes + 1, dtype=np.int64) + np.cumsum(degrees, out=indptr[1:]) + return indptr, dst, eid, degrees + + +def remove_ticks(vertices, edges, tick_length, radii=None): + """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 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 pruned. + + Parameters + ---------- + vertices: + Float array with shape ``(V, D)`` of skeleton vertex coordinates. + edges: + Integer array with shape ``(E, 2)`` indexing ``vertices``. + tick_length: + Maximum branch length (physical) that may be pruned. + radii: + Optional per-vertex radii, carried through the same remapping. + + Returns + ------- + vertices, edges, radii: + The pruned arrays, reindexed with unused vertices dropped. ``radii`` is + ``None`` when no input radii were given. + """ + vertices = np.asarray(vertices, dtype=np.float64) + edges = np.asarray(edges, dtype=np.int64) + num_nodes = len(vertices) + if len(edges) == 0: + return vertices, edges.copy(), radii + + indptr, dst, eid, degrees = _adjacency(num_nodes, edges) + + # Distance supergraph: sid -> [end_a, end_b, length, edge_ids]. + supers = {} + incident = defaultdict(set) + edge_used = np.zeros(len(edges), dtype=bool) + sid = 0 + for node in map(int, np.where(degrees != 2)[0]): + for k in range(indptr[node], indptr[node + 1]): + if edge_used[eid[k]]: + continue + prev, cur = node, int(dst[k]) + path = [int(eid[k])] + length = float(np.linalg.norm(vertices[cur] - vertices[node])) + while degrees[cur] == 2: + s, e = indptr[cur], indptr[cur + 1] + nbrs, eids = dst[s:e], eid[s:e] + pick = 0 if int(nbrs[0]) != prev else 1 + path.append(int(eids[pick])) + nxt = int(nbrs[pick]) + length += float(np.linalg.norm(vertices[nxt] - vertices[cur])) + prev, cur = cur, nxt + for pe in path: + edge_used[pe] = True + supers[sid] = [node, cur, length, path] + incident[node].add(sid) + incident[cur].add(sid) + sid += 1 + + dropped = set() + while True: + best, best_len = None, tick_length + for s, (a, b, length, _) in supers.items(): + terminal_a, terminal_b = len(incident[a]) == 1, len(incident[b]) == 1 + if (terminal_a ^ terminal_b) and length < best_len: + best, best_len = s, length + if best is None: + break + a, b, length, path = supers.pop(best) + incident[a].discard(best) + incident[b].discard(best) + dropped.update(path) + for node in (a, b): + if len(incident[node]) == 2: + s1, s2 = incident[node] + a1, b1, l1, p1 = supers.pop(s1) + a2, b2, l2, p2 = supers.pop(s2) + far1 = b1 if a1 == node else a1 + far2 = b2 if a2 == node else a2 + for x in (far1, far2, node): + incident[x].discard(s1) + incident[x].discard(s2) + supers[sid] = [far1, far2, l1 + l2, p1 + p2] + incident[far1].add(sid) + incident[far2].add(sid) + sid += 1 + + if dropped: + keep = np.ones(len(edges), dtype=bool) + keep[list(dropped)] = False + edges = edges[keep] + else: + edges = edges.copy() + + vertices, edges, radii = _compact(vertices, edges, radii) + return vertices, edges, radii + + +def _endpoint_tangent(endpoint, indptr, dst, degrees, vertices, span): + if degrees[endpoint] == 0: + return None + prev, cur = endpoint, int(dst[indptr[endpoint]]) + for _ in range(span - 1): + if degrees[cur] != 2: + break + s, e = indptr[cur], indptr[cur + 1] + nbrs = dst[s:e] + nxt = int(nbrs[0]) if int(nbrs[0]) != prev else int(nbrs[1]) + prev, cur = cur, nxt + direction = vertices[endpoint] - vertices[cur] # outward: interior -> tip + norm = np.linalg.norm(direction) + return direction / norm if norm > 0 else None + + +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 ``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 joined at most once and each + endpoint is used once. + + Parameters + ---------- + vertices, edges: + Skeleton graph; ``edges`` indexes ``vertices``. + dist: + Maximum gap (physical) across which endpoints may be joined. + min_join_angle: + Minimum straightness (degrees) of the joined path; 180 is perfectly + collinear. + direction_span: + Nodes over which each endpoint's tangent is measured. + radii: + Optional per-vertex radii, returned unchanged. + + Returns + ------- + vertices, edges, radii: + ``vertices`` and ``radii`` unchanged; ``edges`` has the join edges added. + """ + from scipy.spatial import cKDTree + + vertices = np.asarray(vertices, dtype=np.float64) + edges = np.asarray(edges, dtype=np.int64) + n = len(vertices) + if len(edges) == 0: + return vertices, edges.copy(), radii + + indptr, dst, _, degrees = _adjacency(n, edges) + labels = connected_components(skeleton_to_graph(vertices, edges)) + endpoints = np.where(degrees <= 1)[0] + if len(endpoints) < 2: + return vertices, edges.copy(), radii + + tol = np.deg2rad(180.0 - min_join_angle) + tree = cKDTree(vertices[endpoints]) + candidates = [] + for ia, ib in tree.query_pairs(dist): + a, b = int(endpoints[ia]), int(endpoints[ib]) + if labels[a] == labels[b]: + continue + gap = vertices[b] - vertices[a] + dist = float(np.linalg.norm(gap)) + if dist == 0.0: + continue + ta = _endpoint_tangent(a, indptr, dst, degrees, vertices, direction_span) + tb = _endpoint_tangent(b, indptr, dst, degrees, vertices, direction_span) + if ta is None or tb is None: + continue + gdir = gap / dist + bend_a = np.arccos(np.clip(ta @ gdir, -1.0, 1.0)) + bend_b = np.arccos(np.clip(tb @ -gdir, -1.0, 1.0)) + if bend_a <= tol and bend_b <= tol: + candidates.append((dist, a, b)) + + candidates.sort() + uf = UnionFind(int(labels.max()) + 1) + + used, new_edges = set(), [] + for _, a, b in candidates: + if a in used or b in used: + continue + if uf.find(int(labels[a])) == uf.find(int(labels[b])): + continue + uf.merge(int(labels[a]), int(labels[b])) + new_edges.append([a, b]) + used.add(a) + used.add(b) + + if new_edges: + edges = np.concatenate([edges, np.asarray(new_edges, dtype=np.int64)], axis=0) + else: + edges = edges.copy() + return vertices, edges, radii + + +def draw_instances(vertices, edges, labels, shape, radius=1): + """Rasterize a labeled skeleton graph into a dense instance volume. + + Each edge is drawn as a line between its two vertices and dilated by a ball + of the given radius. Every voxel on a component's tubes is set to that + component's label plus one, so background stays zero. + + Parameters + ---------- + vertices: + Float array with shape ``(V, 3)`` of vertex coordinates in voxel index + space matching ``shape`` (``(z, y, x)`` order). + edges: + Integer array with shape ``(E, 2)`` indexing ``vertices``. + labels: + Per-vertex integer labels, e.g. from + :func:`bioimage_cpp.graph.connected_components`. + shape: + Output volume shape ``(Z, Y, X)``. + radius: + Tube radius in voxels. + + Returns + ------- + numpy.ndarray + Integer volume of ``shape`` with background ``0`` and each tube voxel + set to ``labels[endpoint] + 1``. + """ + vertices = np.asarray(vertices) + edges = np.asarray(edges) + labels = np.asarray(labels) + shape = tuple(int(s) for s in shape) + + if len(edges) == 0: + return np.zeros(shape, dtype=np.uint16) + + r = int(radius) + zz, yy, xx = np.ogrid[-r:r + 1, -r:r + 1, -r:r + 1] + offsets = np.stack(np.where(zz ** 2 + yy ** 2 + xx ** 2 <= r * r), axis=1) - r + + vi = np.rint(vertices).astype(np.int64) + centers, center_labels = [], [] + for a, b in edges: + 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) + center_labels = np.concatenate(center_labels) + + dtype = np.uint16 if int(center_labels.max()) < 2 ** 16 else np.uint32 + volume = np.zeros(shape, dtype=dtype) + shape_arr = np.asarray(shape) + for offset in offsets: + coords = centers + offset + inside = ((coords >= 0) & (coords < shape_arr)).all(axis=1) + c = coords[inside] + volume[c[:, 0], c[:, 1], c[:, 2]] = center_labels[inside] + return volume + + +__all__ = ["clean_graph", "draw_instances", "join_close_components", "remove_ticks"] diff --git a/tests/skeleton/test_postprocessing.py b/tests/skeleton/test_postprocessing.py new file mode 100644 index 0000000..65f408d --- /dev/null +++ b/tests/skeleton/test_postprocessing.py @@ -0,0 +1,182 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic +from bioimage_cpp.skeleton.postprocessing import split_degree3, split_degree4 + + +def _graph(vertices, edges): + vertices = np.asarray(vertices, dtype=np.float64) + edges = np.asarray(edges, dtype=np.int64) + return vertices, edges, bic.skeleton.skeleton_to_graph(vertices, edges) + + +def _component_count(vertices, edges): + labels = bic.graph.connected_components(bic.skeleton.skeleton_to_graph(vertices, edges)) + return len(np.unique(labels)) + + +def test_split_degree3_splits_odd_arm(): + # Node 0 is a T: a straight through-pair along x (nodes 1, 2) plus an odd arm along y (node 3). + vertices, _, graph = _graph( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], + [[0, 1], [0, 2], [0, 3]], + ) + assert split_degree3(0, graph, vertices, direction_span=1, min_branch_angle=30.0) == [2] + + +def test_split_degree3_keeps_collinear_arm(): + # The third arm now runs almost parallel to the through pair, below min_branch_angle. + vertices, _, graph = _graph( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, -1.0], [0.0, 1.0, 10.0]], + [[0, 1], [0, 2], [0, 3]], + ) + assert split_degree3(0, graph, vertices, direction_span=1, min_branch_angle=30.0) is None + + +def test_split_degree4_splits_through_pair(): + # A '+' crossing: edges 0,1 are the x through-pair, edges 2,3 the y through-pair. + vertices, _, graph = _graph( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, -1.0, 0.0]], + [[0, 1], [0, 2], [0, 3], [0, 4]], + ) + result = split_degree4(0, graph, vertices, direction_span=1, min_through_angle=170.0) + assert set(result) in ({0, 1}, {2, 3}) + + +def test_split_degree4_keeps_non_collinear_crossing(): + # Four arms at 0, 60, 120, 180 degrees: no pair is collinear enough to split. + angle = np.deg2rad([0.0, 60.0, 120.0, 180.0]) + arms = np.stack([np.zeros_like(angle), np.sin(angle), np.cos(angle)], axis=1) + vertices, _, graph = _graph( + np.concatenate([[[0.0, 0.0, 0.0]], arms]), + [[0, 1], [0, 2], [0, 3], [0, 4]], + ) + assert split_degree4(0, graph, vertices, direction_span=1, min_through_angle=170.0) is None + + +@pytest.mark.parametrize( + "tick_length, n_vertices, n_edges", + [(1.5, 5, 4), (0.5, 6, 5)], +) +def test_remove_ticks_prunes_spurs_below_threshold(tick_length, n_vertices, n_edges): + # Backbone 0-1-2-3-4 along x with a length-1 spur (node 5) off the branch node 2. + vertices = np.array( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0], + [0.0, 0.0, 4.0], [0.0, 1.0, 2.0]], + ) + edges = np.array([[0, 1], [1, 2], [2, 3], [3, 4], [2, 5]], dtype=np.int64) + + out_vertices, out_edges, _ = bic.skeleton.remove_ticks(vertices, edges, tick_length=tick_length) + + assert len(out_vertices) == n_vertices + assert len(out_edges) == n_edges + + +def test_remove_ticks_keeps_standalone_path(): + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]]) + edges = np.array([[0, 1], [1, 2]], dtype=np.int64) + + out_vertices, out_edges, _ = bic.skeleton.remove_ticks(vertices, edges, tick_length=100.0) + + assert len(out_vertices) == 3 + assert len(out_edges) == 2 + + +def test_remove_ticks_empty_edges(): + vertices = np.zeros((3, 3), dtype=np.float64) + edges = np.empty((0, 2), dtype=np.int64) + + out_vertices, out_edges, _ = bic.skeleton.remove_ticks(vertices, edges, tick_length=10.0) + + assert len(out_vertices) == 3 + assert len(out_edges) == 0 + + +def test_join_close_components_links_collinear_endpoints(): + # Two collinear fragments along x with a gap of 2 between node 1 and node 2. + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 3.0], [0.0, 0.0, 4.0]]) + edges = np.array([[0, 1], [2, 3]], dtype=np.int64) + + _, out_edges, _ = bic.skeleton.join_close_components(vertices, edges, dist=2.5) + + assert len(out_edges) == 3 + assert (1, 2) in {tuple(sorted(map(int, edge))) for edge in out_edges} + + +def test_join_close_components_ignores_beyond_dist(): + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 3.0], [0.0, 0.0, 4.0]]) + edges = np.array([[0, 1], [2, 3]], dtype=np.int64) + + _, out_edges, _ = bic.skeleton.join_close_components(vertices, edges, dist=1.5) + + assert len(out_edges) == 2 + + +def test_join_close_components_ignores_non_collinear(): + # Two fragments offset in y so the gap bends away from each endpoint tangent. + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 3.0, 2.0], [0.0, 4.0, 2.0]]) + edges = np.array([[0, 1], [2, 3]], dtype=np.int64) + + _, out_edges, _ = bic.skeleton.join_close_components(vertices, edges, dist=3.5) + + assert len(out_edges) == 2 + + +def test_join_close_components_skips_same_component(): + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]]) + edges = np.array([[0, 1], [1, 2]], dtype=np.int64) + + _, out_edges, _ = bic.skeleton.join_close_components(vertices, edges, dist=5.0) + + assert len(out_edges) == 2 + + +def test_clean_graph_splits_crossing(): + vertices = np.array( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, -1.0, 0.0]], + ) + edges = np.array([[0, 1], [0, 2], [0, 3], [0, 4]], dtype=np.int64) + radii = np.arange(len(vertices), dtype=np.float64) + + assert _component_count(vertices, edges) == 1 + + out_vertices, out_edges, out_radii = bic.skeleton.clean_graph( + vertices, edges, radii=radii, direction_span=1, tick_length=0.0, join_dist=0.0, + ) + + assert _component_count(out_vertices, out_edges) == 2 + assert len(out_radii) == len(out_vertices) + + +def test_draw_instances_labels_edge(): + vertices = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 3.0]]) + edges = np.array([[0, 1]], dtype=np.int64) + labels = np.array([2, 2]) + + volume = bic.skeleton.draw_instances(vertices, edges, labels, (1, 1, 6), radius=0) + + np.testing.assert_array_equal(volume[0, 0], [3, 3, 3, 3, 0, 0]) + assert set(np.unique(volume)) == {0, 3} + + +def test_draw_instances_dilates_by_radius(): + vertices = np.array([[2.0, 2.0, 0.0], [2.0, 2.0, 4.0]]) + edges = np.array([[0, 1]], dtype=np.int64) + labels = np.array([0, 0]) + + thin = bic.skeleton.draw_instances(vertices, edges, labels, (5, 5, 5), radius=0) + thick = bic.skeleton.draw_instances(vertices, edges, labels, (5, 5, 5), radius=1) + + assert thick.astype(bool).sum() > thin.astype(bool).sum() + + +def test_draw_instances_empty_edges(): + vertices = np.zeros((2, 3), dtype=np.float64) + edges = np.empty((0, 2), dtype=np.int64) + labels = np.zeros(2, dtype=np.int64) + + volume = bic.skeleton.draw_instances(vertices, edges, labels, (2, 3, 4)) + + assert volume.shape == (2, 3, 4) + assert volume.sum() == 0