From 25fa6e8e30c52e48276bd295e30aa00d2358bad7 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Sat, 16 May 2026 21:48:42 +0200 Subject: [PATCH 1/3] Optimize lifted affinity features --- development/graph/_rag_compatibility.py | 35 ++- .../graph/feature_accumulation.hxx | 200 +++++++++++++++--- tests/graph/test_rag_features.py | 83 ++++++++ 3 files changed, 278 insertions(+), 40 deletions(-) diff --git a/development/graph/_rag_compatibility.py b/development/graph/_rag_compatibility.py index aedc231..d8c6561 100644 --- a/development/graph/_rag_compatibility.py +++ b/development/graph/_rag_compatibility.py @@ -85,6 +85,10 @@ def make_watershed_labels( def time_call(function, repeats: int): + # One untimed warm-up call before the measured loop so the first sample + # doesn't carry nanobind tuple-shape caching, numpy ufunc init, code-page + # faults, etc. Mirrors the grid-affinity helper for consistency. + function() timings = [] result = None for _ in range(repeats): @@ -136,7 +140,9 @@ def compare_boundary_features( import bioimage_cpp as bic import nifty.graph.rag as nrag - block_shape = list(labels.shape) + # Don't force blockShape on the nifty side — its default block layout is + # what its parallelism is tuned for, and forcing a single block (== labels + # shape) starves nifty's worker pool. bic_timings, bic_features = time_call( lambda: bic.graph.edge_map_features( bic_rag, labels, boundary_map, number_of_threads=threads @@ -147,7 +153,6 @@ def compare_boundary_features( lambda: nrag.accumulateEdgeMeanAndLength( nifty_rag, np.ascontiguousarray(boundary_map, dtype=np.float32), - blockShape=block_shape, numberOfThreads=threads, ), repeats, @@ -193,9 +198,12 @@ def compare_affinity_features( ), repeats, ) - # The installed Nifty wrapper can crash or fail when an explicit - # numberOfThreads is passed here. The default path is still the reference - # implementation and keeps this compatibility script robust. + # nifty's accumulateAffinityStandartFeatures crashes inside vigra's + # UserRangeHistogram when numberOfThreads=1 (the accumulators never get + # setMinMax). For threads >= 2 we can pass the value through; for the + # single-thread case we fall back to nifty's default (-1, all cores) and + # surface the unfairness in the printed report. + nifty_threads_effective = threads if threads != 1 else -1 nifty_timings, nifty_features_full = time_call( lambda: nrag.accumulateAffinityStandartFeatures( nifty_rag, @@ -203,6 +211,7 @@ def compare_affinity_features( offsets_for_nifty, min_val, max_val, + numberOfThreads=nifty_threads_effective, ), repeats, ) @@ -213,6 +222,8 @@ def compare_affinity_features( np.testing.assert_allclose(aligned_bic, nifty_features, rtol=1.0e-5, atol=1.0e-6) return bic_timings, nifty_timings, { "max_abs_diff": float(np.max(np.abs(aligned_bic - nifty_features))), + "nifty_threads_effective": nifty_threads_effective, + "nifty_threads_requested": threads, } @@ -290,7 +301,15 @@ def run_compatibility_check( print(f"boundary-map size convention: {boundary_summary['size_convention']}") _print_timing("affinity features", affinity_bic_timings, affinity_nifty_timings) print(f"affinity feature max abs diff: {affinity_summary['max_abs_diff']:.6g}") - print("nifty affinity timing uses the wrapper default thread handling") + requested = affinity_summary["nifty_threads_requested"] + effective = affinity_summary["nifty_threads_effective"] + if requested != effective: + print( + f"WARNING: affinity features — requested {requested} thread(s) but " + f"nifty was called with numberOfThreads={effective} " + "(its single-thread path crashes inside vigra's UserRangeHistogram). " + f"bioimage-cpp used {requested} thread(s); the timing comparison is NOT apples-to-apples." + ) def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float]): @@ -303,7 +322,9 @@ def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float def add_common_arguments(parser: argparse.ArgumentParser) -> None: - parser.add_argument("--repeats", type=int, default=3) + # 5 matches the grid-affinity helper; median of 3 is noisy because one + # GC stall in the middle sample becomes the median. + parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--threads", type=int, default=1) parser.add_argument("--watershed-min-distance", type=int, default=5) parser.add_argument("--watershed-grid-spacing", type=int, default=12) diff --git a/include/bioimage_cpp/graph/feature_accumulation.hxx b/include/bioimage_cpp/graph/feature_accumulation.hxx index edf4297..e0403c6 100644 --- a/include/bioimage_cpp/graph/feature_accumulation.hxx +++ b/include/bioimage_cpp/graph/feature_accumulation.hxx @@ -2,6 +2,7 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/profile.hxx" #include "bioimage_cpp/detail/threading.hxx" #include "bioimage_cpp/graph/region_adjacency_graph.hxx" @@ -204,29 +205,123 @@ void scan_edge_map_3d_chunk( } } +// Given an offset along one axis and the axis length, return the half-open +// range of axis coordinates `[lo, hi)` for which `coord + delta` stays in +// `[0, length)`. Returns `lo >= hi` if the offset is larger than the axis. +inline void valid_axis_range( + const std::ptrdiff_t delta, + const std::size_t length, + std::size_t &lo, + std::size_t &hi +) { + if (delta >= 0) { + lo = 0; + const auto d = static_cast(delta); + hi = (d >= length) ? 0 : (length - d); + } else { + const auto d = static_cast(-delta); + lo = (d >= length) ? length : d; + hi = length; + } +} + template -void scan_affinity_chunk( +void scan_affinity_2d_chunk( const RegionAdjacencyGraph &rag, const LabelT *labels, const ValueT *affinities, const std::vector> &offsets, - const std::vector &shape, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::vector &stats ) { - const auto spatial_strides = bioimage_cpp::detail::c_order_strides(shape); - const auto number_of_nodes = static_cast(number_of_pixels(shape)); + const auto number_of_nodes = static_cast(height * width); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { - const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offsets[channel], shape, spatial_strides, target)) { - continue; + const auto &off = offsets[channel]; + std::size_t y_lo_full, y_hi_full, x_lo, x_hi; + valid_axis_range(off[0], height, y_lo_full, y_hi_full); + valid_axis_range(off[1], width, x_lo, x_hi); + const auto y_lo = std::max(y_lo_full, y_begin); + const auto y_hi = std::min(y_hi_full, y_end); + if (y_lo >= y_hi || x_lo >= x_hi) { + continue; + } + const auto offset_stride = + off[0] * static_cast(width) + off[1]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] + ); + } } - const auto edge = edge_for_labels(rag, label_at(labels, node), label_at(labels, target)); - if (edge >= 0) { - stats[static_cast(edge)].add(affinities[channel_offset + node]); + } + } +} + +template +void scan_affinity_3d_chunk( + const RegionAdjacencyGraph &rag, + const LabelT *labels, + const ValueT *affinities, + const std::vector> &offsets, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::vector &stats +) { + const auto slice_size = height * width; + const auto number_of_nodes = static_cast(depth * slice_size); + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + const auto &off = offsets[channel]; + std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; + valid_axis_range(off[0], depth, z_lo_full, z_hi_full); + valid_axis_range(off[1], height, y_lo, y_hi); + valid_axis_range(off[2], width, x_lo, x_hi); + const auto z_lo = std::max(z_lo_full, z_begin); + const auto z_hi = std::min(z_hi_full, z_end); + if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { + continue; + } + const auto offset_stride = + off[0] * static_cast(slice_size) + + off[1] * static_cast(width) + + off[2]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + for (std::size_t z = z_lo; z < z_hi; ++z) { + const auto slice_offset = z * slice_size; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = slice_offset + y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] + ); + } + } } } } @@ -408,40 +503,79 @@ void accumulate_affinity_features( throw std::invalid_argument("out shape must be (number_of_edges, number_of_features)"); } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); - const auto n_threads = detail::normalize_thread_count(number_of_threads, number_of_nodes); + const auto work_items = static_cast(labels.shape[0]); + const auto n_threads = detail::normalize_thread_count(number_of_threads, work_items); const auto number_of_edges = static_cast(rag.number_of_edges()); + BIOIMAGE_PROFILE_INIT(aff_profiler); + const auto run_scan = [&](auto &per_thread_stats) { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:scan"); bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_features::scan_affinity_chunk( - rag, labels.data, affinities.data, offsets, labels.shape, - begin, end, per_thread_stats[thread_id] - ); + if (labels.ndim() == 2) { + detail_features::scan_affinity_2d_chunk( + rag, labels.data, affinities.data, offsets, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread_stats[thread_id] + ); + } else { + detail_features::scan_affinity_3d_chunk( + rag, labels.data, affinities.data, offsets, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread_stats[thread_id] + ); + } } ); }; if (compute_complex_features) { - std::vector> per_thread_stats( - n_threads, - std::vector(number_of_edges) - ); + std::vector> per_thread_stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc"); + per_thread_stats.assign( + n_threads, + std::vector(number_of_edges) + ); + } run_scan(per_thread_stats); - auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges); - detail_features::write_complex_features(stats, out); + std::vector stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge"); + stats = detail_features::merge_stats(per_thread_stats, number_of_edges); + } + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write"); + detail_features::write_complex_features(stats, out); + } } else { - std::vector> per_thread_stats( - n_threads, - std::vector(number_of_edges) - ); + std::vector> per_thread_stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc"); + per_thread_stats.assign( + n_threads, + std::vector(number_of_edges) + ); + } run_scan(per_thread_stats); - auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges); - detail_features::write_simple_features(stats, out); + std::vector stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge"); + stats = detail_features::merge_stats(per_thread_stats, number_of_edges); + } + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write"); + detail_features::write_simple_features(stats, out); + } } + + BIOIMAGE_PROFILE_REPORT(aff_profiler); } } // namespace bioimage_cpp::graph diff --git a/tests/graph/test_rag_features.py b/tests/graph/test_rag_features.py index 263d32d..8f8c42c 100644 --- a/tests/graph/test_rag_features.py +++ b/tests/graph/test_rag_features.py @@ -89,6 +89,89 @@ def test_affinity_features_simple(): ) +def _reference_affinity_features(labels, rag, affinities, offsets): + """Trivial Python reference: per-channel sweep with explicit bounds checks.""" + uv_ids = np.asarray(rag.uv_ids(), dtype=np.uint64) + edge_lookup = {(int(u), int(v)): i for i, (u, v) in enumerate(uv_ids)} + out = np.zeros((len(edge_lookup), 2), dtype=np.float64) + shape = labels.shape + for channel, offset in enumerate(offsets): + for index in np.ndindex(*shape): + target = tuple(int(c) + int(d) for c, d in zip(index, offset)) + if any(t < 0 or t >= s for t, s in zip(target, shape)): + continue + u, v = int(labels[index]), int(labels[target]) + if u == v: + continue + key = (min(u, v), max(u, v)) + edge = edge_lookup.get(key) + if edge is None: + # Pair (u,v) doesn't exist as a RAG edge (e.g. when the offset + # reaches further than direct neighbors). The kernel skips + # these via find_edge -> -1; the reference must match. + continue + out[edge, 0] += float(affinities[(channel,) + index]) + out[edge, 1] += 1.0 + for edge in range(len(edge_lookup)): + if out[edge, 1] > 0: + out[edge, 0] /= out[edge, 1] + return out + + +@pytest.mark.parametrize( + "offsets", + [ + # axis-aligned positive offsets (the only case the previous tests covered) + [[0, 1], [1, 0]], + # negative offsets — must still hit the valid box correctly + [[0, -1], [-1, 0]], + # large magnitude that crops most of the image + [[0, 3], [3, 0]], + # diagonal / multi-axis nonzero — exercises the new valid-box math + [[1, 1], [1, -1], [-1, 1]], + # offset larger than an axis — must produce zero contributions + [[10, 0], [0, 10]], + ], +) +def test_affinity_features_2d_offsets_match_reference(offsets): + rng = np.random.default_rng(42) + labels = rng.integers(0, 5, size=(5, 7)).astype(np.uint32) + rag = bic.graph.region_adjacency_graph(labels) + affinities = rng.standard_normal( + (len(offsets),) + labels.shape + ).astype(np.float32) + + expected = _reference_affinity_features(labels, rag, affinities, offsets) + got = bic.graph.affinity_features(rag, labels, affinities, offsets=offsets) + + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + "offsets", + [ + [[1, 0, 0], [0, 1, 0], [0, 0, 1]], + [[-1, 0, 0], [0, -1, 0], [0, 0, -1]], + # 3D diagonal + [[1, 1, 1], [1, -1, 1], [-1, 1, -1]], + # mixed magnitudes + [[2, 0, 0], [0, 0, -3]], + ], +) +def test_affinity_features_3d_offsets_match_reference(offsets): + rng = np.random.default_rng(7) + labels = rng.integers(0, 4, size=(4, 5, 6)).astype(np.uint32) + rag = bic.graph.region_adjacency_graph(labels) + affinities = rng.standard_normal( + (len(offsets),) + labels.shape + ).astype(np.float32) + + expected = _reference_affinity_features(labels, rag, affinities, offsets) + got = bic.graph.affinity_features(rag, labels, affinities, offsets=offsets) + + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + def test_affinity_features_complex_parallel_matches_single_thread(): labels = np.array([[0, 1, 1], [2, 2, 3], [2, 4, 3]], dtype=np.uint64) rag = bic.graph.region_adjacency_graph(labels) From 78eb1d59b276b6b5999df4fb4ab3bca1a1caebdf Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Sat, 16 May 2026 21:59:28 +0200 Subject: [PATCH 2/3] More optimization --- .../graph/feature_accumulation.hxx | 139 ++++++++---- .../graph/lifted_from_affinities.hxx | 212 +++++++++++++----- 2 files changed, 246 insertions(+), 105 deletions(-) diff --git a/include/bioimage_cpp/graph/feature_accumulation.hxx b/include/bioimage_cpp/graph/feature_accumulation.hxx index e0403c6..4b2db1f 100644 --- a/include/bioimage_cpp/graph/feature_accumulation.hxx +++ b/include/bioimage_cpp/graph/feature_accumulation.hxx @@ -225,6 +225,83 @@ inline void valid_axis_range( } } +// Sweep every (node, target) pair on a 2D grid for which `node + offset` stays +// in bounds, restricted to the half-open y-slab [y_begin, y_end). The body +// receives flat C-order indices for both endpoints and is expected to inline +// at -O2 since this is a header-only template with a fully-known callable +// type at instantiation. +template +void sweep_offset_box_2d( + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, + const Body &body +) { + std::size_t y_lo_full, y_hi_full, x_lo, x_hi; + valid_axis_range(dy, height, y_lo_full, y_hi_full); + valid_axis_range(dx, width, x_lo, x_hi); + const auto y_lo = std::max(y_lo_full, y_begin); + const auto y_hi = std::min(y_hi_full, y_end); + if (y_lo >= y_hi || x_lo >= x_hi) { + return; + } + const auto offset_stride = dy * static_cast(width) + dx; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + body(static_cast(node), target); + } + } +} + +// 3D variant of `sweep_offset_box_2d`. Restricts the sweep to a z-slab. +template +void sweep_offset_box_3d( + const std::ptrdiff_t dz, + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + const Body &body +) { + std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; + valid_axis_range(dz, depth, z_lo_full, z_hi_full); + valid_axis_range(dy, height, y_lo, y_hi); + valid_axis_range(dx, width, x_lo, x_hi); + const auto z_lo = std::max(z_lo_full, z_begin); + const auto z_hi = std::min(z_hi_full, z_end); + if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { + return; + } + const auto slice_size = height * width; + const auto offset_stride = + dz * static_cast(slice_size) + + dy * static_cast(width) + dx; + for (std::size_t z = z_lo; z < z_hi; ++z) { + const auto slice_offset = z * slice_size; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = slice_offset + y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + body(static_cast(node), target); + } + } + } +} + template void scan_affinity_2d_chunk( const RegionAdjacencyGraph &rag, @@ -240,25 +317,11 @@ void scan_affinity_2d_chunk( const auto number_of_nodes = static_cast(height * width); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { const auto &off = offsets[channel]; - std::size_t y_lo_full, y_hi_full, x_lo, x_hi; - valid_axis_range(off[0], height, y_lo_full, y_hi_full); - valid_axis_range(off[1], width, x_lo, x_hi); - const auto y_lo = std::max(y_lo_full, y_begin); - const auto y_hi = std::min(y_hi_full, y_end); - if (y_lo >= y_hi || x_lo >= x_hi) { - continue; - } - const auto offset_stride = - off[0] * static_cast(width) + off[1]; const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::size_t y = y_lo; y < y_hi; ++y) { - const auto row_offset = y * width; - for (std::size_t x = x_lo; x < x_hi; ++x) { - const auto node = row_offset + x; - const auto target = static_cast( - static_cast(node) + offset_stride - ); + sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { const auto u = label_at(labels, node); const auto v = label_at(labels, target); const auto edge = edge_for_labels(rag, u, v); @@ -268,7 +331,7 @@ void scan_affinity_2d_chunk( ); } } - } + ); } } @@ -289,41 +352,21 @@ void scan_affinity_3d_chunk( const auto number_of_nodes = static_cast(depth * slice_size); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { const auto &off = offsets[channel]; - std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; - valid_axis_range(off[0], depth, z_lo_full, z_hi_full); - valid_axis_range(off[1], height, y_lo, y_hi); - valid_axis_range(off[2], width, x_lo, x_hi); - const auto z_lo = std::max(z_lo_full, z_begin); - const auto z_hi = std::min(z_hi_full, z_end); - if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { - continue; - } - const auto offset_stride = - off[0] * static_cast(slice_size) + - off[1] * static_cast(width) + - off[2]; const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::size_t z = z_lo; z < z_hi; ++z) { - const auto slice_offset = z * slice_size; - for (std::size_t y = y_lo; y < y_hi; ++y) { - const auto row_offset = slice_offset + y * width; - for (std::size_t x = x_lo; x < x_hi; ++x) { - const auto node = row_offset + x; - const auto target = static_cast( - static_cast(node) + offset_stride + sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] ); - const auto u = label_at(labels, node); - const auto v = label_at(labels, target); - const auto edge = edge_for_labels(rag, u, v); - if (edge >= 0) { - stats[static_cast(edge)].add( - affinities[channel_offset + node] - ); - } } } - } + ); } } diff --git a/include/bioimage_cpp/graph/lifted_from_affinities.hxx b/include/bioimage_cpp/graph/lifted_from_affinities.hxx index 4fed26c..8f968c4 100644 --- a/include/bioimage_cpp/graph/lifted_from_affinities.hxx +++ b/include/bioimage_cpp/graph/lifted_from_affinities.hxx @@ -65,34 +65,65 @@ inline void validate_affinity_inputs( } template -void discover_lifted_chunk( +void discover_lifted_2d_chunk( const RegionAdjacencyGraph &rag, const LabelT *labels, const std::vector> &offsets, const std::vector &long_range_channels, - const std::vector &shape, - const std::vector &strides, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::unordered_set &out ) { for (const auto channel : long_range_channels) { - const auto &offset = offsets[channel]; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { - continue; + const auto &off = offsets[channel]; + detail_features::sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + if (rag.find_edge(u, v) >= 0) { + return; + } + out.insert(bioimage_cpp::detail::edge_key(u, v)); } - const auto u = detail_features::label_at(labels, node); - const auto v = detail_features::label_at(labels, target); - if (u == v) { - continue; - } - if (rag.find_edge(u, v) >= 0) { - continue; + ); + } +} + +template +void discover_lifted_3d_chunk( + const RegionAdjacencyGraph &rag, + const LabelT *labels, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::unordered_set &out +) { + for (const auto channel : long_range_channels) { + const auto &off = offsets[channel]; + detail_features::sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + if (rag.find_edge(u, v) >= 0) { + return; + } + out.insert(bioimage_cpp::detail::edge_key(u, v)); } - out.insert(bioimage_cpp::detail::edge_key(u, v)); - } + ); } } @@ -133,11 +164,10 @@ std::vector lifted_edges_from_offsets( return {}; } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto work_items = static_cast(labels.shape[0]); const auto n_threads = bioimage_cpp::detail::normalize_thread_count( - number_of_threads, number_of_nodes + number_of_threads, work_items ); - const auto strides = bioimage_cpp::detail::c_order_strides(labels.shape); using EdgeSet = std::unordered_set< bioimage_cpp::detail::Edge, bioimage_cpp::detail::EdgeHash @@ -146,12 +176,24 @@ std::vector lifted_edges_from_offsets( bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_lifted::discover_lifted_chunk( - rag, labels.data, offsets, long_range_channels, - labels.shape, strides, begin, end, per_thread[thread_id] - ); + if (labels.ndim() == 2) { + detail_lifted::discover_lifted_2d_chunk( + rag, labels.data, offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread[thread_id] + ); + } else { + detail_lifted::discover_lifted_3d_chunk( + rag, labels.data, offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread[thread_id] + ); + } } ); @@ -173,39 +215,82 @@ std::vector lifted_edges_from_offsets( namespace detail_lifted { template -void scan_lifted_affinity_chunk( +void scan_lifted_affinity_2d_chunk( const std::unordered_map &lifted_index, const LabelT *labels, const ValueT *affinities, const std::vector> &offsets, const std::vector &long_range_channels, - const std::vector &shape, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::vector &stats ) { - const auto strides = bioimage_cpp::detail::c_order_strides(shape); - const auto number_of_nodes = static_cast(detail_features::number_of_pixels(shape)); + const auto number_of_nodes = static_cast(height * width); for (const auto channel : long_range_channels) { - const auto &offset = offsets[channel]; - const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { - continue; - } - const auto u = detail_features::label_at(labels, node); - const auto v = detail_features::label_at(labels, target); - if (u == v) { - continue; + const auto &off = offsets[channel]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + detail_features::sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + const auto found = lifted_index.find( + bioimage_cpp::detail::edge_key(u, v) + ); + if (found == lifted_index.end()) { + return; + } + stats[found->second].add(affinities[channel_offset + node]); } - const auto found = lifted_index.find(bioimage_cpp::detail::edge_key(u, v)); - if (found == lifted_index.end()) { - continue; + ); + } +} + +template +void scan_lifted_affinity_3d_chunk( + const std::unordered_map + &lifted_index, + const LabelT *labels, + const ValueT *affinities, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::vector &stats +) { + const auto slice_size = height * width; + const auto number_of_nodes = static_cast(depth * slice_size); + for (const auto channel : long_range_channels) { + const auto &off = offsets[channel]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + detail_features::sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + const auto found = lifted_index.find( + bioimage_cpp::detail::edge_key(u, v) + ); + if (found == lifted_index.end()) { + return; + } + stats[found->second].add(affinities[channel_offset + node]); } - stats[found->second].add(affinities[channel_offset + node]); - } + ); } } @@ -297,21 +382,34 @@ void accumulate_lifted_affinity_features( return; } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto work_items = static_cast(labels.shape[0]); const auto n_threads = bioimage_cpp::detail::normalize_thread_count( - number_of_threads, number_of_nodes + number_of_threads, work_items ); const auto run_scan = [&](auto &per_thread_stats) { bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_lifted::scan_lifted_affinity_chunk( - lifted_index, labels.data, affinities.data, - offsets, long_range_channels, labels.shape, - begin, end, per_thread_stats[thread_id] - ); + if (labels.ndim() == 2) { + detail_lifted::scan_lifted_affinity_2d_chunk( + lifted_index, labels.data, affinities.data, + offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread_stats[thread_id] + ); + } else { + detail_lifted::scan_lifted_affinity_3d_chunk( + lifted_index, labels.data, affinities.data, + offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread_stats[thread_id] + ); + } } ); }; From c6aef542f6e8807fbeb8206f8ef36dbaa89d9895 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Sun, 17 May 2026 12:55:32 +0200 Subject: [PATCH 3/3] Add affinity computation --- CMakeLists.txt | 1 + .../affinities/check_compute_affinities.py | 133 ++++++++++ .../affinities/compute_affinities.hxx | 191 ++++++++++++++ src/bindings/affinities.cxx | 232 ++++++++++++++++++ src/bindings/affinities.hxx | 9 + src/bindings/module.cxx | 2 + src/bioimage_cpp/__init__.py | 2 + src/bioimage_cpp/_data.py | 23 ++ src/bioimage_cpp/affinities/__init__.py | 5 + .../affinities/compute_affinities.py | 127 ++++++++++ tests/affinities/test_compute_affinities.py | 154 ++++++++++++ 11 files changed, 879 insertions(+) create mode 100644 development/affinities/check_compute_affinities.py create mode 100644 include/bioimage_cpp/affinities/compute_affinities.hxx create mode 100644 src/bindings/affinities.cxx create mode 100644 src/bindings/affinities.hxx create mode 100644 src/bioimage_cpp/affinities/__init__.py create mode 100644 src/bioimage_cpp/affinities/compute_affinities.py create mode 100644 tests/affinities/test_compute_affinities.py diff --git a/CMakeLists.txt b/CMakeLists.txt index e56cde5..23391cd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -11,6 +11,7 @@ find_package(nanobind CONFIG REQUIRED) nanobind_add_module(_core NB_STATIC + src/bindings/affinities.cxx src/bindings/blocking.cxx src/bindings/module.cxx src/bindings/graph.cxx diff --git a/development/affinities/check_compute_affinities.py b/development/affinities/check_compute_affinities.py new file mode 100644 index 0000000..1056067 --- /dev/null +++ b/development/affinities/check_compute_affinities.py @@ -0,0 +1,133 @@ +"""Cross-check bioimage-cpp's compute_affinities against affogato. + +Benchmarks on the registered ISBI ground-truth segmentation volume +(30 × 512 × 512 = 7.86 M voxels, ~660 distinct labels) using the same +17-channel offset configuration that elf uses for mutex-watershed on this +data. Small enough to fit in memory, big enough that initialization and +allocation effects don't dominate. + +Not part of the pytest suite (per CLAUDE.md: external-library comparisons +live under ``development/``, not under ``tests/``). +""" + +from __future__ import annotations + +import argparse +import sys +from statistics import mean, median +from time import perf_counter + +import numpy as np + +import bioimage_cpp as bic +from bioimage_cpp._data import ISBI_AFFINITY_OFFSETS, load_isbi_gt_segmentation + +try: + import affogato.affinities as affo +except ImportError as error: # pragma: no cover - dev script + sys.stderr.write(f"affogato not installed: {error}\n") + sys.exit(1) + + +# Subsets of ISBI_AFFINITY_OFFSETS. The "nearest neighbours" subset is the +# typical multicut input (3 channels); the "full 17" subset matches +# elf's mutex-watershed proposal generator and exercises long-range offsets. +OFFSET_SUBSETS = { + "nearest": [(-1, 0, 0), (0, -1, 0), (0, 0, -1)], + "full17": list(ISBI_AFFINITY_OFFSETS), +} + + +def time_call(fn, repeats): + timings = [] + result = None + for _ in range(repeats): + start = perf_counter() + result = fn() + timings.append(perf_counter() - start) + return timings, result + + +def run_case(labels, offsets, *, repeats, ignore_label=None): + offsets_list = [list(offset) for offset in offsets] + + bic_timings, (bic_affs, bic_mask) = time_call( + lambda: bic.affinities.compute_affinities( + labels, + offsets_list, + ignore_label=ignore_label, + return_mask=True, + number_of_threads=1, + ), + repeats, + ) + affo_timings, (affo_affs, affo_mask) = time_call( + lambda: affo.compute_affinities( + labels, + offsets_list, + have_ignore_label=ignore_label is not None, + ignore_label=ignore_label if ignore_label is not None else 0, + ), + repeats, + ) + + return { + "n_offsets": len(offsets_list), + "ignore": ignore_label, + "ok_affs": np.array_equal(bic_affs, affo_affs), + "ok_mask": np.array_equal(bic_mask, affo_mask), + "bic_median_s": median(bic_timings), + "affo_median_s": median(affo_timings), + "bic_mean_s": mean(bic_timings), + "affo_mean_s": mean(affo_timings), + "max_abs_diff": float(np.max(np.abs(bic_affs - affo_affs))) if bic_affs.shape == affo_affs.shape else float("nan"), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--timeout", type=float, default=60.0) + args = parser.parse_args() + + labels = load_isbi_gt_segmentation(timeout=args.timeout) + n_labels = int(labels.max()) + 1 + n_voxels = int(np.prod(labels.shape)) + print( + f"labels: shape={labels.shape}, dtype={labels.dtype}, " + f"n_voxels={n_voxels:,}, n_labels={n_labels}", + flush=True, + ) + + rows = [] + for name, offsets in OFFSET_SUBSETS.items(): + for ig in (None, 0): + row = run_case(labels, offsets, repeats=args.repeats, ignore_label=ig) + row["offsets_name"] = name + rows.append(row) + + print() + print( + f"{'offsets':>10} {'n':>3} {'ignore':>6} {'affs':>5} {'mask':>5}" + f" {'bic_s':>9} {'affo_s':>9} {'speedup':>8}" + ) + print("-" * 68) + all_ok = True + for r in rows: + speedup = r["affo_median_s"] / r["bic_median_s"] if r["bic_median_s"] > 0 else float("inf") + print( + f"{r['offsets_name']:>10} {r['n_offsets']:>3d} {str(r['ignore']):>6}" + f" {'OK' if r['ok_affs'] else 'FAIL':>5}" + f" {'OK' if r['ok_mask'] else 'FAIL':>5}" + f" {r['bic_median_s']:>9.4f} {r['affo_median_s']:>9.4f}" + f" {speedup:>7.2f}x" + ) + all_ok = all_ok and r["ok_affs"] and r["ok_mask"] + + if not all_ok: + print("\nFAIL: output mismatch", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/affinities/compute_affinities.hxx b/include/bioimage_cpp/affinities/compute_affinities.hxx new file mode 100644 index 0000000..92950d2 --- /dev/null +++ b/include/bioimage_cpp/affinities/compute_affinities.hxx @@ -0,0 +1,191 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/threading.hxx" + +#include +#include +#include +#include +#include +#include + +// Pairwise affinity computation from a label volume. Mirrors affogato's +// `affinities/affinities.hxx::compute_affinities` but without the xtensor +// dependency, with proper input validation at the binding boundary, and with +// optional output mask + optional per-offset parallelism. +// +// For each spatial coordinate ``c`` and offset index ``oi``, +// ``affs[oi, c] = 1`` iff ``labels[c] == labels[c + offsets[oi]]``, else 0. +// ``mask[oi, c] = 1`` iff the offset stays in bounds and neither endpoint +// equals ``ignore_label``; out-of-bounds and ignore-label positions produce +// ``affs = 0`` and ``mask = 0``. +namespace bioimage_cpp::affinities { + +// Boolean affinities on a 2D label volume. Preconditions (validated in the +// binding layer): +// * labels.ndim() == 2 +// * affs.shape == {n_offsets, labels.shape[0], labels.shape[1]} +// * if mask is non-null: mask->shape == affs.shape +// * each entry of offsets has length 2 +template +void compute_affinities_2d( + const ConstArrayView &labels, + const std::vector> &offsets, + const ArrayView &affs, + const ArrayView *mask, + const std::optional ignore_label, + const std::size_t number_of_threads = 1 +) { + const auto height = labels.shape[0]; + const auto width = labels.shape[1]; + const auto plane = height * width; + const auto number_of_offsets = offsets.size(); + + const auto n_threads = ::bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_offsets + ); + + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_offsets, + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t oi = begin; oi < end; ++oi) { + const auto dy = offsets[oi][0]; + const auto dx = offsets[oi][1]; + + AffT * const affs_channel = + affs.data + static_cast(oi) * plane; + std::uint8_t * const mask_channel = (mask != nullptr) + ? mask->data + static_cast(oi) * plane + : nullptr; + + std::fill_n(affs_channel, plane, AffT{0}); + if (mask_channel != nullptr) { + std::fill_n(mask_channel, plane, std::uint8_t{0}); + } + + // Sub-rectangle where (y, x) AND (y+dy, x+dx) are both in bounds. + const auto y_begin = std::max(0, -dy); + const auto y_end = height - std::max(0, dy); + const auto x_begin = std::max(0, -dx); + const auto x_end = width - std::max(0, dx); + if (y_begin >= y_end || x_begin >= x_end) { + continue; + } + + for (std::ptrdiff_t y = y_begin; y < y_end; ++y) { + const auto ny = y + dy; + const LabelT * const row = labels.data + y * width; + const LabelT * const neighbor_row = labels.data + ny * width; + AffT * const out_row = affs_channel + y * width; + std::uint8_t * const mask_row = (mask_channel != nullptr) + ? mask_channel + y * width : nullptr; + for (std::ptrdiff_t x = x_begin; x < x_end; ++x) { + const LabelT a = row[x]; + const LabelT b = neighbor_row[x + dx]; + if (ignore_label.has_value()) { + const LabelT ig = *ignore_label; + if (a == ig || b == ig) { + continue; // affs/mask already 0 + } + } + out_row[x] = (a == b) ? AffT{1} : AffT{0}; + if (mask_row != nullptr) { + mask_row[x] = 1; + } + } + } + } + } + ); +} + +// Boolean affinities on a 3D label volume. Preconditions identical to the 2D +// case with one extra axis. +template +void compute_affinities_3d( + const ConstArrayView &labels, + const std::vector> &offsets, + const ArrayView &affs, + const ArrayView *mask, + const std::optional ignore_label, + const std::size_t number_of_threads = 1 +) { + const auto depth = labels.shape[0]; + const auto height = labels.shape[1]; + const auto width = labels.shape[2]; + const auto volume = depth * height * width; + const auto plane = height * width; + const auto number_of_offsets = offsets.size(); + + const auto n_threads = ::bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_offsets + ); + + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_offsets, + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t oi = begin; oi < end; ++oi) { + const auto dz = offsets[oi][0]; + const auto dy = offsets[oi][1]; + const auto dx = offsets[oi][2]; + + AffT * const affs_channel = + affs.data + static_cast(oi) * volume; + std::uint8_t * const mask_channel = (mask != nullptr) + ? mask->data + static_cast(oi) * volume + : nullptr; + + std::fill_n(affs_channel, volume, AffT{0}); + if (mask_channel != nullptr) { + std::fill_n(mask_channel, volume, std::uint8_t{0}); + } + + const auto z_begin = std::max(0, -dz); + const auto z_end = depth - std::max(0, dz); + const auto y_begin = std::max(0, -dy); + const auto y_end = height - std::max(0, dy); + const auto x_begin = std::max(0, -dx); + const auto x_end = width - std::max(0, dx); + if (z_begin >= z_end || y_begin >= y_end || x_begin >= x_end) { + continue; + } + + for (std::ptrdiff_t z = z_begin; z < z_end; ++z) { + const auto nz = z + dz; + const LabelT * const slab = labels.data + z * plane; + const LabelT * const neighbor_slab = labels.data + nz * plane; + AffT * const out_slab = affs_channel + z * plane; + std::uint8_t * const mask_slab = (mask_channel != nullptr) + ? mask_channel + z * plane : nullptr; + for (std::ptrdiff_t y = y_begin; y < y_end; ++y) { + const auto ny = y + dy; + const LabelT * const row = slab + y * width; + const LabelT * const neighbor_row = neighbor_slab + ny * width; + AffT * const out_row = out_slab + y * width; + std::uint8_t * const mask_row = (mask_slab != nullptr) + ? mask_slab + y * width : nullptr; + for (std::ptrdiff_t x = x_begin; x < x_end; ++x) { + const LabelT a = row[x]; + const LabelT b = neighbor_row[x + dx]; + if (ignore_label.has_value()) { + const LabelT ig = *ignore_label; + if (a == ig || b == ig) { + continue; + } + } + out_row[x] = (a == b) ? AffT{1} : AffT{0}; + if (mask_row != nullptr) { + mask_row[x] = 1; + } + } + } + } + } + } + ); +} + +} // namespace bioimage_cpp::affinities diff --git a/src/bindings/affinities.cxx b/src/bindings/affinities.cxx new file mode 100644 index 0000000..7b3465c --- /dev/null +++ b/src/bindings/affinities.cxx @@ -0,0 +1,232 @@ +#include "affinities.hxx" + +#include "bioimage_cpp/affinities/compute_affinities.hxx" +#include "bioimage_cpp/array_view.hxx" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +template +using LabelArray = nb::ndarray; +using FloatArray = nb::ndarray; +using UInt8Array = nb::ndarray; + +template +std::vector ndarray_shape(const T &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; +} + +FloatArray make_float_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new float[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return FloatArray(data, shape.size(), shape.data(), owner); +} + +UInt8Array make_uint8_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new std::uint8_t[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return UInt8Array(data, shape.size(), shape.data(), owner); +} + +template +std::vector> validate_offsets( + const std::vector> &offsets +) { + std::vector> result; + result.reserve(offsets.size()); + for (std::size_t i = 0; i < offsets.size(); ++i) { + if (offsets[i].size() != D) { + throw std::invalid_argument( + "each offset must have length matching the spatial ndim, got " + "spatial ndim=" + std::to_string(D) + + ", offset[" + std::to_string(i) + "].length=" + + std::to_string(offsets[i].size()) + ); + } + std::array entry{}; + for (std::size_t axis = 0; axis < D; ++axis) { + entry[axis] = offsets[i][axis]; + } + result.push_back(entry); + } + return result; +} + +template +std::pair compute_affinities_nd_t( + LabelArray labels, + const std::vector> &offsets, + const std::optional ignore_label, + const bool return_mask, + const std::size_t number_of_threads +) { + if (labels.ndim() != D) { + throw std::invalid_argument( + "labels must have ndim=" + std::to_string(D) + + ", got ndim=" + std::to_string(labels.ndim()) + ); + } + if (offsets.empty()) { + throw std::invalid_argument("offsets must not be empty"); + } + auto offsets_typed = validate_offsets(offsets); + + const auto labels_shape = ndarray_shape(labels); + + std::vector out_shape; + out_shape.reserve(D + 1); + out_shape.push_back(offsets.size()); + for (std::size_t axis = 0; axis < D; ++axis) { + out_shape.push_back(static_cast(labels_shape[axis])); + } + + std::vector out_view_shape(out_shape.size()); + for (std::size_t axis = 0; axis < out_shape.size(); ++axis) { + out_view_shape[axis] = static_cast(out_shape[axis]); + } + + auto affs = make_float_array(out_shape); + UInt8Array mask = return_mask + ? make_uint8_array(out_shape) + : UInt8Array(nullptr, 0, nullptr); + + ConstArrayView labels_view{ + labels.data(), + labels_shape, + {}, + }; + ArrayView affs_view{ + affs.data(), + out_view_shape, + {}, + }; + ArrayView mask_view{ + return_mask ? mask.data() : nullptr, + out_view_shape, + {}, + }; + const ArrayView *mask_ptr = return_mask ? &mask_view : nullptr; + + { + nb::gil_scoped_release release; + if constexpr (D == 2) { + affinities::compute_affinities_2d( + labels_view, offsets_typed, affs_view, mask_ptr, + ignore_label, number_of_threads + ); + } else if constexpr (D == 3) { + affinities::compute_affinities_3d( + labels_view, offsets_typed, affs_view, mask_ptr, + ignore_label, number_of_threads + ); + } + } + + return {std::move(affs), std::move(mask)}; +} + +} // namespace + +void bind_affinities(nb::module_ &m) { + m.def( + "_compute_affinities_2d_uint32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_uint64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_int32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_int64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_uint32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_uint64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_int32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_int64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/affinities.hxx b/src/bindings/affinities.hxx new file mode 100644 index 0000000..e261a4b --- /dev/null +++ b/src/bindings/affinities.hxx @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace bioimage_cpp::bindings { + +void bind_affinities(nanobind::module_ &m); + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx index cee23a9..bab0a1e 100644 --- a/src/bindings/module.cxx +++ b/src/bindings/module.cxx @@ -1,3 +1,4 @@ +#include "affinities.hxx" #include "blocking.hxx" #include "graph.hxx" #include "ground_truth.hxx" @@ -10,6 +11,7 @@ namespace nb = nanobind; NB_MODULE(_core, m) { m.doc() = "C++ extension module for bioimage_cpp."; + bioimage_cpp::bindings::bind_affinities(m); bioimage_cpp::bindings::bind_blocking(m); bioimage_cpp::bindings::bind_graph(m); bioimage_cpp::bindings::bind_ground_truth(m); diff --git a/src/bioimage_cpp/__init__.py b/src/bioimage_cpp/__init__.py index f8a5e6f..0fe3a83 100644 --- a/src/bioimage_cpp/__init__.py +++ b/src/bioimage_cpp/__init__.py @@ -2,6 +2,7 @@ from ._version import __version__ from ._core import Block, Blocking, BlockWithHalo +from . import affinities from . import graph from . import ground_truth from . import segmentation @@ -12,6 +13,7 @@ "Block", "Blocking", "BlockWithHalo", + "affinities", "graph", "ground_truth", "segmentation", diff --git a/src/bioimage_cpp/_data.py b/src/bioimage_cpp/_data.py index 7993f68..b234406 100644 --- a/src/bioimage_cpp/_data.py +++ b/src/bioimage_cpp/_data.py @@ -214,3 +214,26 @@ def load_isbi_raw( with h5py.File(affinity_path(timeout=timeout), "r") as f: raw = f["raw"][:] return np.ascontiguousarray(raw) + + +def load_isbi_gt_segmentation( + *, + timeout: Optional[float] = None, +) -> np.ndarray: + """Load the registered ISBI ground-truth segmentation volume. + + The labels are stored in the same HDF5 file as the affinities under key + ``labels/gt_segmentation``. Returned as a C-contiguous ``uint64`` array + with shape ``(30, 512, 512)`` (~7.9 M voxels). + """ + try: + import h5py + except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "h5py is required to load the registered ISBI segmentation file. " + "Install it with `pip install h5py`." + ) from error + + with h5py.File(affinity_path(timeout=timeout), "r") as f: + labels = f["labels/gt_segmentation"][:] + return np.ascontiguousarray(labels) diff --git a/src/bioimage_cpp/affinities/__init__.py b/src/bioimage_cpp/affinities/__init__.py new file mode 100644 index 0000000..0406fc7 --- /dev/null +++ b/src/bioimage_cpp/affinities/__init__.py @@ -0,0 +1,5 @@ +"""Pairwise affinities from label volumes.""" + +from .compute_affinities import compute_affinities + +__all__ = ["compute_affinities"] diff --git a/src/bioimage_cpp/affinities/compute_affinities.py b/src/bioimage_cpp/affinities/compute_affinities.py new file mode 100644 index 0000000..9d2c870 --- /dev/null +++ b/src/bioimage_cpp/affinities/compute_affinities.py @@ -0,0 +1,127 @@ +"""Pairwise boolean affinities from a label volume.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import overload + +import numpy as np + +from .. import _core + + +_COMPUTE_AFFINITIES_2D_BY_DTYPE = { + np.dtype("uint32"): _core._compute_affinities_2d_uint32, + np.dtype("uint64"): _core._compute_affinities_2d_uint64, + np.dtype("int32"): _core._compute_affinities_2d_int32, + np.dtype("int64"): _core._compute_affinities_2d_int64, +} + +_COMPUTE_AFFINITIES_3D_BY_DTYPE = { + np.dtype("uint32"): _core._compute_affinities_3d_uint32, + np.dtype("uint64"): _core._compute_affinities_3d_uint64, + np.dtype("int32"): _core._compute_affinities_3d_int32, + np.dtype("int64"): _core._compute_affinities_3d_int64, +} + + +@overload +def compute_affinities( + labels: np.ndarray, + offsets: Sequence[Sequence[int]] | np.ndarray, + *, + ignore_label: int | None = None, + return_mask: bool = True, + number_of_threads: int = 1, +) -> tuple[np.ndarray, np.ndarray]: ... + + +def compute_affinities( + labels: np.ndarray, + offsets: Sequence[Sequence[int]] | np.ndarray, + *, + ignore_label: int | None = None, + return_mask: bool = True, + number_of_threads: int = 1, +): + """Compute boolean pairwise affinities from a label volume. + + For each spatial coordinate ``c`` and offset index ``oi``, + ``affinities[oi, c]`` is ``1.0`` if ``labels[c] == labels[c + offsets[oi]]`` + (the two voxels are in the same cluster) and ``0.0`` otherwise. + + Parameters + ---------- + labels: + 2D or 3D integer label volume. Supported dtypes are ``uint32``, + ``uint64``, ``int32``, ``int64``. Non-contiguous arrays are copied + to a C-contiguous buffer first. + offsets: + Shape ``(n_offsets, ndim)``. Each offset is a per-axis displacement, + in NumPy axis order, applied at each voxel to find the neighbor. + ignore_label: + If given, any pair where either endpoint has this label produces + ``affinity = 0`` and ``mask = 0`` (treated as out-of-volume). + return_mask: + When ``True`` (default), also return a ``uint8`` validity mask of + the same shape as the affinities: ``1`` for in-bounds non-ignored + pairs, ``0`` otherwise. Set to ``False`` to skip the allocation + when only the affinities are needed. + number_of_threads: + Number of threads to parallelize over the offset channels. + + Returns + ------- + affinities : np.ndarray + ``float32`` array of shape ``(n_offsets, *labels.shape)``. + mask : np.ndarray, only if ``return_mask`` is ``True`` + ``uint8`` array of shape ``(n_offsets, *labels.shape)``. + """ + array = np.ascontiguousarray(labels) + if array.ndim not in (2, 3): + raise ValueError( + "labels must be 2D or 3D, got ndim=" + str(array.ndim) + ) + + table = ( + _COMPUTE_AFFINITIES_2D_BY_DTYPE if array.ndim == 2 + else _COMPUTE_AFFINITIES_3D_BY_DTYPE + ) + try: + run = table[array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in table) + raise TypeError( + f"labels must have one of dtypes ({supported}), got dtype={array.dtype}" + ) from error + + normalized_offsets = [ + [int(value) for value in offset] for offset in np.asarray(offsets).tolist() + ] + if len(normalized_offsets) == 0: + raise ValueError("offsets must not be empty") + if any(len(offset) != array.ndim for offset in normalized_offsets): + raise ValueError( + "each offset must have length matching the spatial ndim, got " + f"spatial ndim={array.ndim}" + ) + + n_threads = int(number_of_threads) + if n_threads < 1: + raise ValueError("number_of_threads must be >= 1") + + if ignore_label is None: + typed_ignore: int | None = None + else: + typed_ignore = int(ignore_label) + + affs, mask = run( + array, + normalized_offsets, + typed_ignore, + bool(return_mask), + n_threads, + ) + if return_mask: + return affs, mask + return affs diff --git a/tests/affinities/test_compute_affinities.py b/tests/affinities/test_compute_affinities.py new file mode 100644 index 0000000..6e25db5 --- /dev/null +++ b/tests/affinities/test_compute_affinities.py @@ -0,0 +1,154 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _numpy_reference(labels, offsets, ignore_label=None): + """Slow but obvious reference: nested Python loops over voxels.""" + labels = np.asarray(labels) + n = len(offsets) + affs = np.zeros((n,) + labels.shape, dtype=np.float32) + mask = np.zeros((n,) + labels.shape, dtype=np.uint8) + for oi, offset in enumerate(offsets): + it = np.ndindex(*labels.shape) + for coord in it: + neighbor = tuple(c + d for c, d in zip(coord, offset)) + if any(n < 0 or n >= s for n, s in zip(neighbor, labels.shape)): + continue + a = labels[coord] + b = labels[neighbor] + if ignore_label is not None and (a == ignore_label or b == ignore_label): + continue + affs[(oi, *coord)] = 1.0 if a == b else 0.0 + mask[(oi, *coord)] = 1 + return affs, mask + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_2d_matches_numpy_reference(dtype): + rng = np.random.default_rng(0) + labels = rng.integers(0, 5, size=(7, 11)).astype(dtype) + offsets = [[0, 1], [1, 0], [1, 1], [2, -3], [-1, 2]] + + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + + assert affs.dtype == np.float32 + assert mask.dtype == np.uint8 + assert affs.shape == (len(offsets), *labels.shape) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_3d_matches_numpy_reference(dtype): + rng = np.random.default_rng(1) + labels = rng.integers(0, 4, size=(4, 5, 6)).astype(dtype) + offsets = [[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 1], [-2, 0, 3]] + + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_ignore_label_masks_pairs_at_ignored_voxels(): + labels = np.array( + [ + [0, 1, 1, 2], + [0, 1, 2, 2], + ], + dtype=np.uint32, + ) + offsets = [[0, 1]] + + affs, mask = bic.affinities.compute_affinities( + labels, offsets, ignore_label=0 + ) + ref_affs, ref_mask = _numpy_reference(labels, offsets, ignore_label=0) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + assert mask[0, 0, 0] == 0 # voxel has ignore label + assert affs[0, 0, 0] == 0.0 + + +def test_offset_completely_out_of_bounds_yields_zero_mask(): + labels = np.ones((3, 3), dtype=np.uint32) + offsets = [[10, 0]] # never in bounds + affs, mask = bic.affinities.compute_affinities(labels, offsets) + assert affs.sum() == 0 + assert mask.sum() == 0 + + +def test_return_mask_false_returns_only_affinities(): + labels = np.array([[0, 0], [1, 1]], dtype=np.uint32) + offsets = [[0, 1], [1, 0]] + + affs = bic.affinities.compute_affinities(labels, offsets, return_mask=False) + assert isinstance(affs, np.ndarray) + assert affs.shape == (2, 2, 2) + assert affs.dtype == np.float32 + + +def test_negative_offsets_work(): + labels = np.array([[1, 2], [1, 2]], dtype=np.uint32) + offsets = [[0, -1], [-1, 0]] + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_threading_does_not_change_output(): + rng = np.random.default_rng(42) + labels = rng.integers(0, 7, size=(8, 12)).astype(np.uint32) + offsets = [[0, 1], [1, 0], [1, 1], [2, 3], [-1, 2]] + + affs_single, mask_single = bic.affinities.compute_affinities( + labels, offsets, number_of_threads=1 + ) + affs_multi, mask_multi = bic.affinities.compute_affinities( + labels, offsets, number_of_threads=4 + ) + np.testing.assert_array_equal(affs_single, affs_multi) + np.testing.assert_array_equal(mask_single, mask_multi) + + +def test_non_contiguous_input_is_handled(): + labels = np.array([[0, 0, 1], [0, 1, 1]], dtype=np.uint32) + # Transpose to produce a non-contiguous view. + labels_T = labels.T + assert not labels_T.flags["C_CONTIGUOUS"] + affs, mask = bic.affinities.compute_affinities(labels_T, [[0, 1]]) + # Should give same result as feeding a contiguous copy. + ref_affs, ref_mask = bic.affinities.compute_affinities( + np.ascontiguousarray(labels_T), [[0, 1]] + ) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_rejects_unsupported_dtype(): + labels = np.zeros((4, 4), dtype=np.float32) + with pytest.raises(TypeError, match="dtypes"): + bic.affinities.compute_affinities(labels, [[0, 1]]) + + +def test_rejects_1d_input(): + labels = np.zeros(8, dtype=np.uint32) + with pytest.raises(ValueError, match="2D or 3D"): + bic.affinities.compute_affinities(labels, [[1]]) + + +def test_rejects_empty_offsets(): + labels = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="offsets"): + bic.affinities.compute_affinities(labels, []) + + +def test_rejects_offset_with_wrong_length(): + labels = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="spatial ndim"): + bic.affinities.compute_affinities(labels, [[0, 1, 2]])