From 366d6852edfcc3623faf4fca9ad9fba58a518993 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 16 Jul 2026 18:16:26 +0200 Subject: [PATCH 1/3] Implement primitives for blockwise skeletonization --- CMakeLists.txt | 1 + MIGRATION_GUIDE.md | 80 +- development/skeleton/PERFORMANCE_NOTES.md | 55 ++ development/skeleton/blockwise_stitching.py | 173 +++++ .../skeleton/compare_blockwise_kimimaro.py | 686 ++++++++++++++++++ .../skeleton/detail/components.hxx | 206 +++++- .../skeleton/distributed/border_targets.hxx | 390 ++++++++++ .../skeleton/distributed/merge.hxx | 413 +++++++++++ include/bioimage_cpp/skeleton/teasar.hxx | 537 ++++++++++---- .../bioimage_cpp/skeleton/teasar_labels.hxx | 139 +++- src/bindings/module.cxx | 1 + src/bindings/skeleton.hxx | 1 + src/bindings/skeleton_distributed.cxx | 537 ++++++++++++++ src/bioimage_cpp/skeleton/__init__.py | 5 +- src/bioimage_cpp/skeleton/distributed.py | 396 ++++++++++ tests/skeleton/test_distributed.py | 339 +++++++++ 16 files changed, 3813 insertions(+), 146 deletions(-) create mode 100644 development/skeleton/blockwise_stitching.py create mode 100644 development/skeleton/compare_blockwise_kimimaro.py create mode 100644 include/bioimage_cpp/skeleton/distributed/border_targets.hxx create mode 100644 include/bioimage_cpp/skeleton/distributed/merge.hxx create mode 100644 src/bindings/skeleton_distributed.cxx create mode 100644 src/bioimage_cpp/skeleton/distributed.py create mode 100644 tests/skeleton/test_distributed.py diff --git a/CMakeLists.txt b/CMakeLists.txt index b7d5196..8ca97ec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -23,6 +23,7 @@ nanobind_add_module(_core src/bindings/mesh.cxx src/bindings/segmentation.cxx src/bindings/skeleton.cxx + src/bindings/skeleton_distributed.cxx src/bindings/transformation.cxx src/bindings/util.cxx src/bindings/utils.cxx diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 831f31f..adf74b7 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -2589,12 +2589,14 @@ Important differences and current scope: distance-to-boundary field, a deterministic two-sweep root, a physical Dijkstra distance-from-root field, penalized repeated Dijkstra paths, and rolling invalidation with radius `scale * radius + constant`. -- This first correctness-oriented version uses a physical axis-aligned - invalidation cube. It does not implement kimimaro's soma handling, hole - filling, border stitching, manual targets, component dust filtering, - cross-section metadata, or postprocessing heuristics. These differences can - change branch positions and vertex counts, so output is not expected to be - vertex-for-vertex identical to kimimaro. +- This correctness-oriented implementation uses a physical axis-aligned + invalidation cube. The ordinary in-core entry points do not automatically + perform block stitching or accept manual targets; the dedicated block APIs + below provide mandatory interface targets and exact consolidation. Soma + handling, hole filling, component dust filtering, cross-section metadata, + and kimimaro's other postprocessing heuristics are not implemented. These + differences can change branch positions and vertex counts, so output is not + expected to be vertex-for-vertex identical to kimimaro. - The C++ core remains dependency-free. Component discovery uses x-runs and a union-find rather than dense component-label images. `number_of_threads=1` is the default and `0` uses hardware concurrency; one shared budget covers @@ -2604,14 +2606,76 @@ Important differences and current scope: heap. Compact IDs avoid full-volume shortest-path fields; the public Dijkstra functions remain dense and exact FP64. -Correctness tests are under `tests/skeleton/test_teasar.py` and -`tests/skeleton/test_teasar_labels.py`. The independent +### Blockwise skeletonization and exact stitching + +`bioimage_cpp.skeleton.distributed` provides the dependency-light primitives +needed by an external block scheduler. It does not perform I/O, enumerate +blocks, launch tasks, or serialize artifacts. The implemented layout uses one +shared voxel plane: a core block with interval `[a, b)` is processed through +`b + 1` when it has a high-side neighbor, while the neighbor begins at `b`. + +```python +dist = bic.skeleton.distributed + +# `left` includes its high-side overlap plane and starts at `left_origin`. +targets = dist.block_border_targets( + left, + faces=[(2, "high")], + origin=left_origin, + spacing=spacing, +) +left_fragment = dist.block_teasar( + left, + open_faces=[(2, "high")], + origin=left_origin, + required_targets=targets, + spacing=spacing, +) + +# Repeat independently for every processing block. Neighboring calls select +# identical global targets on their shared plane. +merged = dist.merge_block_skeletons(block_fragments) +forest = dist.minimum_spanning_forest(merged, spacing=spacing) # optional +vertices, edges, radii = dist.lattice_to_physical(forest, spacing=spacing) +``` + +Block fragments use global `int64` lattice vertices, `uint64` edges, and +`float32` physical radii. Integer coordinates are the exact merge identity; +physical proximity is never used to invent a connection. Duplicate radii are +reduced with `maximum`, and canonical sort-and-unique output makes +`merge_block_skeletons` associative, commutative, idempotent, and suitable for +hierarchical reduction. `merge_block_skeleton_maps` applies the same operation +per exact semantic label, including negative and large `uint64` labels. + +`open_faces` is mandatory even when empty. It identifies artificial cuts where +the distance transform may look through the processing-block boundary; paths +remain restricted to real input foreground. This prevents interface radii from +collapsing to one voxel. When present, the lexicographically greatest required +target on an open face becomes the deterministic component root. + +The labeled counterparts are `block_border_targets_labels` and +`block_teasar_labels`. Border targets are selected per 8-connected same-label +face patch by maximum anisotropic 2D distance transform with deterministic, +orientation-stable plateau resolution. Several valid targets can create cycles +after graph union, so exact consolidation keeps cycles and +`minimum_spanning_forest` removes them only when explicitly called. Blocked +skeletons are connected through their anchors but are not expected to be +vertex-for-vertex identical to whole-volume TEASAR because each block still has +less path-selection context. + +Correctness tests are under `tests/skeleton/test_teasar.py`, +`tests/skeleton/test_teasar_labels.py`, and +`tests/skeleton/test_distributed.py`. The independent Dijkstra benchmark is `development/distance/benchmark_dijkstra.py`; the end-to-end benchmark is `development/skeleton/benchmark_teasar.py`. Use `--suite all --kimimaro` to compare paired binary and multi-label scenarios, `--memory` for fresh-process peak-RSS probes, or `--sequential-backends` for the dense/compact FP64 binary design matrix. Extended density, spacing, and PDRF regimes are in `development/skeleton/benchmark_teasar_sequential.py`. +The development-only serial block harness is +`development/skeleton/blockwise_stitching.py`; the rigorous one-thread layered +comparison against whole and blocked Kimimaro is +`development/skeleton/compare_blockwise_kimimaro.py`. ## I/O and Build Dependencies diff --git a/development/skeleton/PERFORMANCE_NOTES.md b/development/skeleton/PERFORMANCE_NOTES.md index 25bae44..801452d 100644 --- a/development/skeleton/PERFORMANCE_NOTES.md +++ b/development/skeleton/PERFORMANCE_NOTES.md @@ -736,3 +736,58 @@ leave the initial linear mode. Skeleton output remained array-exact across backends and worker counts. Final verification after the optimization and benchmark addition: `1182 passed` with third-party pytest plugin autoload disabled. + +## Blockwise open-boundary correctness and Kimimaro comparison (2026-07-16) + +The first block prototype reused the ordinary per-component zero crop halo for +its DBF. On an artificial cut this made every interface radius exactly one +voxel. Thick branching tubes remained connected, but accumulated 2.1--3.1x the +whole-volume Kimimaro cable length. A controlled Kimimaro run reproduced the +failure when `black_border=True`, isolating the problem from target selection, +exact consolidation, and cycle removal. + +The corrected block path keeps its real foreground path mask unchanged and +uses a separate EDT-only mask with one ghost layer across explicitly declared +`open_faces`. One deterministic face target roots each affected component. +Ghost voxels can never become graph vertices. The complete reference workflow +is: + +```bash +python development/skeleton/compare_blockwise_kimimaro.py \ + --warmup 1 --repeats 7 --check \ + --json /tmp/blockwise-kimimaro.json +``` + +The run used bioimage-cpp 0.7.0, Kimimaro 5.8.1, EDT 3.1.1, NumPy 2.4.6, +SciPy 1.17.1, isotropic spacing, and one thread for every backend. Kimimaro's +blocked fragments used `fix_borders=True`; both block backends then used the +same exact lattice merge and minimum-spanning forest. + +| case | bio block | Kimimaro whole | Kimimaro block | bio / Kimimaro-block cable | direct p95 / Hausdorff | +| --- | ---: | ---: | ---: | ---: | ---: | +| thin line, 16 blocks | 3.52 ms | 1.66 ms | 6.95 ms | 1.000 | 0.00 / 0.00 | +| oblique tube `64^3`, anisotropic, 8 blocks | 4.41 ms | 11.44 ms | 17.78 ms | 1.000 | 0.00 / 0.00 | +| tube `64^3`, 8 blocks | 6.73 ms | 18.34 ms | 46.97 ms | 0.988 | 0.20 / 1.41 | +| tube `96^3`, 8 blocks | 13.04 ms | 48.95 ms | 72.01 ms | 0.916 | 1.41 / 3.74 | +| tube `128^3`, seam-aligned, 8 blocks | 21.70 ms | 102.29 ms | 123.76 ms | 0.988 | 0.00 / 4.24 | +| tube `128^3`, 27 blocks | 24.85 ms | 102.02 ms | 118.78 ms | 0.903 | 1.27 / 5.00 | + +The thin line remains graph-exact. The anisotropic oblique case has identical +blocked vertex geometry and cable length. All 112 expected local interface +anchors were present in both controlled implementations, and their radii +matched Kimimaro exactly. Across all thick cases, final graphs were connected +and acyclic, cable length stayed within 10% of Kimimaro's blocked result, and +every quality gate passed. On the nontrivial timing cases, bioimage-cpp was +2.6--4.7x faster than whole-volume one-thread Kimimaro and 4.0--7.0x faster than +its serial block workflow. + +On the final `128^3` calls, the last measured phase breakdowns were: + +| layout | targets | local TEASAR | merge | forest | +| --- | ---: | ---: | ---: | ---: | +| 8 blocks | 2.39 ms | 17.91 ms | 0.25 ms | 0.09 ms | +| 27 blocks | 4.68 ms | 15.89 ms | 0.35 ms | 0.10 ms | + +Merge and cycle removal remain negligible. More blocks mainly increase face +analysis and Python orchestration. The comparison intentionally does not apply +Kimimaro's heuristic nearest-component joining, dust removal, or tick pruning. diff --git a/development/skeleton/blockwise_stitching.py b/development/skeleton/blockwise_stitching.py new file mode 100644 index 0000000..2de95d9 --- /dev/null +++ b/development/skeleton/blockwise_stitching.py @@ -0,0 +1,173 @@ +"""Minimal serial harness for blockwise TEASAR and exact stitching. + +This is deliberately development-only orchestration. It performs no I/O, +parallel scheduling, retries, or artifact serialization; those responsibilities +belong to downstream block-processing systems. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +import bioimage_cpp as bic + + +dist = bic.skeleton.distributed + + +def _processing_blocks(volume: np.ndarray, block_shape: Sequence[int]): + blocking = bic.utils.Blocking([0, 0, 0], list(volume.shape), list(block_shape)) + for block_id in range(blocking.number_of_blocks): + core = blocking.get_block(block_id) + begin = [int(value) for value in core.begin] + end = [int(value) for value in core.end] + faces = [] + for axis in range(3): + if blocking.get_neighbor_id(block_id, axis, lower=True) >= 0: + faces.append((axis, "low")) + if blocking.get_neighbor_id(block_id, axis, lower=False) >= 0: + faces.append((axis, "high")) + end[axis] += 1 + slices = tuple(slice(begin[axis], end[axis]) for axis in range(3)) + yield blocking, block_id, np.ascontiguousarray(volume[slices]), begin, faces + + +def _unique_coordinates(parts: list[np.ndarray]) -> np.ndarray: + if not parts: + return np.empty((0, 3), dtype=np.int64) + return np.unique(np.concatenate(parts, axis=0), axis=0) + + +def _unique_target_map(parts: list[dict[int, np.ndarray]]) -> dict[int, np.ndarray]: + labels = sorted({label for part in parts for label in part}) + return { + label: _unique_coordinates([part[label] for part in parts if label in part]) + for label in labels + } + + +def _assert_matching_interfaces(blocking, per_face) -> None: + for block_id in range(blocking.number_of_blocks): + for axis in range(3): + neighbor = blocking.get_neighbor_id(block_id, axis, lower=False) + if neighbor < 0: + continue + left = per_face[(block_id, axis, "high")] + right = per_face[(neighbor, axis, "low")] + if isinstance(left, dict): + if left.keys() != right.keys(): + raise AssertionError( + f"target labels disagree across blocks {block_id}/{neighbor}" + ) + for label in left: + np.testing.assert_array_equal(left[label], right[label]) + else: + np.testing.assert_array_equal(left, right) + + +def run_blockwise_binary( + mask: np.ndarray, + block_shape: Sequence[int], + *, + spacing=(1.0, 1.0, 1.0), + remove_cycles: bool = False, +): + """Run the complete binary block pipeline serially and return a lattice graph.""" + fragments = [] + per_face = {} + blocking = None + for blocking, block_id, block, origin, faces in _processing_blocks( + np.asarray(mask), block_shape + ): + face_targets = [] + for axis, side in faces: + targets = dist.block_border_targets( + block, + [(axis, side)], + origin=origin, + spacing=spacing, + number_of_threads=1, + ) + per_face[(block_id, axis, side)] = targets + face_targets.append(targets) + targets = _unique_coordinates(face_targets) + fragments.append( + dist.block_teasar( + block, + open_faces=faces, + origin=origin, + required_targets=targets, + spacing=spacing, + number_of_threads=1, + ) + ) + if blocking is None: + raise ValueError("mask must be a three-dimensional array") + _assert_matching_interfaces(blocking, per_face) + merged = dist.merge_block_skeletons(fragments) + if remove_cycles: + merged = dist.minimum_spanning_forest(merged, spacing=spacing) + return merged + + +def run_blockwise_labels( + labels: np.ndarray, + block_shape: Sequence[int], + *, + background: int = 0, + spacing=(1.0, 1.0, 1.0), + remove_cycles: bool = False, +): + """Run the labeled block pipeline serially and return lattice graphs by label.""" + fragment_maps = [] + per_face = {} + blocking = None + for blocking, block_id, block, origin, faces in _processing_blocks( + np.asarray(labels), block_shape + ): + face_targets = [] + for axis, side in faces: + targets = dist.block_border_targets_labels( + block, + [(axis, side)], + origin=origin, + background=background, + spacing=spacing, + number_of_threads=1, + ) + per_face[(block_id, axis, side)] = targets + face_targets.append(targets) + targets = _unique_target_map(face_targets) + fragment_maps.append( + dist.block_teasar_labels( + block, + open_faces=faces, + origin=origin, + required_targets=targets, + background=background, + spacing=spacing, + number_of_threads=1, + ) + ) + if blocking is None: + raise ValueError("labels must be a three-dimensional array") + _assert_matching_interfaces(blocking, per_face) + merged = dist.merge_block_skeleton_maps(fragment_maps) + if remove_cycles: + merged = { + label: dist.minimum_spanning_forest(fragment, spacing=spacing) + for label, fragment in merged.items() + } + return merged + + +if __name__ == "__main__": + example = np.zeros((17, 17, 25), dtype=np.uint8) + example[8, 8, 2:23] = 1 + result = run_blockwise_binary(example, (9, 9, 8), remove_cycles=True) + print( + f"stitched {len(result[0])} lattice vertices and {len(result[1])} edges " + "from serial processing blocks" + ) diff --git a/development/skeleton/compare_blockwise_kimimaro.py b/development/skeleton/compare_blockwise_kimimaro.py new file mode 100644 index 0000000..0bee627 --- /dev/null +++ b/development/skeleton/compare_blockwise_kimimaro.py @@ -0,0 +1,686 @@ +"""Rigorous serial comparison of blockwise TEASAR against Kimimaro. + +This workflow is deliberately development-only. It uses Kimimaro and SciPy as +reference dependencies, fixes every backend to one thread, applies the same +exact lattice merge and minimum-spanning forest to both block implementations, +and can enforce the quality gates recorded in ``SKEL-PARALLEL.md``. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +import gc +import importlib.metadata +import json +import os +from pathlib import Path +import random +from statistics import median +import sys +from time import perf_counter + +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" + +import numpy as np +from scipy.spatial import cKDTree + +import bioimage_cpp as bic + +from benchmark_teasar import draw_segment, make_branching_tube +from blockwise_stitching import ( + _assert_matching_interfaces, + _processing_blocks, + _unique_coordinates, +) + + +dist = bic.skeleton.distributed +SPACING = (1.0, 1.0, 1.0) +PARAMETERS = { + "scale": 1.5, + "constant": 0.0, + "pdrf_scale": 100000.0, + "pdrf_exponent": 4.0, +} + + +@dataclass(frozen=True) +class Workload: + name: str + volume: np.ndarray + block_shape: tuple[int, int, int] + spacing: tuple[float, float, float] = SPACING + timing_gate: bool = True + + +@dataclass +class BlockResult: + raw: tuple[np.ndarray, np.ndarray, np.ndarray] + forest: tuple[np.ndarray, np.ndarray, np.ndarray] + phases: dict[str, float] + fragments: list[tuple[np.ndarray, np.ndarray, np.ndarray]] + targets: list[np.ndarray] + + +def _kimimaro_parameters(): + return { + "scale": PARAMETERS["scale"], + "const": PARAMETERS["constant"], + "pdrf_scale": PARAMETERS["pdrf_scale"], + "pdrf_exponent": int(PARAMETERS["pdrf_exponent"]), + "soma_detection_threshold": float("inf"), + "soma_acceptance_threshold": float("inf"), + "soma_invalidation_scale": 1.0, + "soma_invalidation_const": 0.0, + } + + +def _kimimaro_call( + volume, *, spacing, fix_borders, extra_targets_before=() +): + import kimimaro + + return kimimaro.skeletonize( + volume, + teasar_params=_kimimaro_parameters(), + anisotropy=spacing, + object_ids=[1], + dust_threshold=0, + progress=False, + fix_branching=True, + fix_borders=fix_borders, + fill_holes=False, + parallel=1, + extra_targets_before=list(extra_targets_before), + ) + + +def _physical_to_lattice(vertices, spacing, backend): + scaled = np.asarray(vertices, dtype=float) / np.asarray(spacing, dtype=float) + rounded = np.rint(scaled) + if not np.allclose(scaled, rounded, rtol=0.0, atol=1e-5): + error = float(np.max(np.abs(scaled - rounded))) + raise RuntimeError( + f"{backend} returned vertices off the voxel lattice by {error:.6g}" + ) + return rounded.astype(np.int64) + + +def _kimimaro_graph(result, spacing, origin=(0, 0, 0)): + if 1 not in result: + return ( + np.empty((0, 3), np.int64), + np.empty((0, 2), np.uint64), + np.empty((0,), np.float32), + ) + skeleton = result[1] + vertices = _physical_to_lattice( + skeleton.vertices, spacing, "Kimimaro" + ) + vertices += np.asarray(origin, dtype=np.int64) + return ( + vertices, + np.asarray(skeleton.edges, dtype=np.uint64), + np.asarray(skeleton.radius, dtype=np.float32), + ) + + +def run_bioimage_block(volume, block_shape, spacing) -> BlockResult: + phases = {name: 0.0 for name in ("targets", "teasar", "merge", "forest")} + fragments = [] + block_targets = [] + per_face = {} + blocking = None + for blocking, block_id, block, origin, faces in _processing_blocks( + volume, block_shape + ): + start = perf_counter() + face_targets = [] + for axis, side in faces: + targets = dist.block_border_targets( + block, + [(axis, side)], + origin=origin, + spacing=spacing, + number_of_threads=1, + ) + per_face[(block_id, axis, side)] = targets + face_targets.append(targets) + targets = _unique_coordinates(face_targets) + phases["targets"] += perf_counter() - start + + start = perf_counter() + fragments.append( + dist.block_teasar( + block, + open_faces=faces, + origin=origin, + required_targets=targets, + spacing=spacing, + number_of_threads=1, + **PARAMETERS, + ) + ) + phases["teasar"] += perf_counter() - start + block_targets.append(targets) + if blocking is None: + raise RuntimeError("comparison volume produced no blocks") + _assert_matching_interfaces(blocking, per_face) + + start = perf_counter() + raw = dist.merge_block_skeletons(fragments) + phases["merge"] += perf_counter() - start + start = perf_counter() + forest = dist.minimum_spanning_forest(raw, spacing=spacing) + phases["forest"] += perf_counter() - start + return BlockResult(raw, forest, phases, fragments, block_targets) + + +def run_kimimaro_block(volume, block_shape, spacing) -> BlockResult: + phases = {name: 0.0 for name in ("targets", "teasar", "merge", "forest")} + fragments = [] + for _, _, block, origin, _ in _processing_blocks(volume, block_shape): + start = perf_counter() + fragments.append( + _kimimaro_graph( + _kimimaro_call( + block, spacing=spacing, fix_borders=True + ), + spacing, + origin=origin, + ) + ) + phases["teasar"] += perf_counter() - start + start = perf_counter() + raw = dist.merge_block_skeletons(fragments) + phases["merge"] += perf_counter() - start + start = perf_counter() + forest = dist.minimum_spanning_forest(raw, spacing=spacing) + phases["forest"] += perf_counter() - start + return BlockResult(raw, forest, phases, fragments, []) + + +def run_bioimage_whole(volume, spacing): + vertices, edges, radii = bic.skeleton.teasar( + volume, + spacing=spacing, + number_of_threads=1, + **PARAMETERS, + ) + lattice_vertices = _physical_to_lattice( + vertices, spacing, "bioimage-cpp" + ) + return lattice_vertices, edges, radii + + +def run_kimimaro_whole(volume, spacing): + return _kimimaro_graph( + _kimimaro_call(volume, spacing=spacing, fix_borders=False), spacing + ) + + +def _number_of_components(number_of_vertices, edges): + parent = np.arange(number_of_vertices, dtype=np.int64) + + def find(node): + while parent[node] != node: + parent[node] = parent[parent[node]] + node = int(parent[node]) + return node + + for first, second in np.asarray(edges, dtype=np.int64): + first_root = find(int(first)) + second_root = find(int(second)) + if first_root != second_root: + parent[second_root] = first_root + return len({find(node) for node in range(number_of_vertices)}) + + +def graph_stats(graph, spacing): + vertices, edges, _ = graph + edge_ids = np.asarray(edges, dtype=np.int64) + components = _number_of_components(len(vertices), edge_ids) + degree = np.bincount(edge_ids.ravel(), minlength=len(vertices)) if len(edges) else np.zeros(len(vertices), dtype=np.int64) + length = 0.0 + if len(edge_ids): + edge_vectors = ( + vertices[edge_ids[:, 0]] - vertices[edge_ids[:, 1]] + ) * np.asarray(spacing) + length = float(np.linalg.norm(edge_vectors, axis=1).sum()) + return { + "vertices": int(len(vertices)), + "edges": int(len(edges)), + "components": int(components), + "cycle_rank": int(len(edges) - len(vertices) + components), + "endpoints": int(np.count_nonzero(degree == 1)), + "branchpoints": int(np.count_nonzero(degree >= 3)), + "contracted_degree_signature": sorted( + int(value) for value in degree if value != 2 + ), + "cable_length": length, + } + + +def distance_metrics(first, second, spacing): + first_vertices = ( + np.asarray(first[0], dtype=float) * np.asarray(spacing) + ) + second_vertices = ( + np.asarray(second[0], dtype=float) * np.asarray(spacing) + ) + if len(first_vertices) == 0 or len(second_vertices) == 0: + return {"median": float("inf"), "p95": float("inf"), "hausdorff": float("inf")} + first_to_second = cKDTree(second_vertices).query(first_vertices)[0] + second_to_first = cKDTree(first_vertices).query(second_vertices)[0] + both = np.concatenate([first_to_second, second_to_first]) + return { + "median": float(np.median(both)), + "p95": float(np.quantile(both, 0.95)), + "hausdorff": float(np.max(both)), + "first_exact_fraction": float(np.mean(first_to_second == 0)), + "second_exact_fraction": float(np.mean(second_to_first == 0)), + } + + +def critical_point_distances(first, second, predicate, spacing): + def points(graph): + vertices, edges, _ = graph + degree = np.bincount( + np.asarray(edges, dtype=np.int64).ravel(), minlength=len(vertices) + ) if len(edges) else np.zeros(len(vertices), dtype=np.int64) + return ( + np.asarray(vertices, dtype=float)[predicate(degree)] + * np.asarray(spacing) + ) + + first_points = points(first) + second_points = points(second) + if len(first_points) == 0 or len(second_points) == 0: + return {"p95": None, "hausdorff": None} + distances = np.concatenate([ + cKDTree(second_points).query(first_points)[0], + cKDTree(first_points).query(second_points)[0], + ]) + return { + "p95": float(np.quantile(distances, 0.95)), + "hausdorff": float(np.max(distances)), + } + + +def radius_metrics(first, second): + first_map = { + tuple(coordinate): float(radius) + for coordinate, radius in zip(first[0], first[2]) + } + second_map = { + tuple(coordinate): float(radius) + for coordinate, radius in zip(second[0], second[2]) + } + shared = sorted(first_map.keys() & second_map.keys()) + errors = np.asarray([ + abs(first_map[coordinate] - second_map[coordinate]) + for coordinate in shared + ]) + return { + "shared_vertices": len(shared), + "mean_absolute_error": float(errors.mean()) if len(errors) else None, + "maximum_absolute_error": float(errors.max()) if len(errors) else None, + } + + +def anchor_radius_errors(volume, block_shape, spacing, bio_result): + errors = [] + compared = 0 + expected = 0 + for (_, _, block, origin, _), fragment, targets in zip( + _processing_blocks(volume, block_shape), + bio_result.fragments, + bio_result.targets, + ): + if len(targets) == 0: + continue + expected += len(targets) + local_targets = targets - np.asarray(origin, dtype=np.int64) + reference = _kimimaro_graph( + _kimimaro_call( + block, + spacing=spacing, + fix_borders=False, + extra_targets_before=local_targets.tolist(), + ), + spacing, + origin=origin, + ) + bio_radii = { + tuple(coordinate): float(radius) + for coordinate, radius in zip(fragment[0], fragment[2]) + } + reference_radii = { + tuple(coordinate): float(radius) + for coordinate, radius in zip(reference[0], reference[2]) + } + for target in targets: + coordinate = tuple(target) + if coordinate not in reference_radii or coordinate not in bio_radii: + continue + compared += 1 + errors.append(abs(bio_radii[coordinate] - reference_radii[coordinate])) + return { + "compared": compared, + "expected": expected, + "maximum_absolute_error": max(errors, default=None), + "mean_absolute_error": float(np.mean(errors)) if errors else None, + } + + +def _canonical_graph(graph): + vertices, edges, radii = graph + radius = {tuple(v): float(r) for v, r in zip(vertices, radii)} + coordinate_edges = set() + for first, second in np.asarray(edges, dtype=np.int64): + uv = sorted((tuple(vertices[first]), tuple(vertices[second]))) + coordinate_edges.add(tuple(uv)) + return radius, coordinate_edges + + +def graphs_exact(first, second): + first_radius, first_edges = _canonical_graph(first) + second_radius, second_edges = _canonical_graph(second) + return first_radius == second_radius and first_edges == second_edges + + +def timed(backends, warmup, repeats): + results = {} + for _ in range(warmup): + for name, function in backends.items(): + results[name] = function() + samples = {name: [] for name in backends} + rng = random.Random(20260716) + gc.disable() + try: + for _ in range(repeats): + names = list(backends) + rng.shuffle(names) + for name in names: + start = perf_counter() + results[name] = backends[name]() + samples[name].append(perf_counter() - start) + finally: + gc.enable() + return results, samples + + +def workloads(quick): + line = np.zeros((17, 17, 25), dtype=np.uint8) + line[8, 8, 2:23] = 1 + output = [ + Workload("thin-line", line, (9, 9, 8), timing_gate=False) + ] + oblique = np.zeros((64, 64, 64), dtype=np.uint8) + draw_segment(oblique, (7, 12, 5), (56, 40, 50), 1) + output.append( + Workload( + "oblique-anisotropic-64-8", + oblique, + (32, 32, 32), + spacing=(1.5, 1.0, 1.0), + ) + ) + specs = [ + ("tube-64-8", 64, 3, (32, 32, 32)), + ("tube-96-8", 96, 4, (48, 48, 48)), + ] + if not quick: + specs.extend([ + ("tube-128-aligned-8", 128, 5, (64, 64, 64)), + ("tube-128-27", 128, 5, (48, 48, 48)), + ]) + output.extend( + Workload(name, make_branching_tube(size, radius), block_shape) + for name, size, radius, block_shape in specs + ) + return output + + +def compare_workload(workload, warmup, repeats): + backends = { + "bioimage_whole": lambda: run_bioimage_whole( + workload.volume, workload.spacing + ), + "bioimage_block": lambda: run_bioimage_block( + workload.volume, workload.block_shape, workload.spacing + ), + "kimimaro_whole": lambda: run_kimimaro_whole( + workload.volume, workload.spacing + ), + "kimimaro_block": lambda: run_kimimaro_block( + workload.volume, workload.block_shape, workload.spacing + ), + } + results, samples = timed(backends, warmup, repeats) + bio_whole = results["bioimage_whole"] + bio_block_result = results["bioimage_block"] + kimi_whole = results["kimimaro_whole"] + kimi_block_result = results["kimimaro_block"] + bio_block = bio_block_result.forest + kimi_block = kimi_block_result.forest + graphs = { + "bioimage_whole": bio_whole, + "bioimage_block_raw": bio_block_result.raw, + "bioimage_block": bio_block, + "kimimaro_whole": kimi_whole, + "kimimaro_block_raw": kimi_block_result.raw, + "kimimaro_block": kimi_block, + } + output = { + "name": workload.name, + "shape": list(workload.volume.shape), + "block_shape": list(workload.block_shape), + "spacing": list(workload.spacing), + "number_of_blocks": int(np.prod(np.ceil( + np.asarray(workload.volume.shape) / np.asarray(workload.block_shape) + ))), + "timing_gate": workload.timing_gate, + "timing_seconds": { + name: { + "median": median(values), + "samples": values, + } + for name, values in samples.items() + }, + "bioimage_block_phases_seconds": bio_block_result.phases, + "kimimaro_block_phases_seconds": kimi_block_result.phases, + "graphs": { + name: graph_stats(graph, workload.spacing) + for name, graph in graphs.items() + }, + "distances": { + "bioimage_whole_to_kimimaro_whole": distance_metrics( + bio_whole, kimi_whole, workload.spacing + ), + "bioimage_block_to_kimimaro_block": distance_metrics( + bio_block, kimi_block, workload.spacing + ), + "bioimage_block_to_kimimaro_whole": distance_metrics( + bio_block, kimi_whole, workload.spacing + ), + "kimimaro_block_to_kimimaro_whole": distance_metrics( + kimi_block, kimi_whole, workload.spacing + ), + }, + "critical_points": { + "endpoints": critical_point_distances( + bio_block, + kimi_block, + lambda degree: degree == 1, + workload.spacing, + ), + "branchpoints": critical_point_distances( + bio_block, + kimi_block, + lambda degree: degree >= 3, + workload.spacing, + ), + }, + "radii": { + "merged_common_vertices": radius_metrics(bio_block, kimi_block), + "identical_anchor_coordinates": anchor_radius_errors( + workload.volume, + workload.block_shape, + workload.spacing, + bio_block_result, + ), + }, + "exact_graph": { + "bioimage_block_to_kimimaro_block": graphs_exact( + bio_block, kimi_block + ), + "bioimage_block_to_kimimaro_whole": graphs_exact( + bio_block, kimi_whole + ), + }, + } + return output + + +def gate_failures(result): + failures = [] + name = result["name"] + graphs = result["graphs"] + bio = graphs["bioimage_block"] + kimi = graphs["kimimaro_block"] + whole = graphs["kimimaro_whole"] + if bio["components"] != whole["components"]: + failures.append(f"{name}: bioimage block component count differs from whole") + if bio["cycle_rank"] != 0: + failures.append(f"{name}: final bioimage graph contains cycles") + if ( + name == "thin-line" and + not result["exact_graph"]["bioimage_block_to_kimimaro_whole"] + ): + failures.append(f"{name}: expected exact graph agreement") + if ( + name == "oblique-anisotropic-64-8" and + not result["exact_graph"]["bioimage_block_to_kimimaro_block"] + ): + failures.append(f"{name}: expected exact blocked graph agreement") + anchor_error = result["radii"]["identical_anchor_coordinates"] + if anchor_error["compared"] != anchor_error["expected"]: + failures.append( + f"{name}: compared radii at {anchor_error['compared']}/" + f"{anchor_error['expected']} interface anchors" + ) + if anchor_error["compared"] and anchor_error["maximum_absolute_error"] > 1e-5: + failures.append( + f"{name}: interface radius error {anchor_error['maximum_absolute_error']:.6g}" + ) + if name != "thin-line": + cable_ratio = bio["cable_length"] / kimi["cable_length"] + if abs(cable_ratio - 1.0) > 0.10: + failures.append( + f"{name}: blocked cable ratio to Kimimaro is {cable_ratio:.4f}" + ) + direct = result["distances"]["bioimage_block_to_kimimaro_block"] + if direct["p95"] > 2.0 or direct["hausdorff"] > 5.0: + failures.append( + f"{name}: direct p95/Hausdorff is " + f"{direct['p95']:.3f}/{direct['hausdorff']:.3f}" + ) + bio_whole_error = abs(bio["cable_length"] / whole["cable_length"] - 1.0) + kimi_whole_error = abs(kimi["cable_length"] / whole["cable_length"] - 1.0) + if bio_whole_error > kimi_whole_error + 0.10: + failures.append(f"{name}: cable error exceeds layered reference gate") + bio_to_whole = result["distances"]["bioimage_block_to_kimimaro_whole"] + kimi_to_whole = result["distances"]["kimimaro_block_to_kimimaro_whole"] + if bio_to_whole["p95"] > kimi_to_whole["p95"] + 1.0: + failures.append(f"{name}: p95 error exceeds layered reference gate") + bio_endpoint_error = abs(bio["endpoints"] - whole["endpoints"]) + kimi_endpoint_error = abs(kimi["endpoints"] - whole["endpoints"]) + if bio_endpoint_error > kimi_endpoint_error + 2: + failures.append(f"{name}: endpoint error exceeds layered reference gate") + if result["timing_gate"]: + timings = result["timing_seconds"] + bio_time = timings["bioimage_block"]["median"] + if bio_time > timings["kimimaro_whole"]["median"]: + failures.append(f"{name}: blocked bioimage is slower than whole Kimimaro") + if bio_time > timings["kimimaro_block"]["median"]: + failures.append(f"{name}: blocked bioimage is slower than blocked Kimimaro") + phases = result["bioimage_block_phases_seconds"] + if (phases["merge"] + phases["forest"]) > 0.10 * sum(phases.values()): + failures.append(f"{name}: merge plus forest exceeds 10% of block time") + return failures + + +def print_summary(result): + timing = result["timing_seconds"] + graphs = result["graphs"] + direct = result["distances"]["bioimage_block_to_kimimaro_block"] + bio = graphs["bioimage_block"] + kimi = graphs["kimimaro_block"] + whole = graphs["kimimaro_whole"] + anchor = result["radii"]["identical_anchor_coordinates"] + print( + f"{result['name']}: " + f"bio-block={1000 * timing['bioimage_block']['median']:.2f} ms, " + f"kimi-whole={1000 * timing['kimimaro_whole']['median']:.2f} ms, " + f"kimi-block={1000 * timing['kimimaro_block']['median']:.2f} ms" + ) + print( + f" length bio/kimi-block/whole=" + f"{bio['cable_length']:.2f}/{kimi['cable_length']:.2f}/" + f"{whole['cable_length']:.2f}; endpoints=" + f"{bio['endpoints']}/{kimi['endpoints']}/{whole['endpoints']}; " + f"direct p95/H={direct['p95']:.2f}/{direct['hausdorff']:.2f}; " + f"anchor max radius error={anchor['maximum_absolute_error']}" + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--quick", action="store_true") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--repeats", type=int, default=7) + parser.add_argument("--check", action="store_true") + parser.add_argument("--json", type=Path) + args = parser.parse_args() + if args.warmup < 0 or args.repeats < 1: + parser.error("--warmup must be >= 0 and --repeats must be >= 1") + + versions = {} + for package in ("bioimage-cpp", "kimimaro", "edt", "numpy", "scipy"): + try: + versions[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + versions[package] = None + report = { + "versions": versions, + "python": sys.version, + "default_spacing": SPACING, + "parameters": PARAMETERS, + "threads": 1, + "warmup": args.warmup, + "repeats": args.repeats, + "workloads": [], + } + failures = [] + for workload in workloads(args.quick): + result = compare_workload(workload, args.warmup, args.repeats) + report["workloads"].append(result) + print_summary(result) + failures.extend(gate_failures(result)) + report["gate_failures"] = failures + if args.json: + args.json.write_text(json.dumps(report, indent=2) + "\n") + if failures: + print("Quality gate failures:", file=sys.stderr) + for failure in failures: + print(f" - {failure}", file=sys.stderr) + if args.check: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/skeleton/detail/components.hxx b/include/bioimage_cpp/skeleton/detail/components.hxx index 5b95576..8794663 100644 --- a/include/bioimage_cpp/skeleton/detail/components.hxx +++ b/include/bioimage_cpp/skeleton/detail/components.hxx @@ -78,6 +78,7 @@ struct ComponentDescriptor { template struct ComponentSet { + std::array input_shape{}; std::vector> runs; std::vector row_offsets; std::vector> components; @@ -88,8 +89,29 @@ struct ComponentSet { struct PreparedTeasarComponent { std::vector padded_shape; std::vector padded_mask; + // Optional EDT-only mask. The path domain remains padded_mask; this mask + // extends foreground across artificial processing-block cuts so those + // cuts do not become false object boundaries. + std::vector distance_mask; std::array input_origin{}; std::size_t foreground_count = 0; + std::vector required_target_voxels; + std::size_t required_root_voxel = std::numeric_limits::max(); +}; + +struct OpenBlockFaces { + // low/high pairs for axes 0, 1, and 2. + std::array values{}; + + bool operator()(const std::size_t axis, const bool high) const { + return values[2 * axis + static_cast(high)]; + } + + bool any() const { + return std::any_of(values.begin(), values.end(), [](const bool value) { + return value; + }); + } }; inline bool intervals_within_one( @@ -135,6 +157,11 @@ ComponentSet extract_components( ); ComponentSet result; + result.input_shape = { + static_cast(z_size), + static_cast(y_size), + static_cast(x_size), + }; result.row_offsets.resize( checked_add_size(row_count, 1, "component row offsets overflow size_t"), 0 @@ -392,7 +419,9 @@ std::size_t padded_component_volume( template PreparedTeasarComponent prepare_component( const ComponentSet &components, - const std::size_t component_id + const std::size_t component_id, + const std::vector> &required_targets = {}, + const OpenBlockFaces *open_faces = nullptr ) { const auto &component = components.components.at(component_id); PreparedTeasarComponent prepared; @@ -441,7 +470,182 @@ PreparedTeasarComponent prepare_component( std::uint8_t{1} ); } + + // The zero halo in padded_mask is part of the TEASAR path-domain + // contract. For block calls, construct a separate distance-only mask by + // copying foreground through declared artificial cuts. The EDT does not + // treat array exteriors as background, so one copied layer is sufficient. + // Corners and edges are extended only if every participating face is open. + const auto input_size = checked_multiply_size( + checked_multiply_size( + static_cast(components.input_shape[0]), + static_cast(components.input_shape[1]), + "component input size overflows size_t" + ), + static_cast(components.input_shape[2]), + "component input size overflows size_t" + ); + const bool all_foreground = components.components.size() == 1 && + static_cast(component.voxel_count) == input_size; + if (open_faces != nullptr && open_faces->any() && !all_foreground) { + std::array extend{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + extend[2 * axis] = (*open_faces)(axis, false) && + component.begin[axis] == 0; + extend[2 * axis + 1] = (*open_faces)(axis, true) && + component.end[axis] == components.input_shape[axis]; + } + if (std::any_of(extend.begin(), extend.end(), [](const bool value) { + return value; + })) { + prepared.distance_mask = prepared.padded_mask; + const auto z_size = prepared.padded_shape[0]; + const auto y_size = prepared.padded_shape[1]; + const auto x_size = prepared.padded_shape[2]; + for (std::ptrdiff_t z = 0; z < z_size; ++z) { + for (std::ptrdiff_t y = 0; y < y_size; ++y) { + for (std::ptrdiff_t x = 0; x < x_size; ++x) { + const std::array coordinate{z, y, x}; + std::array source = coordinate; + bool is_halo = false; + bool permitted = true; + for (std::size_t axis = 0; axis < 3; ++axis) { + if (coordinate[axis] == 0) { + is_halo = true; + permitted = permitted && extend[2 * axis]; + source[axis] = 1; + } else if ( + coordinate[axis] == prepared.padded_shape[axis] - 1 + ) { + is_halo = true; + permitted = permitted && extend[2 * axis + 1]; + source[axis] = prepared.padded_shape[axis] - 2; + } + } + if (!is_halo || !permitted) { + continue; + } + const auto destination = static_cast( + z * static_cast(sz) + + y * static_cast(sy) + x + ); + const auto source_index = static_cast( + source[0] * static_cast(sz) + + source[1] * static_cast(sy) + source[2] + ); + prepared.distance_mask[destination] = + prepared.padded_mask[source_index]; + } + } + } + } + } + + prepared.required_target_voxels.reserve(required_targets.size()); + std::array required_root_coordinate{}; + bool have_required_root = false; + for (const auto &target : required_targets) { + const auto z = static_cast( + target[0] - component.begin[0] + 1 + ); + const auto y = static_cast( + target[1] - component.begin[1] + 1 + ); + const auto x = static_cast( + target[2] - component.begin[2] + 1 + ); + const auto target_voxel = z * sz + y * sy + x; + prepared.required_target_voxels.push_back(target_voxel); + bool on_open_face = false; + if (open_faces != nullptr) { + for (std::size_t axis = 0; axis < 3; ++axis) { + on_open_face = on_open_face || + (target[axis] == 0 && (*open_faces)(axis, false)) || + (target[axis] == components.input_shape[axis] - 1 && + (*open_faces)(axis, true)); + } + } + if ( + on_open_face && + (!have_required_root || target > required_root_coordinate) + ) { + have_required_root = true; + required_root_coordinate = target; + prepared.required_root_voxel = target_voxel; + } + } return prepared; } +template +struct ComponentTarget { + LabelT label{}; + std::array coordinate{}; +}; + +// Assign already validated foreground coordinates to the equality-connected +// run components discovered by extract_components. The O(number_of_runs) +// reverse lookup is allocated only for calls that actually carry targets, so +// ordinary TEASAR dispatch keeps its current memory profile. +template +std::vector>> +assign_targets_to_components( + const ComponentSet &components, + const std::vector> &targets +) { + std::vector>> assigned( + components.components.size() + ); + if (targets.empty()) { + return assigned; + } + + std::vector component_of_run( + components.runs.size(), std::numeric_limits::max() + ); + for (std::size_t component = 0; + component < components.components.size(); ++component) { + const auto &descriptor = components.components[component]; + for (std::size_t offset = 0; offset < descriptor.number_of_runs; ++offset) { + const auto run_id = components.component_run_ids[ + descriptor.run_offset + offset + ]; + component_of_run[run_id] = component; + } + } + + const auto y_size = static_cast(components.input_shape[1]); + for (const auto &target : targets) { + const auto row = static_cast(target.coordinate[0]) * y_size + + static_cast(target.coordinate[1]); + const auto begin = components.row_offsets[row]; + const auto end = components.row_offsets[row + 1]; + auto run = std::lower_bound( + components.runs.begin() + static_cast(begin), + components.runs.begin() + static_cast(end), + target.coordinate[2], + [](const ComponentRun &candidate, const std::ptrdiff_t x) { + return candidate.x_end < x; + } + ); + if ( + run == components.runs.begin() + static_cast(end) || + run->x_begin > target.coordinate[2] || run->label != target.label + ) { + throw std::runtime_error( + "validated required target was not found in a component run" + ); + } + const auto run_id = static_cast( + std::distance(components.runs.begin(), run) + ); + const auto component = component_of_run[run_id]; + if (component == std::numeric_limits::max()) { + throw std::runtime_error("required target component lookup is inconsistent"); + } + assigned[component].push_back(target.coordinate); + } + return assigned; +} + } // namespace bioimage_cpp::skeleton::detail diff --git a/include/bioimage_cpp/skeleton/distributed/border_targets.hxx b/include/bioimage_cpp/skeleton/distributed/border_targets.hxx new file mode 100644 index 0000000..73b0b39 --- /dev/null +++ b/include/bioimage_cpp/skeleton/distributed/border_targets.hxx @@ -0,0 +1,390 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/profile.hxx" +#include "bioimage_cpp/detail/threading.hxx" +#include "bioimage_cpp/distance/distance_transform.hxx" +#include "bioimage_cpp/skeleton/detail/components.hxx" +#include "bioimage_cpp/skeleton/teasar.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::skeleton::distributed { + +struct BlockFace { + std::size_t axis = 0; + bool high = false; + + auto operator<=>(const BlockFace &) const = default; +}; + +namespace detail_border { + +inline std::int64_t checked_global_coordinate( + const std::int64_t origin, + const std::ptrdiff_t local +) { + if (local < 0) { + throw std::runtime_error("local border coordinate became negative"); + } + if ( + static_cast(local) > + static_cast(std::numeric_limits::max()) + ) { + throw std::overflow_error("local border coordinate exceeds int64 range"); + } + const auto local64 = static_cast(local); + if (origin > std::numeric_limits::max() - local64) { + throw std::overflow_error("global border coordinate overflows int64"); + } + return origin + local64; +} + +inline void validate_common( + const std::vector &shape, + const std::array &spacing, + const std::vector &faces +) { + if (shape.size() != 3) { + throw std::invalid_argument("block input must have exactly three dimensions"); + } + for (std::size_t axis = 0; axis < 3; ++axis) { + if (shape[axis] < 0) { + throw std::invalid_argument("block shape entries must be non-negative"); + } + if (!(std::isfinite(spacing[axis]) && spacing[axis] > 0.0)) { + throw std::invalid_argument("spacing values must be positive and finite"); + } + } + for (const auto &face : faces) { + if (face.axis >= 3) { + throw std::invalid_argument("face axis must be in [0, 3)"); + } + } +} + +template +std::vector> border_targets_impl( + const ConstArrayView &input, + const LabelT background, + std::vector faces, + const std::array &origin, + const std::array &spacing, + const std::size_t number_of_threads +) { + static_assert(std::is_integral_v && !std::is_same_v); + validate_common(input.shape, spacing, faces); + std::sort(faces.begin(), faces.end()); + faces.erase(std::unique(faces.begin(), faces.end()), faces.end()); + + BIOIMAGE_PROFILE_INIT(profile) + std::vector> targets; + for (const auto &face : faces) { + if ( + input.shape[face.axis] == 0 || + input.shape[(face.axis + 1) % 3] == 0 || + input.shape[(face.axis + 2) % 3] == 0 + ) { + continue; + } + std::array face_axes{}; + std::size_t cursor = 0; + for (std::size_t axis = 0; axis < 3; ++axis) { + if (axis != face.axis) { + face_axes[cursor++] = axis; + } + } + const auto height = input.shape[face_axes[0]]; + const auto width = input.shape[face_axes[1]]; + const auto face_size = detail::checked_multiply_size( + static_cast(height), + static_cast(width), + "face size overflows size_t" + ); + std::vector packed(face_size); + std::array local{}; + local[face.axis] = face.high ? input.shape[face.axis] - 1 : 0; + for (std::ptrdiff_t first = 0; first < height; ++first) { + local[face_axes[0]] = first; + for (std::ptrdiff_t second = 0; second < width; ++second) { + local[face_axes[1]] = second; + const auto input_index = static_cast( + (local[0] * input.shape[1] + local[1]) * input.shape[2] + + local[2] + ); + auto value = input.data[input_index]; + if constexpr (Binary) { + value = value == background ? LabelT{0} : LabelT{1}; + } + packed[ + static_cast(first * width + second) + ] = value; + } + } + + ConstArrayView packed_view{ + packed.data(), {1, height, width}, {} + }; + auto components = [&] { + if constexpr (Binary) { + return detail::extract_components( + packed_view, LabelT{0}, profile + ); + } else { + return detail::extract_labeled_components( + packed_view, background, profile + ); + } + }(); + + for (const auto &component : components.components) { + const auto component_height = component.end[1] - component.begin[1]; + const auto component_width = component.end[2] - component.begin[2]; + const std::vector padded_shape{ + component_height + 2, component_width + 2 + }; + const auto padded_size = detail::checked_shape_size( + padded_shape, "padded face component size overflows size_t" + ); + std::vector component_mask(padded_size, 0); + const auto padded_width = padded_shape[1]; + long double first_sum = 0.0L; + long double second_sum = 0.0L; + std::size_t count = 0; + for (std::size_t offset = 0; + offset < component.number_of_runs; ++offset) { + const auto run_id = components.component_run_ids[ + component.run_offset + offset + ]; + const auto &run = components.runs[run_id]; + const auto row = run.y - component.begin[1] + 1; + for (auto x = run.x_begin; x <= run.x_end; ++x) { + const auto column = x - component.begin[2] + 1; + component_mask[static_cast( + row * padded_width + column + )] = 1; + first_sum += static_cast(run.y); + second_sum += static_cast(x); + ++count; + } + } + if (count == 0) { + throw std::runtime_error("face component has no foreground pixels"); + } + + std::vector distances(padded_size, 0.0f); + ConstArrayView component_view{ + component_mask.data(), padded_shape, {} + }; + ArrayView distance_view{ + distances.data(), padded_shape, {} + }; + distance::distance_transform( + component_view, + {spacing[face_axes[0]], spacing[face_axes[1]]}, + {distance_view, {}, {}}, + bioimage_cpp::detail::normalize_thread_count( + number_of_threads, count + ) + ); + + const auto mean_first = first_sum / static_cast(count); + const auto mean_second = second_sum / static_cast(count); + const auto face_center_first = + static_cast(height) / 2.0L; + const auto face_center_second = + static_cast(width) / 2.0L; + // Resolve half-grid centroids toward the centre of the full face. + // This makes the choice invariant when neighboring blocks view the + // shared plane from opposite sides. + const auto centroid_first = mean_first >= face_center_first + ? std::floor(mean_first) : std::ceil(mean_first); + const auto centroid_second = mean_second >= face_center_second + ? std::floor(mean_second) : std::ceil(mean_second); + float best_distance = -1.0f; + long double best_centroid_distance = + std::numeric_limits::infinity(); + long double best_face_center_distance = + std::numeric_limits::infinity(); + long double best_corner_distance = + std::numeric_limits::infinity(); + long double best_edge_distance = + std::numeric_limits::infinity(); + VoxelCoordinate best_coordinate{}; + bool have_best = false; + for (std::size_t offset = 0; + offset < component.number_of_runs; ++offset) { + const auto run_id = components.component_run_ids[ + component.run_offset + offset + ]; + const auto &run = components.runs[run_id]; + const auto row = run.y - component.begin[1] + 1; + for (auto x = run.x_begin; x <= run.x_end; ++x) { + const auto column = x - component.begin[2] + 1; + const auto edt = distances[static_cast( + row * padded_width + column + )]; + const auto first_delta = + (static_cast(run.y) - centroid_first) * + static_cast(spacing[face_axes[0]]); + const auto second_delta = + (static_cast(x) - centroid_second) * + static_cast(spacing[face_axes[1]]); + const auto centroid_distance = + first_delta * first_delta + second_delta * second_delta; + const auto face_first_delta = + (static_cast(run.y) - face_center_first) * + static_cast(spacing[face_axes[0]]); + const auto face_second_delta = + (static_cast(x) - face_center_second) * + static_cast(spacing[face_axes[1]]); + const auto face_center_distance = + face_first_delta * face_first_delta + + face_second_delta * face_second_delta; + const auto corner_distance = [&] { + long double result = + std::numeric_limits::infinity(); + for (const auto corner_first : { + -0.5L, + static_cast(height) - 0.5L, + }) { + for (const auto corner_second : { + -0.5L, + static_cast(width) - 0.5L, + }) { + const auto first = + (static_cast(run.y) - corner_first) * + static_cast(spacing[face_axes[0]]); + const auto second = + (static_cast(x) - corner_second) * + static_cast(spacing[face_axes[1]]); + result = std::min( + result, first * first + second * second + ); + } + } + return result; + }(); + const auto edge_distance = std::min({ + static_cast(spacing[face_axes[0]]) * + (static_cast(run.y) - 0.5L), + static_cast(spacing[face_axes[0]]) * + (static_cast(height) - 0.5L - + static_cast(run.y)), + static_cast(spacing[face_axes[1]]) * + (static_cast(x) - 0.5L), + static_cast(spacing[face_axes[1]]) * + (static_cast(width) - 0.5L - + static_cast(x)), + }); + local[face_axes[0]] = run.y; + local[face_axes[1]] = x; + VoxelCoordinate global{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + global[axis] = checked_global_coordinate( + origin[axis], local[axis] + ); + } + if ( + !have_best || edt > best_distance || + (edt == best_distance && + (centroid_distance < best_centroid_distance || + (centroid_distance == best_centroid_distance && + (face_center_distance < best_face_center_distance || + (face_center_distance == best_face_center_distance && + (corner_distance < best_corner_distance || + (corner_distance == best_corner_distance && + (edge_distance < best_edge_distance || + (edge_distance == best_edge_distance && + global < best_coordinate))))))))) + ) { + have_best = true; + best_distance = edt; + best_centroid_distance = centroid_distance; + best_face_center_distance = face_center_distance; + best_corner_distance = corner_distance; + best_edge_distance = edge_distance; + best_coordinate = global; + } + } + } + if (!have_best) { + throw std::runtime_error("failed to select a face target"); + } + targets.push_back({ + Binary ? LabelT{1} : component.label, + best_coordinate, + }); + } + } + + std::sort( + targets.begin(), targets.end(), + [](const auto &first, const auto &second) { + if (first.label != second.label) { + return first.label < second.label; + } + return first.coordinate < second.coordinate; + } + ); + targets.erase( + std::unique( + targets.begin(), targets.end(), + [](const auto &first, const auto &second) { + return first.label == second.label && + first.coordinate == second.coordinate; + } + ), + targets.end() + ); + BIOIMAGE_PROFILE_REPORT(profile) + return targets; +} + +} // namespace detail_border + +inline std::vector block_border_targets( + const ConstArrayView &mask, + std::vector faces, + const std::array &origin, + const std::array &spacing, + const std::size_t number_of_threads = 1 +) { + auto labeled = detail_border::border_targets_impl( + mask, std::uint8_t{0}, std::move(faces), origin, spacing, + number_of_threads + ); + std::vector targets; + targets.reserve(labeled.size()); + for (const auto &target : labeled) { + targets.push_back(target.coordinate); + } + return targets; +} + +template +std::vector> block_border_targets_labels( + const ConstArrayView &labels, + const LabelT background, + std::vector faces, + const std::array &origin, + const std::array &spacing, + const std::size_t number_of_threads = 1 +) { + return detail_border::border_targets_impl( + labels, background, std::move(faces), origin, spacing, + number_of_threads + ); +} + +} // namespace bioimage_cpp::skeleton::distributed diff --git a/include/bioimage_cpp/skeleton/distributed/merge.hxx b/include/bioimage_cpp/skeleton/distributed/merge.hxx new file mode 100644 index 0000000..90ba889 --- /dev/null +++ b/include/bioimage_cpp/skeleton/distributed/merge.hxx @@ -0,0 +1,413 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/skeleton/teasar.hxx" +#include "bioimage_cpp/util/union_find.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::skeleton::distributed { + +using LatticeSkeletonGraph = bioimage_cpp::skeleton::LatticeSkeletonGraph; + +struct LatticeSkeletonView { + ConstArrayView vertices; + ConstArrayView edges; + ConstArrayView radii; +}; + +inline void validate_lattice_skeleton( + const LatticeSkeletonGraph &graph, + const std::size_t fragment_index = 0 +) { + if (graph.radii.size() != graph.vertices.size()) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " radii length must match vertices" + ); + } + for (std::size_t vertex = 0; vertex < graph.radii.size(); ++vertex) { + if (!(std::isfinite(graph.radii[vertex]) && graph.radii[vertex] >= 0.0f)) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " radii must be finite and non-negative" + ); + } + } + const auto number_of_vertices = static_cast( + graph.vertices.size() + ); + for (const auto &edge : graph.edges) { + if (edge[0] >= number_of_vertices || edge[1] >= number_of_vertices) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " contains an edge endpoint outside its vertex range" + ); + } + } +} + +inline void validate_lattice_skeleton( + const LatticeSkeletonView &graph, + const std::size_t fragment_index = 0 +) { + if (graph.vertices.ndim() != 2 || graph.vertices.shape[1] != 3) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " vertices must have shape (n, 3)" + ); + } + if (graph.edges.ndim() != 2 || graph.edges.shape[1] != 2) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " edges must have shape (n, 2)" + ); + } + if (graph.radii.ndim() != 1 || graph.radii.shape[0] != graph.vertices.shape[0]) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " radii length must match vertices" + ); + } + const auto number_of_vertices = static_cast( + graph.vertices.shape[0] + ); + for (std::ptrdiff_t vertex = 0; vertex < graph.radii.shape[0]; ++vertex) { + const auto radius = graph.radii.data[vertex]; + if (!(std::isfinite(radius) && radius >= 0.0f)) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " radii must be finite and non-negative" + ); + } + } + for (std::ptrdiff_t edge = 0; edge < graph.edges.shape[0]; ++edge) { + if ( + graph.edges.data[edge * 2] >= number_of_vertices || + graph.edges.data[edge * 2 + 1] >= number_of_vertices + ) { + throw std::invalid_argument( + "fragment " + std::to_string(fragment_index) + + " contains an edge endpoint outside its vertex range" + ); + } + } +} + +namespace detail_merge { + +inline std::size_t number_of_vertices(const LatticeSkeletonGraph &graph) { + return graph.vertices.size(); +} + +inline std::size_t number_of_vertices(const LatticeSkeletonView &graph) { + return static_cast(graph.vertices.shape[0]); +} + +inline std::size_t number_of_edges(const LatticeSkeletonGraph &graph) { + return graph.edges.size(); +} + +inline std::size_t number_of_edges(const LatticeSkeletonView &graph) { + return static_cast(graph.edges.shape[0]); +} + +inline VoxelCoordinate vertex( + const LatticeSkeletonGraph &graph, + const std::size_t index +) { + return graph.vertices[index]; +} + +inline VoxelCoordinate vertex( + const LatticeSkeletonView &graph, + const std::size_t index +) { + return { + graph.vertices.data[index * 3], + graph.vertices.data[index * 3 + 1], + graph.vertices.data[index * 3 + 2], + }; +} + +inline float radius(const LatticeSkeletonGraph &graph, const std::size_t index) { + return graph.radii[index]; +} + +inline float radius(const LatticeSkeletonView &graph, const std::size_t index) { + return graph.radii.data[index]; +} + +inline std::array edge( + const LatticeSkeletonGraph &graph, + const std::size_t index +) { + return graph.edges[index]; +} + +inline std::array edge( + const LatticeSkeletonView &graph, + const std::size_t index +) { + return {graph.edges.data[index * 2], graph.edges.data[index * 2 + 1]}; +} + +template +LatticeSkeletonGraph merge_block_skeletons_impl( + const std::vector &fragments +) { + struct VertexRecord { + VoxelCoordinate coordinate{}; + std::size_t fragment = 0; + std::size_t old_vertex = 0; + float radius = 0.0f; + }; + + std::size_t total_vertices = 0; + std::size_t total_edges = 0; + for (std::size_t fragment = 0; fragment < fragments.size(); ++fragment) { + validate_lattice_skeleton(fragments[fragment], fragment); + if ( + number_of_vertices(fragments[fragment]) > + std::numeric_limits::max() - total_vertices + ) { + throw std::overflow_error("merged vertex count overflows size_t"); + } + total_vertices += number_of_vertices(fragments[fragment]); + if ( + number_of_edges(fragments[fragment]) > + std::numeric_limits::max() - total_edges + ) { + throw std::overflow_error("merged edge count overflows size_t"); + } + total_edges += number_of_edges(fragments[fragment]); + } + + std::vector records; + records.reserve(total_vertices); + for (std::size_t fragment = 0; fragment < fragments.size(); ++fragment) { + const auto &part = fragments[fragment]; + for (std::size_t vertex_id = 0; + vertex_id < number_of_vertices(part); ++vertex_id) { + records.push_back({ + vertex(part, vertex_id), fragment, vertex_id, + radius(part, vertex_id) + }); + } + } + std::sort( + records.begin(), records.end(), + [](const VertexRecord &first, const VertexRecord &second) { + if (first.coordinate != second.coordinate) { + return first.coordinate < second.coordinate; + } + if (first.fragment != second.fragment) { + return first.fragment < second.fragment; + } + return first.old_vertex < second.old_vertex; + } + ); + + std::vector> old_to_new(fragments.size()); + for (std::size_t fragment = 0; fragment < fragments.size(); ++fragment) { + old_to_new[fragment].resize(number_of_vertices(fragments[fragment])); + } + LatticeSkeletonGraph output; + output.vertices.reserve(records.size()); + output.radii.reserve(records.size()); + std::size_t begin = 0; + while (begin < records.size()) { + auto end = begin + 1; + float radius = records[begin].radius; + while ( + end < records.size() && + records[end].coordinate == records[begin].coordinate + ) { + radius = std::max(radius, records[end].radius); + ++end; + } + if (output.vertices.size() >= std::numeric_limits::max()) { + throw std::overflow_error("merged vertex id exceeds uint64 range"); + } + const auto new_vertex = static_cast( + output.vertices.size() + ); + output.vertices.push_back(records[begin].coordinate); + output.radii.push_back(radius); + for (auto record = begin; record < end; ++record) { + old_to_new[records[record].fragment][records[record].old_vertex] = + new_vertex; + } + begin = end; + } + + output.edges.reserve(total_edges); + for (std::size_t fragment = 0; fragment < fragments.size(); ++fragment) { + for (std::size_t edge_id = 0; + edge_id < number_of_edges(fragments[fragment]); ++edge_id) { + const auto old_edge = edge(fragments[fragment], edge_id); + auto first = old_to_new[fragment][ + static_cast(old_edge[0]) + ]; + auto second = old_to_new[fragment][ + static_cast(old_edge[1]) + ]; + if (first == second) { + continue; + } + if (second < first) { + std::swap(first, second); + } + output.edges.push_back({first, second}); + } + } + std::sort(output.edges.begin(), output.edges.end()); + output.edges.erase( + std::unique(output.edges.begin(), output.edges.end()), + output.edges.end() + ); + return output; +} + +} // namespace detail_merge + +inline LatticeSkeletonGraph merge_block_skeletons( + const std::vector &fragments +) { + return detail_merge::merge_block_skeletons_impl(fragments); +} + +inline LatticeSkeletonGraph merge_block_skeletons( + const std::vector &fragments +) { + return detail_merge::merge_block_skeletons_impl(fragments); +} + +inline LatticeSkeletonGraph minimum_spanning_forest( + const LatticeSkeletonGraph &graph, + const std::array &spacing +) { + validate_lattice_skeleton(graph); + for (const auto value : spacing) { + if (!(std::isfinite(value) && value > 0.0)) { + throw std::invalid_argument("spacing values must be positive and finite"); + } + } + struct WeightedEdge { + long double weight = 0.0L; + VoxelCoordinate first_coordinate{}; + VoxelCoordinate second_coordinate{}; + std::uint64_t first = 0; + std::uint64_t second = 0; + }; + std::vector weighted; + weighted.reserve(graph.edges.size()); + for (const auto &edge : graph.edges) { + if (edge[0] == edge[1]) { + continue; + } + auto first = edge[0]; + auto second = edge[1]; + auto first_coordinate = graph.vertices[static_cast(first)]; + auto second_coordinate = graph.vertices[static_cast(second)]; + if ( + second_coordinate < first_coordinate || + (second_coordinate == first_coordinate && second < first) + ) { + std::swap(first, second); + std::swap(first_coordinate, second_coordinate); + } + long double weight = 0.0L; + for (std::size_t axis = 0; axis < 3; ++axis) { + const auto delta = + (static_cast(first_coordinate[axis]) - + static_cast(second_coordinate[axis])) * + static_cast(spacing[axis]); + weight += delta * delta; + } + if (!std::isfinite(weight)) { + throw std::overflow_error("physical edge length overflowed"); + } + weighted.push_back({ + weight, first_coordinate, second_coordinate, first, second + }); + } + std::sort( + weighted.begin(), weighted.end(), + [](const WeightedEdge &first, const WeightedEdge &second) { + if (first.weight != second.weight) { + return first.weight < second.weight; + } + if (first.first_coordinate != second.first_coordinate) { + return first.first_coordinate < second.first_coordinate; + } + if (first.second_coordinate != second.second_coordinate) { + return first.second_coordinate < second.second_coordinate; + } + if (first.first != second.first) { + return first.first < second.first; + } + return first.second < second.second; + } + ); + + util::UnionFind union_find(graph.vertices.size()); + LatticeSkeletonGraph output; + output.vertices = graph.vertices; + output.radii = graph.radii; + output.edges.reserve(std::min(graph.edges.size(), graph.vertices.size())); + for (const auto &edge : weighted) { + const auto first_root = union_find.find(edge.first); + const auto second_root = union_find.find(edge.second); + if (first_root == second_root) { + continue; + } + union_find.unite_roots(first_root, second_root); + auto first = edge.first; + auto second = edge.second; + if (second < first) { + std::swap(first, second); + } + output.edges.push_back({first, second}); + } + std::sort(output.edges.begin(), output.edges.end()); + return output; +} + +inline SkeletonGraph lattice_to_physical( + const LatticeSkeletonGraph &graph, + const std::array &spacing +) { + validate_lattice_skeleton(graph); + for (const auto value : spacing) { + if (!(std::isfinite(value) && value > 0.0)) { + throw std::invalid_argument("spacing values must be positive and finite"); + } + } + SkeletonGraph output; + output.vertices.reserve(graph.vertices.size()); + output.radii = graph.radii; + output.edges = graph.edges; + for (const auto &coordinate : graph.vertices) { + std::array physical{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + physical[axis] = static_cast(coordinate[axis]) * spacing[axis]; + if (!std::isfinite(physical[axis])) { + throw std::overflow_error("physical skeleton coordinate overflowed"); + } + } + output.vertices.push_back(physical); + } + return output; +} + +} // namespace bioimage_cpp::skeleton::distributed diff --git a/include/bioimage_cpp/skeleton/teasar.hxx b/include/bioimage_cpp/skeleton/teasar.hxx index 6daccc1..500b694 100644 --- a/include/bioimage_cpp/skeleton/teasar.hxx +++ b/include/bioimage_cpp/skeleton/teasar.hxx @@ -39,12 +39,26 @@ struct TeasarOptions { std::size_t number_of_threads = 1; }; +using VoxelCoordinate = std::array; + struct SkeletonGraph { std::vector> vertices; std::vector> edges; std::vector radii; }; +struct LatticeSkeletonGraph { + std::vector vertices; + std::vector> edges; + std::vector radii; +}; + +template +struct LabeledVoxelTarget { + LabelT label{}; + VoxelCoordinate coordinate{}; +}; + // Kept public at the C++ level so development benchmarks can compare the // sequential implementations. The Python API always uses TeasarBackend::Auto. enum class TeasarBackend { @@ -134,7 +148,7 @@ inline void invalidation_bounds( // Skeletonize a binary 3D mask with the core TEASAR procedure. A non-empty // input must contain exactly one 26-connected foreground component. -inline SkeletonGraph teasar_dense_impl( +inline LatticeSkeletonGraph teasar_dense_impl( const ConstArrayView &mask, detail::PreparedTeasarComponent *prepared, const TeasarOptions &options, @@ -145,17 +159,23 @@ inline SkeletonGraph teasar_dense_impl( } BIOIMAGE_PROFILE_INIT(profile) - SkeletonGraph graph; + LatticeSkeletonGraph graph; std::size_t foreground_count = 0; std::array input_origin{0, 0, 0}; std::vector shape; std::vector padded_mask; + std::vector distance_mask; + std::vector required_targets; + std::size_t required_root = std::numeric_limits::max(); std::size_t first_foreground = std::numeric_limits::max(); if (prepared != nullptr) { foreground_count = prepared->foreground_count; input_origin = prepared->input_origin; shape = std::move(prepared->padded_shape); padded_mask = std::move(prepared->padded_mask); + distance_mask = std::move(prepared->distance_mask); + required_targets = std::move(prepared->required_target_voxels); + required_root = prepared->required_root_voxel; for (std::size_t index = 0; index < padded_mask.size(); ++index) { if (padded_mask[index] != 0) { first_foreground = index; @@ -212,7 +232,11 @@ inline SkeletonGraph teasar_dense_impl( std::vector dbf(n, 0.0f); { BIOIMAGE_PROFILE_SCOPE(profile, "distance_transform") - ConstArrayView padded_view{padded_mask.data(), shape, {}}; + const auto &distance_input = distance_mask.empty() + ? padded_mask : distance_mask; + ConstArrayView padded_view{ + distance_input.data(), shape, {} + }; ArrayView distances_view{dbf.data(), shape, {}}; distance::distance_transform( padded_view, @@ -228,28 +252,50 @@ inline SkeletonGraph teasar_dense_impl( distance::DijkstraCostMode::Physical, effective_threads, }; - distance::DijkstraResult first_field; distance::DijkstraResult root_field; std::size_t root = first_foreground; { BIOIMAGE_PROFILE_SCOPE(profile, "root_dijkstra") ConstArrayView padded_view{padded_mask.data(), shape, {}}; - first_field = distance::dijkstra_distance_field( - padded_view, {first_foreground}, physical_options + if (required_root != std::numeric_limits::max()) { + if (required_root >= n || padded_mask[required_root] == 0) { + throw std::runtime_error( + "prepared required root is not foreground" + ); + } + root = required_root; + } else { + auto first_field = distance::dijkstra_distance_field( + padded_view, {first_foreground}, physical_options + ); + for (std::size_t index = 0; index < n; ++index) { + if ( + padded_mask[index] != 0 && + !std::isfinite(first_field.distances[index]) + ) { + throw std::invalid_argument( + "mask foreground must contain exactly one 26-connected component" + ); + } + } + root = detail_teasar::farthest_foreground( + padded_mask, first_field.distances + ); + } + root_field = distance::dijkstra_distance_field( + padded_view, {root}, physical_options ); for (std::size_t index = 0; index < n; ++index) { - if (padded_mask[index] != 0 && !std::isfinite(first_field.distances[index])) { + if ( + padded_mask[index] != 0 && + !std::isfinite(root_field.distances[index]) + ) { throw std::invalid_argument( "mask foreground must contain exactly one 26-connected component" ); } } - root = detail_teasar::farthest_foreground(padded_mask, first_field.distances); - root_field = distance::dijkstra_distance_field( - padded_view, {root}, physical_options - ); } - std::vector().swap(first_field.distances); double dbf_max = 0.0; double daf_max = 0.0; @@ -291,9 +337,9 @@ inline SkeletonGraph teasar_dense_impl( ); const auto vertex_id = static_cast(graph.vertices.size()); graph.vertices.push_back({ - static_cast(coords[0] - 1 + input_origin[0]) * options.spacing[0], - static_cast(coords[1] - 1 + input_origin[1]) * options.spacing[1], - static_cast(coords[2] - 1 + input_origin[2]) * options.spacing[2], + static_cast(coords[0] - 1 + input_origin[0]), + static_cast(coords[1] - 1 + input_origin[1]), + static_cast(coords[2] - 1 + input_origin[2]), }); graph.radii.push_back(dbf[voxel]); vertex_of_voxel[voxel] = static_cast(vertex_id); @@ -308,21 +354,8 @@ inline SkeletonGraph teasar_dense_impl( const distance::DijkstraOptions node_options{ 3, {}, distance::DijkstraCostMode::Node, effective_threads }; - - while (active_count > 0) { - std::size_t target = std::numeric_limits::max(); - double target_distance = -1.0; - for (std::size_t index = 0; index < n; ++index) { - if (active[index] != 0 && root_field.distances[index] > target_distance) { - target = index; - target_distance = root_field.distances[index]; - } - } - if (target == std::numeric_limits::max()) { - throw std::runtime_error("TEASAR active-voxel accounting became inconsistent"); - } - - std::vector path; + std::vector path; + const auto trace_target = [&](const std::size_t target) { { BIOIMAGE_PROFILE_SCOPE(profile, "path_dijkstra") path = distance::dijkstra_path( @@ -377,6 +410,41 @@ inline SkeletonGraph teasar_dense_impl( } } } + }; + + for (const auto target : required_targets) { + if (target >= n || padded_mask[target] == 0) { + throw std::runtime_error("prepared required target is not foreground"); + } + } + std::sort( + required_targets.begin(), required_targets.end(), + [&](const std::size_t first, const std::size_t second) { + if (root_field.distances[first] != root_field.distances[second]) { + return root_field.distances[first] > root_field.distances[second]; + } + return first < second; + } + ); + for (const auto target : required_targets) { + if (vertex_of_voxel[target] < 0) { + trace_target(target); + } + } + + while (active_count > 0) { + std::size_t target = std::numeric_limits::max(); + double target_distance = -1.0; + for (std::size_t index = 0; index < n; ++index) { + if (active[index] != 0 && root_field.distances[index] > target_distance) { + target = index; + target_distance = root_field.distances[index]; + } + } + if (target == std::numeric_limits::max()) { + throw std::runtime_error("TEASAR active-voxel accounting became inconsistent"); + } + trace_target(target); } if (report_profile) { @@ -385,23 +453,25 @@ inline SkeletonGraph teasar_dense_impl( return graph; } -inline SkeletonGraph teasar_dense( +inline LatticeSkeletonGraph teasar_dense( const ConstArrayView &mask, const TeasarOptions &options = {} ) { return teasar_dense_impl(mask, nullptr, options, true); } -inline SkeletonGraph teasar_dense_prepared( +inline LatticeSkeletonGraph teasar_dense_prepared( detail::PreparedTeasarComponent prepared, const TeasarOptions &options ) { const ConstArrayView unused{}; - return teasar_dense_impl(unused, &prepared, options, false); + return teasar_dense_impl( + unused, &prepared, options, false + ); } template -inline SkeletonGraph teasar_compact_impl( +inline LatticeSkeletonGraph teasar_compact_impl( const ConstArrayView &mask, detail::PreparedTeasarComponent *prepared, const TeasarOptions &options, @@ -412,7 +482,7 @@ inline SkeletonGraph teasar_compact_impl( } BIOIMAGE_PROFILE_INIT(profile) - SkeletonGraph graph; + LatticeSkeletonGraph graph; std::array crop_begin{0, 0, 0}; std::array crop_end{0, 0, 0}; std::size_t foreground_count = 0; @@ -420,11 +490,17 @@ inline SkeletonGraph teasar_compact_impl( std::vector strides; std::size_t n = 0; std::vector padded_mask; + std::vector distance_mask; + std::vector required_targets; + std::size_t required_root = std::numeric_limits::max(); if (prepared != nullptr) { crop_begin = prepared->input_origin; foreground_count = prepared->foreground_count; shape = std::move(prepared->padded_shape); padded_mask = std::move(prepared->padded_mask); + distance_mask = std::move(prepared->distance_mask); + required_targets = std::move(prepared->required_target_voxels); + required_root = prepared->required_root_voxel; n = padded_mask.size(); strides = bioimage_cpp::detail::c_order_strides(shape); if (foreground_count == 0) { @@ -490,7 +566,11 @@ inline SkeletonGraph teasar_compact_impl( auto dbf = std::make_unique_for_overwrite(n); { BIOIMAGE_PROFILE_SCOPE(profile, "distance_transform") - ConstArrayView padded_view{padded_mask.data(), shape, {}}; + const auto &distance_input = distance_mask.empty() + ? padded_mask : distance_mask; + ConstArrayView padded_view{ + distance_input.data(), shape, {} + }; ArrayView distances_view{dbf.get(), shape, {}}; distance::distance_transform( padded_view, @@ -509,9 +589,16 @@ inline SkeletonGraph teasar_compact_impl( if (prepared == nullptr) { return teasar_dense(mask, options); } + detail::PreparedTeasarComponent dense_prepared; + dense_prepared.padded_shape = std::move(shape); + dense_prepared.padded_mask = std::move(padded_mask); + dense_prepared.distance_mask = std::move(distance_mask); + dense_prepared.input_origin = crop_begin; + dense_prepared.foreground_count = foreground_count; + dense_prepared.required_target_voxels = std::move(required_targets); + dense_prepared.required_root_voxel = required_root; return teasar_dense_prepared( - {shape, std::move(padded_mask), crop_begin, foreground_count}, - options + std::move(dense_prepared), options ); } } @@ -533,31 +620,60 @@ inline SkeletonGraph teasar_compact_impl( } detail::CompactDijkstraWorkspace dijkstra_workspace; - std::vector first_field; std::vector root_field; std::uint32_t root = 0; { BIOIMAGE_PROFILE_SCOPE(profile, "root_dijkstra") - detail::compact_physical_distance_field( - domain, 0, dijkstra_workspace, first_field - ); - Distance farthest_distance = Distance{-1}; - for (std::uint32_t node = 0; node < domain.size(); ++node) { - if (!std::isfinite(first_field[node])) { - throw std::invalid_argument( - "mask foreground must contain exactly one 26-connected component" + if (required_root != std::numeric_limits::max()) { + if (required_root > std::numeric_limits::max()) { + throw std::runtime_error( + "prepared required root exceeds compact index range" ); } - if (first_field[node] > farthest_distance) { - root = node; - farthest_distance = first_field[node]; + const auto it = std::lower_bound( + domain.compact_to_full.begin(), domain.compact_to_full.end(), + static_cast(required_root) + ); + if ( + it == domain.compact_to_full.end() || + *it != static_cast(required_root) + ) { + throw std::runtime_error( + "prepared required root is not foreground" + ); + } + root = static_cast( + std::distance(domain.compact_to_full.begin(), it) + ); + } else { + std::vector first_field; + detail::compact_physical_distance_field( + domain, 0, dijkstra_workspace, first_field + ); + Distance farthest_distance = Distance{-1}; + for (std::uint32_t node = 0; node < domain.size(); ++node) { + if (!std::isfinite(first_field[node])) { + throw std::invalid_argument( + "mask foreground must contain exactly one 26-connected component" + ); + } + if (first_field[node] > farthest_distance) { + root = node; + farthest_distance = first_field[node]; + } } } detail::compact_physical_distance_field( domain, root, dijkstra_workspace, root_field ); + for (const auto distance : root_field) { + if (!std::isfinite(distance)) { + throw std::invalid_argument( + "mask foreground must contain exactly one 26-connected component" + ); + } + } } - std::vector().swap(first_field); Distance daf_max = Distance{0}; for (std::uint32_t node = 0; node < domain.size(); ++node) { @@ -599,9 +715,9 @@ inline SkeletonGraph teasar_compact_impl( ); const auto vertex_id = static_cast(graph.vertices.size()); graph.vertices.push_back({ - static_cast(coords[0] - 1 + crop_begin[0]) * options.spacing[0], - static_cast(coords[1] - 1 + crop_begin[1]) * options.spacing[1], - static_cast(coords[2] - 1 + crop_begin[2]) * options.spacing[2], + static_cast(coords[0] - 1 + crop_begin[0]), + static_cast(coords[1] - 1 + crop_begin[1]), + static_cast(coords[2] - 1 + crop_begin[2]), }); graph.radii.push_back(compact_dbf[node]); vertex_of_node[node] = static_cast(vertex_id); @@ -615,70 +731,7 @@ inline SkeletonGraph teasar_compact_impl( detail::RowIntervalUnion invalidated_rows( n / static_cast(shape[2]), shape[2] ); - std::vector ordered_targets; - std::size_t ordered_target_cursor = 0; - std::size_t linear_target_selections = 0; - const auto linear_target_limit = std::max( - 16, std::bit_width(domain.size()) - ); - constexpr std::size_t ordered_target_minimum_nodes = 1U << 16; - const bool allow_ordered_targets = - domain.size() >= ordered_target_minimum_nodes; - bool targets_ordered = false; - while (active_count > 0) { - auto target = detail::kNoCompactNode; - { - BIOIMAGE_PROFILE_SCOPE(profile, "target_selection") - if ( - !targets_ordered && - (!allow_ordered_targets || - linear_target_selections < linear_target_limit) - ) { - Distance target_distance = Distance{-1}; - for (std::uint32_t node = 0; node < domain.size(); ++node) { - const auto full = domain.compact_to_full[node]; - if (active[full] != 0 && root_field[node] > target_distance) { - target = node; - target_distance = root_field[node]; - } - } - ++linear_target_selections; - } else { - if (!targets_ordered) { - ordered_targets.reserve(active_count); - for (std::uint32_t node = 0; node < domain.size(); ++node) { - if (active[domain.compact_to_full[node]] != 0) { - ordered_targets.push_back(node); - } - } - std::sort( - ordered_targets.begin(), ordered_targets.end(), - [&](const std::uint32_t first, const std::uint32_t second) { - if (root_field[first] != root_field[second]) { - return root_field[first] > root_field[second]; - } - return first < second; - } - ); - targets_ordered = true; - } - while ( - ordered_target_cursor < ordered_targets.size() && - active[ - domain.compact_to_full[ordered_targets[ordered_target_cursor]] - ] == 0 - ) { - ++ordered_target_cursor; - } - if (ordered_target_cursor < ordered_targets.size()) { - target = ordered_targets[ordered_target_cursor++]; - } - } - } - if (target == detail::kNoCompactNode) { - throw std::runtime_error("TEASAR active-voxel accounting became inconsistent"); - } - + const auto trace_target = [&](const std::uint32_t target) { { BIOIMAGE_PROFILE_SCOPE(profile, "path_dijkstra") detail::compact_node_cost_path( @@ -747,6 +800,107 @@ inline SkeletonGraph teasar_compact_impl( } } } + }; + + std::vector compact_required_targets; + compact_required_targets.reserve(required_targets.size()); + for (const auto full_target : required_targets) { + if (full_target >= n) { + throw std::runtime_error("prepared required target is out of bounds"); + } + const auto it = std::lower_bound( + domain.compact_to_full.begin(), domain.compact_to_full.end(), + static_cast(full_target) + ); + if ( + it == domain.compact_to_full.end() || + *it != static_cast(full_target) + ) { + throw std::runtime_error("prepared required target is not foreground"); + } + compact_required_targets.push_back(static_cast( + std::distance(domain.compact_to_full.begin(), it) + )); + } + std::sort( + compact_required_targets.begin(), compact_required_targets.end(), + [&](const std::uint32_t first, const std::uint32_t second) { + if (root_field[first] != root_field[second]) { + return root_field[first] > root_field[second]; + } + return first < second; + } + ); + for (const auto target : compact_required_targets) { + if (vertex_of_node[target] < 0) { + trace_target(target); + } + } + + std::vector ordered_targets; + std::size_t ordered_target_cursor = 0; + std::size_t linear_target_selections = 0; + const auto linear_target_limit = std::max( + 16, std::bit_width(domain.size()) + ); + constexpr std::size_t ordered_target_minimum_nodes = 1U << 16; + const bool allow_ordered_targets = + domain.size() >= ordered_target_minimum_nodes; + bool targets_ordered = false; + while (active_count > 0) { + auto target = detail::kNoCompactNode; + { + BIOIMAGE_PROFILE_SCOPE(profile, "target_selection") + if ( + !targets_ordered && + (!allow_ordered_targets || + linear_target_selections < linear_target_limit) + ) { + Distance target_distance = Distance{-1}; + for (std::uint32_t node = 0; node < domain.size(); ++node) { + const auto full = domain.compact_to_full[node]; + if (active[full] != 0 && root_field[node] > target_distance) { + target = node; + target_distance = root_field[node]; + } + } + ++linear_target_selections; + } else { + if (!targets_ordered) { + ordered_targets.reserve(active_count); + for (std::uint32_t node = 0; node < domain.size(); ++node) { + if (active[domain.compact_to_full[node]] != 0) { + ordered_targets.push_back(node); + } + } + std::sort( + ordered_targets.begin(), ordered_targets.end(), + [&](const std::uint32_t first, const std::uint32_t second) { + if (root_field[first] != root_field[second]) { + return root_field[first] > root_field[second]; + } + return first < second; + } + ); + targets_ordered = true; + } + while ( + ordered_target_cursor < ordered_targets.size() && + active[ + domain.compact_to_full[ordered_targets[ordered_target_cursor]] + ] == 0 + ) { + ++ordered_target_cursor; + } + if (ordered_target_cursor < ordered_targets.size()) { + target = ordered_targets[ordered_target_cursor++]; + } + } + } + if (target == detail::kNoCompactNode) { + throw std::runtime_error("TEASAR active-voxel accounting became inconsistent"); + } + trace_target(target); } if (report_profile) { @@ -756,7 +910,7 @@ inline SkeletonGraph teasar_compact_impl( } template -inline SkeletonGraph teasar_compact( +inline LatticeSkeletonGraph teasar_compact( const ConstArrayView &mask, const TeasarOptions &options ) { @@ -766,7 +920,7 @@ inline SkeletonGraph teasar_compact( } template -inline SkeletonGraph teasar_compact_prepared( +inline LatticeSkeletonGraph teasar_compact_prepared( detail::PreparedTeasarComponent prepared, const TeasarOptions &options ) { @@ -779,8 +933,8 @@ inline SkeletonGraph teasar_compact_prepared( namespace detail_teasar { inline void append_skeleton_graph( - SkeletonGraph &destination, - SkeletonGraph &&source + LatticeSkeletonGraph &destination, + LatticeSkeletonGraph &&source ) { if (destination.vertices.size() > std::numeric_limits::max()) { throw std::overflow_error("skeleton vertex offset exceeds uint64 range"); @@ -809,8 +963,8 @@ inline void append_skeleton_graph( } } -inline SkeletonGraph assemble_skeleton_graphs( - std::vector &graphs, +inline LatticeSkeletonGraph assemble_skeleton_graphs( + std::vector &graphs, const std::vector &component_ids ) { std::size_t vertices = 0; @@ -825,7 +979,7 @@ inline SkeletonGraph assemble_skeleton_graphs( "assembled skeleton edge count overflows size_t" ); } - SkeletonGraph result; + LatticeSkeletonGraph result; result.vertices.reserve(vertices); result.radii.reserve(vertices); result.edges.reserve(edges); @@ -835,6 +989,24 @@ inline SkeletonGraph assemble_skeleton_graphs( return result; } +inline SkeletonGraph lattice_to_physical( + LatticeSkeletonGraph graph, + const std::array &spacing +) { + SkeletonGraph result; + result.vertices.reserve(graph.vertices.size()); + for (const auto &coordinate : graph.vertices) { + result.vertices.push_back({ + static_cast(coordinate[0]) * spacing[0], + static_cast(coordinate[1]) * spacing[1], + static_cast(coordinate[2]) * spacing[2], + }); + } + result.edges = std::move(graph.edges); + result.radii = std::move(graph.radii); + return result; +} + template std::string component_context( const detail::ComponentDescriptor &component, @@ -959,13 +1131,15 @@ std::vector component_thread_budgets( } template -std::vector skeletonize_components( +std::vector skeletonize_components( const detail::ComponentSet &components, const TeasarOptions &options, - const bool include_label_in_errors + const bool include_label_in_errors, + const std::vector>> *required_targets = nullptr, + const detail::OpenBlockFaces *open_faces = nullptr ) { const auto count = components.components.size(); - std::vector results(count); + std::vector results(count); if (count == 0) { return results; } @@ -996,7 +1170,12 @@ std::vector skeletonize_components( const std::size_t local_budget ) { try { - auto prepared = detail::prepare_component(components, component_id); + auto prepared = required_targets == nullptr + ? detail::prepare_component(components, component_id) + : detail::prepare_component( + components, component_id, required_targets->at(component_id), + open_faces + ); auto local_options = options; local_options.number_of_threads = local_budget; results[component_id] = teasar_compact_prepared< @@ -1052,18 +1231,28 @@ inline SkeletonGraph teasar_with_backend( const TeasarOptions &options, const TeasarBackend backend ) { + LatticeSkeletonGraph lattice; switch (backend) { case TeasarBackend::Auto: case TeasarBackend::CompactOnTheFlyFloat64: - return teasar_compact( + lattice = teasar_compact( mask, options ); + break; case TeasarBackend::DenseFloat64: - return teasar_dense(mask, options); + lattice = teasar_dense(mask, options); + break; case TeasarBackend::CompactCsrFloat64: - return teasar_compact(mask, options); + lattice = teasar_compact( + mask, options + ); + break; + default: + throw std::invalid_argument("invalid TEASAR backend"); } - throw std::invalid_argument("invalid TEASAR backend"); + return detail_teasar::lattice_to_physical( + std::move(lattice), options.spacing + ); } inline SkeletonGraph teasar( @@ -1073,7 +1262,7 @@ inline SkeletonGraph teasar( detail_teasar::validate_options(mask, options); BIOIMAGE_PROFILE_INIT(profile) auto components = detail::extract_binary_components(mask, profile); - std::vector results; + std::vector results; { BIOIMAGE_PROFILE_SCOPE(profile, "component_teasar") results = detail_teasar::skeletonize_components( @@ -1082,7 +1271,85 @@ inline SkeletonGraph teasar( } std::vector component_ids(results.size()); std::iota(component_ids.begin(), component_ids.end(), std::size_t{0}); - SkeletonGraph output; + LatticeSkeletonGraph output; + { + BIOIMAGE_PROFILE_SCOPE(profile, "forest_assembly") + output = detail_teasar::assemble_skeleton_graphs(results, component_ids); + } + BIOIMAGE_PROFILE_REPORT(profile) + return detail_teasar::lattice_to_physical( + std::move(output), options.spacing + ); +} + +inline LatticeSkeletonGraph teasar_block( + const ConstArrayView &mask, + std::vector required_targets, + const detail::OpenBlockFaces &open_faces, + const TeasarOptions &options = {} +) { + detail_teasar::validate_options(mask, options); + std::vector> local_targets; + local_targets.reserve(required_targets.size()); + for (std::size_t row = 0; row < required_targets.size(); ++row) { + const auto &target = required_targets[row]; + std::array coordinate{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + if ( + target[axis] < 0 || + static_cast(target[axis]) >= + static_cast(mask.shape[axis]) + ) { + throw std::invalid_argument( + "required_targets row " + std::to_string(row) + + " is out of bounds at axis " + std::to_string(axis) + ); + } + coordinate[axis] = static_cast(target[axis]); + } + const auto flat = static_cast( + (coordinate[0] * mask.shape[1] + coordinate[1]) * mask.shape[2] + + coordinate[2] + ); + if (mask.data[flat] == 0) { + throw std::invalid_argument( + "required_targets row " + std::to_string(row) + + " must lie on foreground" + ); + } + local_targets.push_back({std::uint8_t{1}, coordinate}); + } + std::sort( + local_targets.begin(), local_targets.end(), + [](const auto &first, const auto &second) { + return first.coordinate < second.coordinate; + } + ); + local_targets.erase( + std::unique( + local_targets.begin(), local_targets.end(), + [](const auto &first, const auto &second) { + return first.coordinate == second.coordinate; + } + ), + local_targets.end() + ); + + BIOIMAGE_PROFILE_INIT(profile) + auto components = detail::extract_binary_components(mask, profile); + auto component_targets = detail::assign_targets_to_components( + components, local_targets + ); + std::vector results; + { + BIOIMAGE_PROFILE_SCOPE(profile, "component_teasar") + results = detail_teasar::skeletonize_components( + components, options, false, &component_targets, &open_faces + ); + } + std::vector component_ids(results.size()); + std::iota(component_ids.begin(), component_ids.end(), std::size_t{0}); + LatticeSkeletonGraph output; { BIOIMAGE_PROFILE_SCOPE(profile, "forest_assembly") output = detail_teasar::assemble_skeleton_graphs(results, component_ids); diff --git a/include/bioimage_cpp/skeleton/teasar_labels.hxx b/include/bioimage_cpp/skeleton/teasar_labels.hxx index 20dffd5..b329e03 100644 --- a/include/bioimage_cpp/skeleton/teasar_labels.hxx +++ b/include/bioimage_cpp/skeleton/teasar_labels.hxx @@ -5,6 +5,10 @@ #include #include +#include +#include +#include +#include #include #include #include @@ -17,6 +21,12 @@ struct LabeledSkeleton { SkeletonGraph skeleton; }; +template +struct LabeledLatticeSkeleton { + LabelT label{}; + LatticeSkeletonGraph skeleton; +}; + template std::vector> teasar_labels( const ConstArrayView &labels, @@ -31,7 +41,7 @@ std::vector> teasar_labels( auto components = detail::extract_labeled_components( labels, background, profile ); - std::vector results; + std::vector results; { BIOIMAGE_PROFILE_SCOPE(profile, "component_teasar") results = detail_teasar::skeletonize_components( @@ -56,6 +66,133 @@ std::vector> teasar_labels( ); std::vector> output; + { + BIOIMAGE_PROFILE_SCOPE(profile, "forest_assembly") + std::size_t begin = 0; + while (begin < component_ids.size()) { + const auto label = components.components[component_ids[begin]].label; + auto end = begin + 1; + while ( + end < component_ids.size() && + components.components[component_ids[end]].label == label + ) { + ++end; + } + std::vector label_components( + component_ids.begin() + static_cast(begin), + component_ids.begin() + static_cast(end) + ); + output.push_back({ + label, + detail_teasar::lattice_to_physical( + detail_teasar::assemble_skeleton_graphs( + results, label_components + ), + options.spacing + ), + }); + begin = end; + } + } + BIOIMAGE_PROFILE_REPORT(profile) + return output; +} + +template +std::vector> teasar_labels_block( + const ConstArrayView &labels, + const LabelT background, + std::vector> required_targets, + const detail::OpenBlockFaces &open_faces, + const TeasarOptions &options = {} +) { + static_assert(std::is_integral_v && !std::is_same_v); + ConstArrayView shape_only{nullptr, labels.shape, {}}; + detail_teasar::validate_options(shape_only, options); + + std::vector> local_targets; + local_targets.reserve(required_targets.size()); + for (std::size_t row = 0; row < required_targets.size(); ++row) { + const auto &target = required_targets[row]; + std::array coordinate{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + if ( + target.coordinate[axis] < 0 || + static_cast(target.coordinate[axis]) >= + static_cast(labels.shape[axis]) + ) { + throw std::invalid_argument( + "required_targets row " + std::to_string(row) + + " is out of bounds at axis " + std::to_string(axis) + ); + } + coordinate[axis] = static_cast( + target.coordinate[axis] + ); + } + const auto flat = static_cast( + (coordinate[0] * labels.shape[1] + coordinate[1]) * labels.shape[2] + + coordinate[2] + ); + if (target.label == background || labels.data[flat] != target.label) { + throw std::invalid_argument( + "required_targets row " + std::to_string(row) + + " does not match its semantic label" + ); + } + local_targets.push_back({target.label, coordinate}); + } + std::sort( + local_targets.begin(), local_targets.end(), + [](const auto &first, const auto &second) { + if (first.label != second.label) { + return first.label < second.label; + } + return first.coordinate < second.coordinate; + } + ); + local_targets.erase( + std::unique( + local_targets.begin(), local_targets.end(), + [](const auto &first, const auto &second) { + return first.label == second.label && + first.coordinate == second.coordinate; + } + ), + local_targets.end() + ); + + BIOIMAGE_PROFILE_INIT(profile) + auto components = detail::extract_labeled_components( + labels, background, profile + ); + auto component_targets = detail::assign_targets_to_components( + components, local_targets + ); + std::vector results; + { + BIOIMAGE_PROFILE_SCOPE(profile, "component_teasar") + results = detail_teasar::skeletonize_components( + components, options, true, &component_targets, &open_faces + ); + } + + std::vector component_ids(results.size()); + std::iota(component_ids.begin(), component_ids.end(), std::size_t{0}); + std::sort( + component_ids.begin(), component_ids.end(), + [&](const std::size_t first, const std::size_t second) { + const auto &first_component = components.components[first]; + const auto &second_component = components.components[second]; + if (first_component.label != second_component.label) { + return first_component.label < second_component.label; + } + return first_component.first_flat_index < + second_component.first_flat_index; + } + ); + + std::vector> output; { BIOIMAGE_PROFILE_SCOPE(profile, "forest_assembly") std::size_t begin = 0; diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx index 77ef8c8..1c8b56e 100644 --- a/src/bindings/module.cxx +++ b/src/bindings/module.cxx @@ -30,6 +30,7 @@ NB_MODULE(_core, m) { bioimage_cpp::bindings::bind_mesh(m); bioimage_cpp::bindings::bind_segmentation(m); bioimage_cpp::bindings::bind_skeleton(m); + bioimage_cpp::bindings::bind_skeleton_distributed(m); bioimage_cpp::bindings::bind_transformation(m); bioimage_cpp::bindings::bind_util(m); bioimage_cpp::bindings::bind_utils(m); diff --git a/src/bindings/skeleton.hxx b/src/bindings/skeleton.hxx index 5b52f93..c596b26 100644 --- a/src/bindings/skeleton.hxx +++ b/src/bindings/skeleton.hxx @@ -5,5 +5,6 @@ namespace bioimage_cpp::bindings { void bind_skeleton(nanobind::module_ &m); +void bind_skeleton_distributed(nanobind::module_ &m); } // namespace bioimage_cpp::bindings diff --git a/src/bindings/skeleton_distributed.cxx b/src/bindings/skeleton_distributed.cxx new file mode 100644 index 0000000..ca05585 --- /dev/null +++ b/src/bindings/skeleton_distributed.cxx @@ -0,0 +1,537 @@ +#include "skeleton.hxx" +#include "ndarray.hxx" + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/skeleton/distributed/border_targets.hxx" +#include "bioimage_cpp/skeleton/distributed/merge.hxx" +#include "bioimage_cpp/skeleton/teasar.hxx" +#include "bioimage_cpp/skeleton/teasar_labels.hxx" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +using UInt8Input = nb::ndarray; +using ConstInt64Array = nb::ndarray; +using ConstUInt64Array = nb::ndarray; +using ConstFloatArray = nb::ndarray; + +std::vector array_shape(const auto &array) { + std::vector shape(array.ndim()); + for (std::size_t axis = 0; axis < array.ndim(); ++axis) { + shape[axis] = static_cast(array.shape(axis)); + } + return shape; +} + +std::array origin_array( + const std::vector &origin +) { + if (origin.size() != 3) { + throw std::invalid_argument("origin must contain exactly three values"); + } + return {origin[0], origin[1], origin[2]}; +} + +std::array spacing_array(const std::vector &spacing) { + if (spacing.size() != 3) { + throw std::invalid_argument("spacing must contain exactly three values"); + } + return {spacing[0], spacing[1], spacing[2]}; +} + +std::vector face_vector( + const std::vector &axes, + const std::vector &high +) { + if (axes.size() != high.size()) { + throw std::invalid_argument("face axes and sides must have equal length"); + } + std::vector faces; + faces.reserve(axes.size()); + for (std::size_t index = 0; index < axes.size(); ++index) { + if (high[index] > 1) { + throw std::invalid_argument("face side flags must be zero or one"); + } + faces.push_back({axes[index], high[index] != 0}); + } + return faces; +} + +skeleton::detail::OpenBlockFaces open_face_policy( + const std::vector &axes, + const std::vector &high +) { + skeleton::detail::OpenBlockFaces output; + for (const auto &face : face_vector(axes, high)) { + output.values[2 * face.axis + static_cast(face.high)] = true; + } + return output; +} + +auto coordinate_array(const std::vector &coordinates) { + auto output = detail::make_array_for_overwrite( + {coordinates.size(), 3} + ); + for (std::size_t row = 0; row < coordinates.size(); ++row) { + for (std::size_t axis = 0; axis < 3; ++axis) { + output.data()[row * 3 + axis] = coordinates[row][axis]; + } + } + return output; +} + +auto edge_array( + const std::vector> &edges +) { + auto output = detail::make_array_for_overwrite( + {edges.size(), 2} + ); + for (std::size_t row = 0; row < edges.size(); ++row) { + output.data()[row * 2] = edges[row][0]; + output.data()[row * 2 + 1] = edges[row][1]; + } + return output; +} + +nb::tuple lattice_graph_to_tuple( + const skeleton::distributed::LatticeSkeletonGraph &graph +) { + return nb::make_tuple( + coordinate_array(graph.vertices), + edge_array(graph.edges), + detail::copy_vector_to_array(graph.radii) + ); +} + +nb::tuple physical_graph_to_tuple(const skeleton::SkeletonGraph &graph) { + auto vertices = detail::make_array_for_overwrite( + {graph.vertices.size(), 3} + ); + for (std::size_t row = 0; row < graph.vertices.size(); ++row) { + for (std::size_t axis = 0; axis < 3; ++axis) { + vertices.data()[row * 3 + axis] = graph.vertices[row][axis]; + } + } + return nb::make_tuple( + vertices, edge_array(graph.edges), detail::copy_vector_to_array(graph.radii) + ); +} + +std::vector global_targets_to_local( + ConstInt64Array targets, + const std::array &origin, + const std::vector &shape, + const std::string &argument_name +) { + if (targets.ndim() != 2 || targets.shape(1) != 3) { + throw std::invalid_argument(argument_name + " must have shape (n, 3)"); + } + std::vector local(targets.shape(0)); + for (std::size_t row = 0; row < targets.shape(0); ++row) { + for (std::size_t axis = 0; axis < 3; ++axis) { + const auto global = targets.data()[row * 3 + axis]; + if (global < origin[axis]) { + throw std::invalid_argument( + argument_name + " row " + std::to_string(row) + + " lies below the block origin" + ); + } + const auto difference = static_cast(global) - + static_cast(origin[axis]); + if (difference >= static_cast(shape[axis])) { + throw std::invalid_argument( + argument_name + " row " + std::to_string(row) + + " lies outside the block" + ); + } + local[row][axis] = static_cast(difference); + } + } + return local; +} + +std::int64_t checked_globalize( + const std::int64_t local, + const std::int64_t origin +) { + if (local < 0 || origin > std::numeric_limits::max() - local) { + throw std::overflow_error("global skeleton coordinate overflows int64"); + } + return origin + local; +} + +skeleton::distributed::LatticeSkeletonGraph globalize_graph( + const skeleton::LatticeSkeletonGraph &graph, + const std::array &origin +) { + skeleton::distributed::LatticeSkeletonGraph output; + output.vertices.reserve(graph.vertices.size()); + output.edges = graph.edges; + output.radii = graph.radii; + for (const auto &local : graph.vertices) { + skeleton::VoxelCoordinate global{}; + for (std::size_t axis = 0; axis < 3; ++axis) { + global[axis] = checked_globalize(local[axis], origin[axis]); + } + output.vertices.push_back(global); + } + return output; +} + +skeleton::TeasarOptions teasar_options( + const std::vector &spacing, + const double scale, + const double constant, + const double pdrf_scale, + const double pdrf_exponent, + const std::size_t number_of_threads +) { + const auto values = spacing_array(spacing); + return { + values, scale, constant, pdrf_scale, pdrf_exponent, + number_of_threads + }; +} + +auto block_border_targets_uint8( + UInt8Input mask, + const std::vector &axes, + const std::vector &high, + const std::vector &origin, + const std::vector &spacing, + const std::size_t number_of_threads +) { + ConstArrayView view{mask.data(), array_shape(mask), {}}; + std::vector targets; + { + nb::gil_scoped_release release; + targets = skeleton::distributed::block_border_targets( + view, face_vector(axes, high), origin_array(origin), + spacing_array(spacing), number_of_threads + ); + } + return coordinate_array(targets); +} + +template +nb::dict block_border_targets_labels_t( + nb::ndarray labels, + const LabelT background, + const std::vector &axes, + const std::vector &high, + const std::vector &origin, + const std::vector &spacing, + const std::size_t number_of_threads +) { + ConstArrayView view{labels.data(), array_shape(labels), {}}; + std::vector> targets; + { + nb::gil_scoped_release release; + targets = skeleton::distributed::block_border_targets_labels( + view, background, face_vector(axes, high), origin_array(origin), + spacing_array(spacing), number_of_threads + ); + } + nb::dict output; + std::size_t begin = 0; + while (begin < targets.size()) { + const auto label = targets[begin].label; + auto end = begin + 1; + while (end < targets.size() && targets[end].label == label) { + ++end; + } + std::vector coordinates; + coordinates.reserve(end - begin); + for (auto index = begin; index < end; ++index) { + coordinates.push_back(targets[index].coordinate); + } + if constexpr (std::is_signed_v) { + output[nb::int_(static_cast(label))] = + coordinate_array(coordinates); + } else { + output[nb::int_(static_cast(label))] = + coordinate_array(coordinates); + } + begin = end; + } + return output; +} + +nb::tuple block_teasar_uint8( + UInt8Input mask, + ConstInt64Array targets, + const std::vector &open_axes, + const std::vector &open_high, + const std::vector &origin, + const std::vector &spacing, + const double scale, + const double constant, + const double pdrf_scale, + const double pdrf_exponent, + const std::size_t number_of_threads +) { + const auto shape = array_shape(mask); + const auto origin_values = origin_array(origin); + auto local_targets = global_targets_to_local( + targets, origin_values, shape, "required_targets" + ); + ConstArrayView view{mask.data(), shape, {}}; + skeleton::LatticeSkeletonGraph result; + { + nb::gil_scoped_release release; + result = skeleton::teasar_block( + view, std::move(local_targets), + open_face_policy(open_axes, open_high), + teasar_options( + spacing, scale, constant, pdrf_scale, pdrf_exponent, + number_of_threads + ) + ); + } + return lattice_graph_to_tuple(globalize_graph(result, origin_values)); +} + +template +nb::dict block_teasar_labels_t( + nb::ndarray labels, + const LabelT background, + nb::dict target_map, + const std::vector &open_axes, + const std::vector &open_high, + const std::vector &origin, + const std::vector &spacing, + const double scale, + const double constant, + const double pdrf_scale, + const double pdrf_exponent, + const std::size_t number_of_threads +) { + const auto shape = array_shape(labels); + const auto origin_values = origin_array(origin); + std::vector> local_targets; + for (auto [key, value] : target_map) { + const auto label = nb::cast(key); + const auto array = nb::cast(value); + auto coordinates = global_targets_to_local( + array, origin_values, shape, "required_targets" + ); + for (const auto &coordinate : coordinates) { + local_targets.push_back({label, coordinate}); + } + } + ConstArrayView view{labels.data(), shape, {}}; + std::vector> results; + { + nb::gil_scoped_release release; + results = skeleton::teasar_labels_block( + view, background, std::move(local_targets), + open_face_policy(open_axes, open_high), + teasar_options( + spacing, scale, constant, pdrf_scale, pdrf_exponent, + number_of_threads + ) + ); + } + nb::dict output; + for (const auto &result : results) { + auto graph = globalize_graph(result.skeleton, origin_values); + if constexpr (std::is_signed_v) { + output[nb::int_(static_cast(result.label))] = + lattice_graph_to_tuple(graph); + } else { + output[nb::int_(static_cast(result.label))] = + lattice_graph_to_tuple(graph); + } + } + return output; +} + +skeleton::distributed::LatticeSkeletonGraph fragment_from_handle( + nb::handle handle, + const std::size_t fragment_index +) { + const auto tuple = nb::cast(handle); + if (tuple.size() != 3) { + throw std::invalid_argument("each fragment must be a 3-tuple"); + } + const auto vertices = nb::cast(tuple[0]); + const auto edges = nb::cast(tuple[1]); + const auto radii = nb::cast(tuple[2]); + if (vertices.ndim() != 2 || vertices.shape(1) != 3) { + throw std::invalid_argument("fragment vertices must have shape (n, 3)"); + } + if (edges.ndim() != 2 || edges.shape(1) != 2) { + throw std::invalid_argument("fragment edges must have shape (n, 2)"); + } + if (radii.ndim() != 1 || radii.shape(0) != vertices.shape(0)) { + throw std::invalid_argument("fragment radii must have shape (n_vertices,)"); + } + skeleton::distributed::LatticeSkeletonGraph graph; + graph.vertices.resize(vertices.shape(0)); + graph.radii.assign(radii.data(), radii.data() + radii.shape(0)); + graph.edges.resize(edges.shape(0)); + for (std::size_t row = 0; row < vertices.shape(0); ++row) { + for (std::size_t axis = 0; axis < 3; ++axis) { + graph.vertices[row][axis] = vertices.data()[row * 3 + axis]; + } + } + for (std::size_t row = 0; row < edges.shape(0); ++row) { + graph.edges[row] = {edges.data()[row * 2], edges.data()[row * 2 + 1]}; + } + skeleton::distributed::validate_lattice_skeleton(graph, fragment_index); + return graph; +} + +struct LatticeFragmentArrays { + ConstInt64Array vertices; + ConstUInt64Array edges; + ConstFloatArray radii; +}; + +LatticeFragmentArrays fragment_arrays_from_handle(nb::handle handle) { + const auto tuple = nb::cast(handle); + if (tuple.size() != 3) { + throw std::invalid_argument("each fragment must be a 3-tuple"); + } + return { + nb::cast(tuple[0]), + nb::cast(tuple[1]), + nb::cast(tuple[2]), + }; +} + +skeleton::distributed::LatticeSkeletonView fragment_view( + const LatticeFragmentArrays &arrays +) { + return { + {arrays.vertices.data(), array_shape(arrays.vertices), {}}, + {arrays.edges.data(), array_shape(arrays.edges), {}}, + {arrays.radii.data(), array_shape(arrays.radii), {}}, + }; +} + +nb::tuple merge_block_skeletons_binding(nb::list fragments) { + std::vector arrays; + arrays.reserve(fragments.size()); + for (std::size_t index = 0; index < fragments.size(); ++index) { + arrays.push_back(fragment_arrays_from_handle(fragments[index])); + } + std::vector views; + views.reserve(arrays.size()); + for (const auto &fragment : arrays) { + views.push_back(fragment_view(fragment)); + } + skeleton::distributed::LatticeSkeletonGraph result; + { + nb::gil_scoped_release release; + result = skeleton::distributed::merge_block_skeletons(views); + } + return lattice_graph_to_tuple(result); +} + +nb::tuple minimum_spanning_forest_binding( + nb::tuple fragment, + const std::vector &spacing +) { + auto graph = fragment_from_handle(fragment, 0); + skeleton::distributed::LatticeSkeletonGraph result; + { + nb::gil_scoped_release release; + result = skeleton::distributed::minimum_spanning_forest( + graph, spacing_array(spacing) + ); + } + return lattice_graph_to_tuple(result); +} + +nb::tuple lattice_to_physical_binding( + nb::tuple fragment, + const std::vector &spacing +) { + auto graph = fragment_from_handle(fragment, 0); + skeleton::SkeletonGraph result; + { + nb::gil_scoped_release release; + result = skeleton::distributed::lattice_to_physical( + graph, spacing_array(spacing) + ); + } + return physical_graph_to_tuple(result); +} + +} // namespace + +void bind_skeleton_distributed(nb::module_ &m) { + m.def( + "_block_border_targets_uint8", &block_border_targets_uint8, + nb::arg("mask"), nb::arg("axes"), nb::arg("high"), + nb::arg("origin"), nb::arg("spacing"), nb::arg("n_threads") + ); + m.def( + "_block_teasar_uint8", &block_teasar_uint8, + nb::arg("mask"), nb::arg("required_targets"), + nb::arg("open_axes"), nb::arg("open_high"), nb::arg("origin"), + nb::arg("spacing"), nb::arg("scale"), nb::arg("constant"), + nb::arg("pdrf_scale"), nb::arg("pdrf_exponent"), nb::arg("n_threads") + ); + +#define BIC_BIND_BLOCK_LABELS(name, type) \ + m.def( \ + "_block_border_targets_labels_" name, \ + &block_border_targets_labels_t, \ + nb::arg("labels"), nb::arg("background"), nb::arg("axes"), \ + nb::arg("high"), nb::arg("origin"), nb::arg("spacing"), \ + nb::arg("n_threads") \ + ); \ + m.def( \ + "_block_teasar_labels_" name, \ + &block_teasar_labels_t, \ + nb::arg("labels"), nb::arg("background"), \ + nb::arg("required_targets"), nb::arg("open_axes"), \ + nb::arg("open_high"), nb::arg("origin"), \ + nb::arg("spacing"), nb::arg("scale"), nb::arg("constant"), \ + nb::arg("pdrf_scale"), nb::arg("pdrf_exponent"), \ + nb::arg("n_threads") \ + ) + + BIC_BIND_BLOCK_LABELS("uint8", std::uint8_t); + BIC_BIND_BLOCK_LABELS("uint16", std::uint16_t); + BIC_BIND_BLOCK_LABELS("uint32", std::uint32_t); + BIC_BIND_BLOCK_LABELS("uint64", std::uint64_t); + BIC_BIND_BLOCK_LABELS("int32", std::int32_t); + BIC_BIND_BLOCK_LABELS("int64", std::int64_t); + +#undef BIC_BIND_BLOCK_LABELS + + m.def( + "_merge_block_skeletons", &merge_block_skeletons_binding, + nb::arg("fragments") + ); + m.def( + "_minimum_spanning_forest", &minimum_spanning_forest_binding, + nb::arg("fragment"), nb::arg("spacing") + ); + m.def( + "_lattice_to_physical", &lattice_to_physical_binding, + nb::arg("fragment"), nb::arg("spacing") + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/skeleton/__init__.py b/src/bioimage_cpp/skeleton/__init__.py index ef5aae9..5c80fb8 100644 --- a/src/bioimage_cpp/skeleton/__init__.py +++ b/src/bioimage_cpp/skeleton/__init__.py @@ -202,4 +202,7 @@ def teasar_labels( return run(labels_c, labels_array.dtype.type(background_value), *options) -__all__ = ["skeleton_to_graph", "teasar", "teasar_labels"] +from . import distributed + + +__all__ = ["distributed", "skeleton_to_graph", "teasar", "teasar_labels"] diff --git a/src/bioimage_cpp/skeleton/distributed.py b/src/bioimage_cpp/skeleton/distributed.py new file mode 100644 index 0000000..65ef80a --- /dev/null +++ b/src/bioimage_cpp/skeleton/distributed.py @@ -0,0 +1,396 @@ +"""Low-level primitives for blockwise TEASAR skeletonization and stitching. + +The caller owns block iteration, storage, scheduling, retries, and reduction +orchestration. A processing block should contain its non-overlapping core plus +one voxel on every high side that has a neighbor. The resulting shared plane is +passed to :func:`block_border_targets`; its global targets are then mandatory +rails for :func:`block_teasar` in both neighboring blocks. The same faces must +be supplied as ``open_faces`` so their artificial cuts do not collapse the +distance-to-boundary radius. + +Block artifacts are ``(vertices, edges, radii)`` tuples. Unlike ordinary +TEASAR output, ``vertices`` contains global ``int64`` lattice coordinates. +This makes :func:`merge_block_skeletons` an exact canonical set union and lets +its output feed another hierarchical merge without floating-point recovery. +Call :func:`lattice_to_physical` only after the final merge and optional cycle +removal. +""" + +from __future__ import annotations + +import operator +from collections.abc import Mapping, Sequence +from typing import TypeAlias + +import numpy as np + +from .. import _core +from .._validation import strict_index +from ..distance._distance import _as_binary_input, _normalize_sampling, _normalize_threads +from . import _normalize_teasar_options + + +SkeletonFragment: TypeAlias = tuple[np.ndarray, np.ndarray, np.ndarray] + +_LABEL_DTYPES = ( + np.dtype("uint8"), + np.dtype("uint16"), + np.dtype("uint32"), + np.dtype("uint64"), + np.dtype("int32"), + np.dtype("int64"), +) + +_BORDER_TARGET_LABELS = { + dtype: getattr(_core, f"_block_border_targets_labels_{dtype.name}") + for dtype in _LABEL_DTYPES +} +_BLOCK_TEASAR_LABELS = { + dtype: getattr(_core, f"_block_teasar_labels_{dtype.name}") + for dtype in _LABEL_DTYPES +} + + +def _normalize_origin(origin) -> list[int]: + try: + values = [operator.index(value) for value in origin] + except TypeError as error: + raise TypeError("origin must contain exactly three integers") from error + if len(values) != 3: + raise ValueError(f"origin must contain exactly three values, got {len(values)}") + info = np.iinfo(np.int64) + if any(value < info.min or value > info.max for value in values): + raise ValueError("origin values must fit int64") + return values + + +def _normalize_faces(faces) -> tuple[list[int], list[int]]: + axes: list[int] = [] + high: list[int] = [] + try: + entries = list(faces) + except TypeError as error: + raise TypeError("faces must be a sequence of (axis, side) pairs") from error + for index, entry in enumerate(entries): + try: + axis, side = entry + except (TypeError, ValueError) as error: + raise ValueError(f"faces[{index}] must be an (axis, side) pair") from error + try: + axis_value = operator.index(axis) + except TypeError as error: + raise TypeError(f"faces[{index}] axis must be an integer") from error + if axis_value < 0 or axis_value >= 3: + raise ValueError(f"faces[{index}] axis must be in [0, 3)") + if side not in ("low", "high"): + raise ValueError(f"faces[{index}] side must be 'low' or 'high'") + axes.append(axis_value) + high.append(side == "high") + return axes, high + + +def _normalize_integer_coordinates(values, name: str) -> np.ndarray: + array = np.asarray(values) + if array.size == 0: + if not ( + (array.ndim == 1 and array.shape == (0,)) + or (array.ndim == 2 and array.shape[1] == 3) + ): + raise ValueError(f"{name} must have shape (n, 3), got shape={array.shape}") + return np.empty((0, 3), dtype=np.int64) + if array.ndim != 2 or array.shape[1] != 3: + raise ValueError(f"{name} must have shape (n, 3), got shape={array.shape}") + if not np.issubdtype(array.dtype, np.integer): + raise TypeError(f"{name} must contain integers, got dtype={array.dtype}") + info = np.iinfo(np.int64) + if np.issubdtype(array.dtype, np.unsignedinteger): + if np.any(array > info.max): + raise ValueError(f"{name} values must fit int64") + elif np.any(array < info.min) or np.any(array > info.max): + raise ValueError(f"{name} values must fit int64") + return np.ascontiguousarray(array, dtype=np.int64) + + +def _normalize_labels(labels, function: str): + array = np.asarray(labels) + if array.ndim != 3: + raise ValueError(f"{function}: labels must have ndim 3, got ndim={array.ndim}") + if array.dtype not in _LABEL_DTYPES: + supported = ", ".join(str(dtype) for dtype in _LABEL_DTYPES) + raise TypeError( + f"{function}: labels must have one of native-endian dtypes " + f"({supported}), got dtype={array.dtype}" + ) + return np.ascontiguousarray(array) + + +def _normalize_background(background, dtype: np.dtype): + value = strict_index(background, "background") + info = np.iinfo(dtype) + if value < info.min or value > info.max: + raise ValueError(f"background={value} is outside the range of dtype {dtype}") + return dtype.type(value) + + +def block_border_targets( + mask: np.ndarray, + faces, + *, + origin=(0, 0, 0), + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), + number_of_threads: int = 1, +) -> np.ndarray: + """Select deterministic global stitch targets on requested block faces. + + ``faces`` contains ``(axis, "low"|"high")`` pairs. One target is selected + for every 8-connected foreground patch using an anisotropic 2D distance + transform. Plateaus are resolved by component centroid, face centre, + corner, edge, and finally global coordinate. The sorted unique result has + shape ``(T, 3)`` and dtype ``int64``. + """ + function = "block_border_targets" + binary = _as_binary_input(mask, function) + if binary.ndim != 3: + raise ValueError(f"{function}: mask must have ndim 3, got ndim={binary.ndim}") + axes, high = _normalize_faces(faces) + return _core._block_border_targets_uint8( + binary, + axes, + high, + _normalize_origin(origin), + _normalize_sampling(spacing, 3, function, name="spacing"), + _normalize_threads(number_of_threads, function), + ) + + +def block_border_targets_labels( + labels: np.ndarray, + faces, + *, + origin=(0, 0, 0), + background: int = 0, + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), + number_of_threads: int = 1, +) -> dict[int, np.ndarray]: + """Select deterministic global stitch targets per semantic label. + + Face connectivity requires equal non-background label values. The returned + dictionary preserves original integer labels and stores a sorted global + ``int64 (T, 3)`` coordinate array for each label. + """ + function = "block_border_targets_labels" + array = _normalize_labels(labels, function) + axes, high = _normalize_faces(faces) + return _BORDER_TARGET_LABELS[array.dtype]( + array, + _normalize_background(background, array.dtype), + axes, + high, + _normalize_origin(origin), + _normalize_sampling(spacing, 3, function, name="spacing"), + _normalize_threads(number_of_threads, function), + ) + + +def block_teasar( + mask: np.ndarray, + *, + open_faces, + origin=(0, 0, 0), + required_targets=None, + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), + scale: float = 1.5, + constant: float = 0.0, + pdrf_scale: float = 100000.0, + pdrf_exponent: float = 4.0, + number_of_threads: int = 1, +) -> SkeletonFragment: + """Skeletonize a binary processing block into a global lattice fragment. + + ``open_faces`` declares artificial processing-block cuts. Their foreground + is extended only for the distance-to-boundary transform; paths remain + confined to real input voxels. ``required_targets`` contains global + coordinates and is normally the union of these faces' targets. One target + on an open face becomes a deterministic component root and every other + target is forced onto a rail. + """ + function = "block_teasar" + binary = _as_binary_input(mask, function) + if binary.ndim != 3: + raise ValueError(f"{function}: mask must have ndim 3, got ndim={binary.ndim}") + targets = _normalize_integer_coordinates( + [] if required_targets is None else required_targets, + "required_targets", + ) + open_axes, open_high = _normalize_faces(open_faces) + options = _normalize_teasar_options( + function, spacing, scale, constant, pdrf_scale, pdrf_exponent, + number_of_threads, + ) + return _core._block_teasar_uint8( + binary, targets, open_axes, open_high, _normalize_origin(origin), *options + ) + + +def block_teasar_labels( + labels: np.ndarray, + *, + open_faces, + origin=(0, 0, 0), + required_targets: Mapping[int, np.ndarray] | None = None, + background: int = 0, + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), + scale: float = 1.5, + constant: float = 0.0, + pdrf_scale: float = 100000.0, + pdrf_exponent: float = 4.0, + number_of_threads: int = 1, +) -> dict[int, SkeletonFragment]: + """Skeletonize a labeled processing block into global lattice fragments. + + ``open_faces`` declares artificial processing-block cuts for DBF purposes. + ``required_targets`` maps original labels to global coordinate arrays. A + coordinate must contain exactly its mapping key in ``labels``. One open-face + target roots each affected component. The result has one global lattice + forest per original non-background label. + """ + function = "block_teasar_labels" + array = _normalize_labels(labels, function) + if required_targets is None: + target_map = {} + elif not isinstance(required_targets, Mapping): + raise TypeError("required_targets must be a mapping from labels to coordinates") + else: + target_map = {} + info = np.iinfo(array.dtype) + for key, coordinates in required_targets.items(): + label = strict_index(key, "required target label") + if label < info.min or label > info.max: + raise ValueError( + f"required target label {label} is outside dtype {array.dtype}" + ) + target_map[int(label)] = _normalize_integer_coordinates( + coordinates, f"required_targets[{label}]" + ) + options = _normalize_teasar_options( + function, spacing, scale, constant, pdrf_scale, pdrf_exponent, + number_of_threads, + ) + open_axes, open_high = _normalize_faces(open_faces) + return _BLOCK_TEASAR_LABELS[array.dtype]( + array, + _normalize_background(background, array.dtype), + target_map, + open_axes, + open_high, + _normalize_origin(origin), + *options, + ) + + +def _normalize_fragment(fragment, name: str = "fragment") -> SkeletonFragment: + try: + vertices, edges, radii = fragment + except (TypeError, ValueError) as error: + raise ValueError(f"{name} must be a (vertices, edges, radii) tuple") from error + vertices_array = _normalize_integer_coordinates(vertices, f"{name} vertices") + + edges_array = np.asarray(edges) + if edges_array.size == 0: + if not ( + (edges_array.ndim == 1 and edges_array.shape == (0,)) + or (edges_array.ndim == 2 and edges_array.shape[1] == 2) + ): + raise ValueError(f"{name} edges must have shape (n, 2)") + edges_array = np.empty((0, 2), dtype=np.uint64) + else: + if edges_array.ndim != 2 or edges_array.shape[1] != 2: + raise ValueError(f"{name} edges must have shape (n, 2)") + if not np.issubdtype(edges_array.dtype, np.integer): + raise TypeError(f"{name} edges must contain integers") + if np.issubdtype(edges_array.dtype, np.signedinteger) and np.any(edges_array < 0): + raise ValueError(f"{name} edges must be non-negative") + edges_array = np.ascontiguousarray(edges_array, dtype=np.uint64) + + radii_array = np.asarray(radii) + if radii_array.ndim != 1 or radii_array.shape[0] != vertices_array.shape[0]: + raise ValueError(f"{name} radii must have shape (n_vertices,)") + if not np.issubdtype(radii_array.dtype, np.floating): + raise TypeError(f"{name} radii must have a floating dtype") + radii_array = np.ascontiguousarray(radii_array, dtype=np.float32) + return vertices_array, edges_array, radii_array + + +def merge_block_skeletons(fragments) -> SkeletonFragment: + """Exactly consolidate global lattice fragments for one semantic object. + + Equal coordinates are unified, duplicate radii reduce by maximum, and + remapped edges are canonicalized, sorted, and deduplicated. Isolated + vertices are retained. The operation is associative, commutative, and + idempotent, so its output is directly reusable in a reduction tree. + """ + normalized = [ + _normalize_fragment(fragment, f"fragments[{index}]") + for index, fragment in enumerate(fragments) + ] + return _core._merge_block_skeletons(normalized) + + +def merge_block_skeleton_maps( + fragment_maps, +) -> dict[int, SkeletonFragment]: + """Consolidate block fragment dictionaries without mixing semantic labels.""" + maps = list(fragment_maps) + if any(not isinstance(mapping, Mapping) for mapping in maps): + raise TypeError("fragment_maps must contain mappings") + labels = sorted({strict_index(label, "fragment label") for mapping in maps for label in mapping}) + return { + label: merge_block_skeletons( + [mapping[label] for mapping in maps if label in mapping] + ) + for label in labels + } + + +def minimum_spanning_forest( + fragment: SkeletonFragment, + *, + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), +) -> SkeletonFragment: + """Remove cycles with a deterministic minimum-physical-length forest. + + This explicit postprocessor preserves all vertices and graph connected + components. Exact consolidation never removes cycles on its own. + """ + return _core._minimum_spanning_forest( + _normalize_fragment(fragment), + _normalize_sampling(spacing, 3, "minimum_spanning_forest", name="spacing"), + ) + + +def lattice_to_physical( + fragment: SkeletonFragment, + *, + spacing: float | Sequence[float] | None = (1.0, 1.0, 1.0), +) -> SkeletonFragment: + """Finalize a global lattice fragment as a physical-coordinate skeleton. + + The returned tuple matches ordinary TEASAR dtypes: float64 physical + vertices, uint64 edges, and float32 radii. + """ + return _core._lattice_to_physical( + _normalize_fragment(fragment), + _normalize_sampling(spacing, 3, "lattice_to_physical", name="spacing"), + ) + + +__all__ = [ + "block_border_targets", + "block_border_targets_labels", + "block_teasar", + "block_teasar_labels", + "lattice_to_physical", + "merge_block_skeleton_maps", + "merge_block_skeletons", + "minimum_spanning_forest", +] diff --git a/tests/skeleton/test_distributed.py b/tests/skeleton/test_distributed.py new file mode 100644 index 0000000..4c42668 --- /dev/null +++ b/tests/skeleton/test_distributed.py @@ -0,0 +1,339 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + +from development.skeleton.blockwise_stitching import ( + run_blockwise_binary, + run_blockwise_labels, +) + + +dist = bic.skeleton.distributed + + +def _contains(vertices, coordinate): + return np.any(np.all(vertices == np.asarray(coordinate), axis=1)) + + +def _number_of_components(number_of_vertices, edges): + parents = list(range(number_of_vertices)) + + def find(node): + while parents[node] != node: + parents[node] = parents[parents[node]] + node = parents[node] + return node + + for first, second in edges: + first_root, second_root = find(int(first)), find(int(second)) + if first_root != second_root: + parents[second_root] = first_root + return len({find(node) for node in range(number_of_vertices)}) + + +def test_neighboring_faces_select_the_same_global_target(): + volume = np.zeros((7, 9, 11), dtype=np.uint8) + volume[1:6, 2:8, 1:10] = 1 + left = np.ascontiguousarray(volume[:, :, :6]) + right = np.ascontiguousarray(volume[:, :, 5:]) + + left_target = dist.block_border_targets( + left, [(2, "high")], origin=(10, 20, 30), spacing=(2.0, 1.5, 1.0) + ) + right_target = dist.block_border_targets( + right, [(2, "low")], origin=(10, 20, 35), spacing=(2.0, 1.5, 1.0) + ) + np.testing.assert_array_equal(left_target, right_target) + np.testing.assert_array_equal(left_target, np.array([[13, 24, 35]])) + + +def test_border_targets_binary_and_label_semantics_and_corner_deduplication(): + labels = np.zeros((4, 4, 4), dtype=np.int64) + labels[1:3, 1, -1] = -7 + labels[1:3, 2, -1] = 2**40 + labels[1, 1, -2:] = -7 # same voxel is visible on two requested faces + + labeled = dist.block_border_targets_labels( + labels, [(1, "low"), (2, "high"), (2, "high")], background=0 + ) + assert set(labeled) == {-7, 2**40} + assert all(values.dtype == np.int64 for values in labeled.values()) + + binary = dist.block_border_targets(labels, [(2, "high")]) + assert binary.shape[0] == 1 # touching nonzero labels are one binary patch + assert sum(len(values) for values in labeled.values()) == 2 + + +def test_block_teasar_required_target_and_existing_teasar_equivalence(): + mask = np.zeros((9, 9, 13), dtype=np.uint8) + mask[2:7, 2:7, 1:12] = 1 + target = np.array([[106, 206, 310]], dtype=np.int64) + + block = dist.block_teasar( + mask, open_faces=(), origin=(100, 200, 300), required_targets=target, + spacing=(2.0, 1.0, 0.5) + ) + assert _contains(block[0], target[0]) + assert block[0].dtype == np.int64 + + unconstrained = dist.block_teasar( + mask, open_faces=(), origin=(0, 0, 0), spacing=(2, 1, 0.5) + ) + physical = dist.lattice_to_physical(unconstrained, spacing=(2, 1, 0.5)) + ordinary = bic.skeleton.teasar(mask, spacing=(2, 1, 0.5)) + for actual, expected in zip(physical, ordinary): + np.testing.assert_array_equal(actual, expected) + + +def test_required_target_order_and_thread_determinism(): + mask = np.zeros((11, 11, 15), dtype=np.uint8) + mask[2:9, 2:9, 1:14] = 1 + targets = np.array([[2, 2, 2], [8, 8, 13], [2, 8, 7]], dtype=np.int64) + first = dist.block_teasar( + mask, open_faces=(), required_targets=targets, number_of_threads=1 + ) + second = dist.block_teasar( + mask, open_faces=(), required_targets=targets[::-1].copy(), + number_of_threads=4 + ) + for actual, expected in zip(first, second): + np.testing.assert_array_equal(actual, expected) + assert all(_contains(first[0], target) for target in targets) + + +def test_open_face_preserves_interface_radius_and_never_emits_ghosts(): + mask = np.zeros((9, 9, 13), dtype=np.uint8) + mask[2:7, 2:7, 1:] = 1 + target = dist.block_border_targets(mask, [(2, "high")]) + assert target.shape == (1, 3) + + closed = dist.block_teasar( + mask, open_faces=(), required_targets=target + ) + opened = dist.block_teasar( + mask, open_faces=[(2, "high")], required_targets=target + ) + + def target_radius(graph): + index = np.flatnonzero(np.all(graph[0] == target[0], axis=1)) + assert index.shape == (1,) + return float(graph[2][index[0]]) + + assert target_radius(closed) == pytest.approx(1.0) + assert target_radius(opened) > 2.0 + np.testing.assert_array_equal(opened[0][0], target[0]) # interface root + assert np.all(opened[0] >= 0) + assert np.all(opened[0] < np.asarray(mask.shape)) + assert np.all(mask[tuple(opened[0].T)] != 0) + + +def test_lexicographically_greatest_open_target_is_deterministic_root(): + mask = np.zeros((7, 7, 11), dtype=np.uint8) + mask[2:5, 2:5, :] = 1 + targets = np.array([[3, 3, 0], [3, 3, 10]], dtype=np.int64) + first = dist.block_teasar( + mask, + open_faces=[(2, "low"), (2, "high")], + required_targets=targets, + number_of_threads=1, + ) + second = dist.block_teasar( + mask, + open_faces=[(2, "high"), (2, "low"), (2, "high")], + required_targets=targets[::-1], + number_of_threads=4, + ) + np.testing.assert_array_equal(first[0][0], targets[1]) + for actual, expected in zip(first, second): + np.testing.assert_array_equal(actual, expected) + + +def test_all_foreground_block_uses_finite_closed_boundary_fallback(): + mask = np.ones((5, 6, 7), dtype=np.uint8) + target = np.array([[4, 5, 6]], dtype=np.int64) + graph = dist.block_teasar( + mask, + open_faces=[ + (0, "low"), (0, "high"), + (1, "low"), (1, "high"), + (2, "low"), (2, "high"), + ], + required_targets=target, + ) + assert np.all(np.isfinite(graph[2])) + assert np.all(graph[2] >= 1.0) + assert np.all(graph[0] >= 0) + assert np.all(graph[0] < np.asarray(mask.shape)) + + +def test_labeled_required_targets_preserve_labels_and_reject_mismatch(): + labels = np.zeros((7, 7, 9), dtype=np.uint64) + labels[1:4, 1:4, 1:8] = 5 + labels[4:6, 3:6, 1:8] = 2**63 + 9 + targets = { + 5: np.array([[2, 2, 7]], dtype=np.int64), + 2**63 + 9: np.array([[5, 4, 7]], dtype=np.int64), + } + result = dist.block_teasar_labels( + labels, open_faces=(), required_targets=targets + ) + assert list(result) == [5, 2**63 + 9] + for label, target in targets.items(): + assert _contains(result[label][0], target[0]) + + with pytest.raises(ValueError, match="does not match"): + dist.block_teasar_labels( + labels, open_faces=(), + required_targets={5: np.array([[5, 4, 7]], dtype=np.int64)} + ) + + +def test_labeled_open_face_uses_label_specific_distance_boundary(): + labels = np.zeros((9, 9, 13), dtype=np.int64) + labels[2:7, 2:7, 1:] = -5 + targets = dist.block_border_targets_labels( + labels, [(2, "high")], background=0 + ) + result = dist.block_teasar_labels( + labels, + open_faces=[(2, "high")], + required_targets=targets, + background=0, + ) + graph = result[-5] + target = targets[-5][0] + np.testing.assert_array_equal(graph[0][0], target) + index = np.flatnonzero(np.all(graph[0] == target, axis=1)) + assert float(graph[2][index[0]]) > 2.0 + assert np.all(labels[tuple(graph[0].T)] == -5) + + +@pytest.mark.parametrize( + "targets, message", + [ + (np.array([[-1, 0, 0]], dtype=np.int64), "below the block origin"), + (np.array([[99, 0, 0]], dtype=np.int64), "outside the block"), + (np.array([[0, 0, 0]], dtype=np.int64), "foreground"), + ], +) +def test_required_target_validation(targets, message): + mask = np.zeros((3, 3, 3), dtype=np.uint8) + mask[1, 1, 1] = 1 + with pytest.raises(ValueError, match=message): + dist.block_teasar(mask, open_faces=(), required_targets=targets) + + +def test_open_faces_is_explicit_and_validated(): + mask = np.ones((3, 3, 3), dtype=np.uint8) + with pytest.raises(TypeError, match="open_faces"): + dist.block_teasar(mask) + with pytest.raises(ValueError, match="side"): + dist.block_teasar(mask, open_faces=[(0, "outside")]) + + +def test_exact_merge_reduces_vertices_radii_and_edges(): + first = ( + np.array([[0, 0, 0], [0, 0, 1], [9, 9, 9]], dtype=np.int64), + np.array([[0, 1]], dtype=np.uint64), + np.array([1.0, 2.0, 0.5], dtype=np.float32), + ) + second = ( + np.array([[0, 0, 1], [0, 0, 2]], dtype=np.int64), + np.array([[1, 0], [0, 0]], dtype=np.uint64), + np.array([3.0, 4.0], dtype=np.float32), + ) + merged = dist.merge_block_skeletons([first, second]) + np.testing.assert_array_equal( + merged[0], np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [9, 9, 9]]) + ) + np.testing.assert_array_equal(merged[1], np.array([[0, 1], [1, 2]])) + np.testing.assert_array_equal(merged[2], np.array([1, 3, 4, 0.5], np.float32)) + + +def test_merge_is_associative_commutative_and_idempotent(): + parts = [] + for offset in range(3): + parts.append( + ( + np.array([[0, 0, offset], [0, 0, offset + 1]], np.int64), + np.array([[0, 1]], np.uint64), + np.array([offset + 1, offset + 2], np.float32), + ) + ) + direct = dist.merge_block_skeletons(parts) + grouped = dist.merge_block_skeletons( + [dist.merge_block_skeletons(parts[:2]), parts[2]] + ) + reversed_ = dist.merge_block_skeletons(parts[::-1]) + duplicate = dist.merge_block_skeletons([direct, direct]) + for candidate in (grouped, reversed_, duplicate): + for actual, expected in zip(candidate, direct): + np.testing.assert_array_equal(actual, expected) + + +def test_labeled_merge_keeps_equal_coordinates_separate(): + fragment = ( + np.array([[1, 2, 3]], np.int64), + np.empty((0, 2), np.uint64), + np.array([1], np.float32), + ) + merged = dist.merge_block_skeleton_maps( + [{-4: fragment}, {2**63 + 2: fragment, -4: fragment}] + ) + assert list(merged) == [-4, 2**63 + 2] + assert merged[-4][0].shape == (1, 3) + assert merged[2**63 + 2][0].shape == (1, 3) + + +def test_minimum_spanning_forest_is_deterministic_and_preserves_vertices(): + graph = ( + np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [9, 9, 9]], np.int64), + np.array([[2, 3], [1, 3], [0, 2], [0, 1]], np.uint64), + np.ones(5, np.float32), + ) + forest = dist.minimum_spanning_forest(graph) + np.testing.assert_array_equal(forest[0], graph[0]) + np.testing.assert_array_equal(forest[2], graph[2]) + np.testing.assert_array_equal( + forest[1], np.array([[0, 1], [0, 2], [1, 3]], np.uint64) + ) + again = dist.minimum_spanning_forest(forest) + for actual, expected in zip(again, forest): + np.testing.assert_array_equal(actual, expected) + assert _number_of_components(len(forest[0]), forest[1]) == 2 + + +def test_serial_binary_harness_stitches_multiple_blocks(): + mask = np.zeros((11, 11, 19), dtype=np.uint8) + mask[5, 5, 1:18] = 1 + graph = run_blockwise_binary(mask, (6, 6, 6), remove_cycles=True) + assert graph[0].shape[0] > 0 + assert _number_of_components(len(graph[0]), graph[1]) == 1 + assert len(graph[1]) == len(graph[0]) - 1 + + +def test_serial_labeled_harness_keeps_touching_labels_separate(): + labels = np.zeros((9, 9, 17), dtype=np.int64) + labels[3, 4, 1:16] = -3 + labels[4, 4, 1:16] = 8 + graphs = run_blockwise_labels(labels, (5, 5, 6), remove_cycles=True) + assert list(graphs) == [-3, 8] + for graph in graphs.values(): + assert _number_of_components(len(graph[0]), graph[1]) == 1 + assert len(graph[1]) == len(graph[0]) - 1 + + +def test_empty_merge_and_invalid_fragment(): + empty = dist.merge_block_skeletons([]) + assert empty[0].shape == (0, 3) + assert empty[1].shape == (0, 2) + assert empty[2].shape == (0,) + + invalid = ( + np.array([[0, 0, 0]], np.int64), + np.array([[0, 1]], np.uint64), + np.array([1], np.float32), + ) + with pytest.raises(ValueError, match="outside its vertex range"): + dist.merge_block_skeletons([invalid]) From a8cb50f76c1a35512707baebb0a40b53ea246261 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 16 Jul 2026 18:21:42 +0200 Subject: [PATCH 2/3] Fix skeleton test --- tests/skeleton/test_distributed.py | 66 ++++++++++++++++++++++++++---- 1 file changed, 57 insertions(+), 9 deletions(-) diff --git a/tests/skeleton/test_distributed.py b/tests/skeleton/test_distributed.py index 4c42668..ec842c0 100644 --- a/tests/skeleton/test_distributed.py +++ b/tests/skeleton/test_distributed.py @@ -3,11 +3,6 @@ import bioimage_cpp as bic -from development.skeleton.blockwise_stitching import ( - run_blockwise_binary, - run_blockwise_labels, -) - dist = bic.skeleton.distributed @@ -304,20 +299,73 @@ def test_minimum_spanning_forest_is_deterministic_and_preserves_vertices(): assert _number_of_components(len(forest[0]), forest[1]) == 2 -def test_serial_binary_harness_stitches_multiple_blocks(): +def test_two_block_binary_pipeline_stitches_shared_target(): mask = np.zeros((11, 11, 19), dtype=np.uint8) mask[5, 5, 1:18] = 1 - graph = run_blockwise_binary(mask, (6, 6, 6), remove_cycles=True) + left = np.ascontiguousarray(mask[:, :, :11]) + right = np.ascontiguousarray(mask[:, :, 10:]) + left_targets = dist.block_border_targets( + left, [(2, "high")], origin=(0, 0, 0) + ) + right_targets = dist.block_border_targets( + right, [(2, "low")], origin=(0, 0, 10) + ) + np.testing.assert_array_equal(left_targets, right_targets) + fragments = [ + dist.block_teasar( + left, + open_faces=[(2, "high")], + origin=(0, 0, 0), + required_targets=left_targets, + ), + dist.block_teasar( + right, + open_faces=[(2, "low")], + origin=(0, 0, 10), + required_targets=right_targets, + ), + ] + graph = dist.minimum_spanning_forest( + dist.merge_block_skeletons(fragments) + ) assert graph[0].shape[0] > 0 assert _number_of_components(len(graph[0]), graph[1]) == 1 assert len(graph[1]) == len(graph[0]) - 1 -def test_serial_labeled_harness_keeps_touching_labels_separate(): +def test_two_block_labeled_pipeline_keeps_touching_labels_separate(): labels = np.zeros((9, 9, 17), dtype=np.int64) labels[3, 4, 1:16] = -3 labels[4, 4, 1:16] = 8 - graphs = run_blockwise_labels(labels, (5, 5, 6), remove_cycles=True) + left = np.ascontiguousarray(labels[:, :, :10]) + right = np.ascontiguousarray(labels[:, :, 9:]) + left_targets = dist.block_border_targets_labels( + left, [(2, "high")], origin=(0, 0, 0) + ) + right_targets = dist.block_border_targets_labels( + right, [(2, "low")], origin=(0, 0, 9) + ) + assert left_targets.keys() == right_targets.keys() + for label in left_targets: + np.testing.assert_array_equal(left_targets[label], right_targets[label]) + fragments = [ + dist.block_teasar_labels( + left, + open_faces=[(2, "high")], + origin=(0, 0, 0), + required_targets=left_targets, + ), + dist.block_teasar_labels( + right, + open_faces=[(2, "low")], + origin=(0, 0, 9), + required_targets=right_targets, + ), + ] + graphs = { + label: dist.minimum_spanning_forest(graph) + for label, graph in dist.merge_block_skeleton_maps(fragments).items() + } assert list(graphs) == [-3, 8] for graph in graphs.values(): assert _number_of_components(len(graph[0]), graph[1]) == 1 From 24ec0fdb14d3eda54b8d57d7b69aeb3cb1b8bea2 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 16 Jul 2026 19:50:46 +0200 Subject: [PATCH 3/3] Address claude code review --- examples/skeleton/teasar.py | 2 +- .../skeleton/detail/compact_grid_dijkstra.hxx | 22 +++++++++++++ .../skeleton/distributed/border_targets.hxx | 4 +-- include/bioimage_cpp/skeleton/teasar.hxx | 32 +++---------------- tests/skeleton/test_distributed.py | 19 +++++++++++ 5 files changed, 49 insertions(+), 30 deletions(-) diff --git a/examples/skeleton/teasar.py b/examples/skeleton/teasar.py index dd7f55d..32812be 100644 --- a/examples/skeleton/teasar.py +++ b/examples/skeleton/teasar.py @@ -11,7 +11,7 @@ def main(): path = Path(__file__).with_name("00004_gt_mask.mrc") - full_crop = False + full_crop = True if full_crop: bb = np.s_[:] else: diff --git a/include/bioimage_cpp/skeleton/detail/compact_grid_dijkstra.hxx b/include/bioimage_cpp/skeleton/detail/compact_grid_dijkstra.hxx index 232f06a..c156f10 100644 --- a/include/bioimage_cpp/skeleton/detail/compact_grid_dijkstra.hxx +++ b/include/bioimage_cpp/skeleton/detail/compact_grid_dijkstra.hxx @@ -52,6 +52,28 @@ struct CompactGridDomain { [[nodiscard]] bool has_full_lookup() const noexcept { return !full_to_compact.empty(); } + + // CSR domains release the dense reverse map after building adjacency. + // Keep full-index conversion independent of that storage decision. + [[nodiscard]] std::uint32_t compact_node_from_full( + const std::size_t full + ) const noexcept { + if (has_full_lookup()) { + return full < full_to_compact.size() + ? full_to_compact[full] : kNoCompactNode; + } + if (full > static_cast(kNoCompactNode)) { + return kNoCompactNode; + } + const auto full32 = static_cast(full); + const auto it = std::lower_bound( + compact_to_full.begin(), compact_to_full.end(), full32 + ); + if (it == compact_to_full.end() || *it != full32) { + return kNoCompactNode; + } + return static_cast(it - compact_to_full.begin()); + } }; inline void build_compact_neighbors( diff --git a/include/bioimage_cpp/skeleton/distributed/border_targets.hxx b/include/bioimage_cpp/skeleton/distributed/border_targets.hxx index 73b0b39..b70b2ad 100644 --- a/include/bioimage_cpp/skeleton/distributed/border_targets.hxx +++ b/include/bioimage_cpp/skeleton/distributed/border_targets.hxx @@ -277,12 +277,12 @@ std::vector> border_targets_impl( }(); const auto edge_distance = std::min({ static_cast(spacing[face_axes[0]]) * - (static_cast(run.y) - 0.5L), + (static_cast(run.y) + 0.5L), static_cast(spacing[face_axes[0]]) * (static_cast(height) - 0.5L - static_cast(run.y)), static_cast(spacing[face_axes[1]]) * - (static_cast(x) - 0.5L), + (static_cast(x) + 0.5L), static_cast(spacing[face_axes[1]]) * (static_cast(width) - 0.5L - static_cast(x)), diff --git a/include/bioimage_cpp/skeleton/teasar.hxx b/include/bioimage_cpp/skeleton/teasar.hxx index 500b694..cfc952a 100644 --- a/include/bioimage_cpp/skeleton/teasar.hxx +++ b/include/bioimage_cpp/skeleton/teasar.hxx @@ -625,26 +625,12 @@ inline LatticeSkeletonGraph teasar_compact_impl( { BIOIMAGE_PROFILE_SCOPE(profile, "root_dijkstra") if (required_root != std::numeric_limits::max()) { - if (required_root > std::numeric_limits::max()) { - throw std::runtime_error( - "prepared required root exceeds compact index range" - ); - } - const auto it = std::lower_bound( - domain.compact_to_full.begin(), domain.compact_to_full.end(), - static_cast(required_root) - ); - if ( - it == domain.compact_to_full.end() || - *it != static_cast(required_root) - ) { + root = domain.compact_node_from_full(required_root); + if (root == detail::kNoCompactNode) { throw std::runtime_error( "prepared required root is not foreground" ); } - root = static_cast( - std::distance(domain.compact_to_full.begin(), it) - ); } else { std::vector first_field; detail::compact_physical_distance_field( @@ -808,19 +794,11 @@ inline LatticeSkeletonGraph teasar_compact_impl( if (full_target >= n) { throw std::runtime_error("prepared required target is out of bounds"); } - const auto it = std::lower_bound( - domain.compact_to_full.begin(), domain.compact_to_full.end(), - static_cast(full_target) - ); - if ( - it == domain.compact_to_full.end() || - *it != static_cast(full_target) - ) { + const auto target = domain.compact_node_from_full(full_target); + if (target == detail::kNoCompactNode) { throw std::runtime_error("prepared required target is not foreground"); } - compact_required_targets.push_back(static_cast( - std::distance(domain.compact_to_full.begin(), it) - )); + compact_required_targets.push_back(target); } std::sort( compact_required_targets.begin(), compact_required_targets.end(), diff --git a/tests/skeleton/test_distributed.py b/tests/skeleton/test_distributed.py index ec842c0..c9baa9c 100644 --- a/tests/skeleton/test_distributed.py +++ b/tests/skeleton/test_distributed.py @@ -60,6 +60,25 @@ def test_border_targets_binary_and_label_semantics_and_corner_deduplication(): assert sum(len(values) for values in labeled.values()) == 2 +def test_border_target_edge_tiebreak_uses_physical_low_edge_distance(): + face = np.array( + [ + [0, 0, 1, 0, 0, 0, 1, 0, 1, 0], + [0, 1, 0, 0, 0, 1, 1, 1, 1, 0], + [1, 0, 0, 1, 1, 1, 0, 0, 1, 0], + ], + dtype=np.uint8, + ) + block = np.zeros((3, 10, 2), dtype=np.uint8) + block[:, :, -1] = face + + targets = dist.block_border_targets( + block, [(2, "high")], spacing=(2.0, 1.0, 1.0) + ) + + np.testing.assert_array_equal(targets, [[1, 1, 1], [2, 4, 1]]) + + def test_block_teasar_required_target_and_existing_teasar_equivalence(): mask = np.zeros((9, 9, 13), dtype=np.uint8) mask[2:7, 2:7, 1:12] = 1