diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 7c4277f..c775763 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -1513,6 +1513,56 @@ Notes: - `number_of_threads=0` uses the library default; pass a positive integer for a fixed thread count. +### Mapping RAG Edges to Pixel Coordinates + +Nifty's `ragCoordinates` scans the label volume once and caches, per edge, the +pixel coordinates of the boundary between the two adjacent regions, so that +per-edge values (e.g. edge probabilities) can be painted back onto a volume. + +Nifty: + +```python +import nifty.graph.rag as nrag + +rag = nrag.gridRag(labels) +rag_coords = nrag.ragCoordinates(rag) +storage = rag_coords.storageLengths() +volume = rag_coords.edgesToVolume(edge_values, edgeDirection=0) +``` + +bioimage-cpp: + +```python +rag = bic.graph.region_adjacency_graph(labels) +rag_coords = bic.graph.rag_coordinates(rag, labels) + +storage = rag_coords.storage_lengths() +coords = rag_coords.edge_coordinates(edge_id) # (n_points, ndim) +volume = rag_coords.edges_to_volume(edge_values, edge_direction=0) +``` + +Notes: + +- `labels` must be the over-segmentation used to construct `rag`, and is passed + explicitly (nifty's RAG holds an internal reference to it; ours does not). + Supported label dtypes: `uint32`, `uint64`, `int32`, `int64`. +- A boundary "contact" is a pair of directly adjacent pixels with different + labels. Each contact contributes two coordinates to its edge: the + lower-coordinate pixel followed by its `+axis` neighbor. `storage_lengths` + therefore reports `2 * n_contacts` per edge, and coordinates are stored in + scan order (NumPy axis order, C-contiguous). +- `edge_direction` selects which side(s) to report / paint: `0` = both (default), + `1` = lower-coordinate pixel only, `2` = higher-coordinate pixel only. +- `edges_to_volume` returns a volume of the label shape and dtype matching + `edge_values` (supported: `float32`, `float64`, `uint32`, `uint64`). + Non-boundary pixels are set to `ignore_value`. Painting is sequential in + ascending edge id, so where several edges' boundaries coincide on a pixel the + highest edge id wins — a deterministic, race-free rule (nifty's parallel + `edgesToVolume` does not guarantee a tie-break order). +- The cached object can be reused across many `edges_to_volume` calls without + re-scanning the labels. `number_of_threads=0` (on `rag_coordinates`) uses the + library default. + ### Accumulating Labels on a RAG Nifty's `gridRagAccumulateLabels` projects a second label volume onto a RAG diff --git a/include/bioimage_cpp/graph/rag_coordinates.hxx b/include/bioimage_cpp/graph/rag_coordinates.hxx new file mode 100644 index 0000000..cfe56a8 --- /dev/null +++ b/include/bioimage_cpp/graph/rag_coordinates.hxx @@ -0,0 +1,381 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/threading.hxx" +#include "bioimage_cpp/graph/region_adjacency_graph.hxx" + +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +namespace detail_rag_coordinates { + +// Visit every forward-neighbor contact (a pair of adjacent pixels with +// different labels) whose lower pixel lies in rows [begin, end) of a 2D label +// image. `on_contact(edge_id, low_flat, high_flat)` is called for each, where +// `low_flat` is the current pixel and `high_flat` the +axis neighbor. +template +void scan_contacts_2d( + const T *data, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, + const RegionAdjacencyGraph &rag, + OnContact &&on_contact +) { + for (std::size_t y = y_begin; y < y_end; ++y) { + const auto row_offset = y * width; + for (std::size_t x = 0; x < width; ++x) { + const auto pixel = row_offset + x; + const auto u = static_cast(data[pixel]); + if (x + 1 < width) { + const auto v = static_cast(data[pixel + 1]); + if (u != v) { + const auto edge = rag.find_edge(u, v); + if (edge >= 0) { + on_contact(static_cast(edge), pixel, pixel + 1); + } + } + } + if (y + 1 < height) { + const auto v = static_cast(data[pixel + width]); + if (u != v) { + const auto edge = rag.find_edge(u, v); + if (edge >= 0) { + on_contact(static_cast(edge), pixel, pixel + width); + } + } + } + } + } +} + +// 3D analogue of scan_contacts_2d; chunks by the leading (z) axis. +template +void scan_contacts_3d( + const T *data, + 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 RegionAdjacencyGraph &rag, + OnContact &&on_contact +) { + const auto slice_size = height * width; + for (std::size_t z = z_begin; z < z_end; ++z) { + const auto slice_offset = z * slice_size; + for (std::size_t y = 0; y < height; ++y) { + const auto row_offset = slice_offset + y * width; + for (std::size_t x = 0; x < width; ++x) { + const auto pixel = row_offset + x; + const auto u = static_cast(data[pixel]); + if (x + 1 < width) { + const auto v = static_cast(data[pixel + 1]); + if (u != v) { + const auto edge = rag.find_edge(u, v); + if (edge >= 0) { + on_contact(static_cast(edge), pixel, pixel + 1); + } + } + } + if (y + 1 < height) { + const auto v = static_cast(data[pixel + width]); + if (u != v) { + const auto edge = rag.find_edge(u, v); + if (edge >= 0) { + on_contact(static_cast(edge), pixel, pixel + width); + } + } + } + if (z + 1 < depth) { + const auto v = static_cast(data[pixel + slice_size]); + if (u != v) { + const auto edge = rag.find_edge(u, v); + if (edge >= 0) { + on_contact(static_cast(edge), pixel, pixel + slice_size); + } + } + } + } + } + } +} + +} // namespace detail_rag_coordinates + +// Maps the edges of a RegionAdjacencyGraph back to the pixel/voxel coordinates +// of the boundary between the two adjacent regions. Scans the label volume once +// at construction and caches a CSR-style mapping so the result can be reused +// across many `edges_to_volume` calls. +// +// Each boundary "contact" — a pair of directly adjacent pixels with different +// labels — contributes two consecutive flat (C-order) pixel indices to its +// edge: the lower-coordinate pixel followed by its +axis neighbor. Within an +// edge's slice even local indices are therefore lower-side pixels and odd local +// indices are higher-side. The `edge_direction` parameter selects which side(s) +// to report: 0 = both (default), 1 = lower-side only, 2 = higher-side only. +class RagCoordinates { +public: + template + RagCoordinates( + const RegionAdjacencyGraph &rag, + const ConstArrayView &labels, + const std::size_t number_of_threads + ) { + if (labels.ndim() != 2 && labels.ndim() != 3) { + throw std::invalid_argument( + "labels must be a 2D or 3D array, got ndim=" + + std::to_string(labels.ndim()) + ); + } + if (static_cast(labels.ndim()) != rag.shape().size()) { + throw std::invalid_argument("labels ndim must match rag shape"); + } + for (std::size_t axis = 0; axis < rag.shape().size(); ++axis) { + if (labels.shape[axis] != static_cast(rag.shape()[axis])) { + throw std::invalid_argument("labels shape must match rag shape"); + } + } + + shape_ = rag.shape(); + const auto n_edges = static_cast(rag.number_of_edges()); + + const auto work_items = static_cast(labels.shape[0]); + const auto n_threads = + bioimage_cpp::detail::normalize_thread_count(number_of_threads, work_items); + + // Pass 1: count contacts per edge, per thread (each contact = 2 points). + std::vector> per_thread_counts( + n_threads, std::vector(n_edges, 0) + ); + run_scan(labels, rag, n_threads, [&per_thread_counts](std::size_t thread_id) { + std::uint64_t *counts = per_thread_counts[thread_id].data(); + return [counts](std::uint64_t edge, std::uint64_t, std::uint64_t) { + counts[edge] += 2; + }; + }); + + // Build CSR offsets (in points) and per-thread write bases within each + // edge slice so the fill preserves global scan order deterministically. + offsets_.assign(n_edges + 1, 0); + std::vector edge_total(n_edges, 0); + for (std::size_t edge = 0; edge < n_edges; ++edge) { + std::uint64_t total = 0; + for (std::size_t thread_id = 0; thread_id < n_threads; ++thread_id) { + total += per_thread_counts[thread_id][edge]; + } + edge_total[edge] = total; + offsets_[edge + 1] = offsets_[edge] + total; + } + points_.assign(static_cast(offsets_[n_edges]), 0); + + // cursor[thread][edge] = next write position (point index) for that + // thread within the edge's slice. + std::vector> cursor( + n_threads, std::vector(n_edges, 0) + ); + for (std::size_t edge = 0; edge < n_edges; ++edge) { + std::uint64_t base = offsets_[edge]; + for (std::size_t thread_id = 0; thread_id < n_threads; ++thread_id) { + cursor[thread_id][edge] = base; + base += per_thread_counts[thread_id][edge]; + } + } + + // Pass 2: fill points at the computed cursors. + std::uint64_t *points = points_.data(); + run_scan(labels, rag, n_threads, [&cursor, points](std::size_t thread_id) { + std::uint64_t *cur = cursor[thread_id].data(); + return [cur, points](std::uint64_t edge, std::uint64_t low, std::uint64_t high) { + points[cur[edge]++] = low; + points[cur[edge]++] = high; + }; + }); + } + + [[nodiscard]] std::size_t ndim() const { + return shape_.size(); + } + + [[nodiscard]] const std::vector &shape() const { + return shape_; + } + + [[nodiscard]] std::uint64_t number_of_edges() const { + return offsets_.empty() ? 0 : static_cast(offsets_.size() - 1); + } + + // Number of stored points (= 2 * number of contacts) per edge. + [[nodiscard]] std::vector storage_lengths() const { + const auto n_edges = static_cast(number_of_edges()); + std::vector lengths(n_edges, 0); + for (std::size_t edge = 0; edge < n_edges; ++edge) { + lengths[edge] = offsets_[edge + 1] - offsets_[edge]; + } + return lengths; + } + + // Decoded boundary coordinates of one edge, flattened as n_points * ndim in + // C-order. `edge_direction` selects which side(s) are reported. + [[nodiscard]] std::vector edge_coordinates( + const std::uint64_t edge, + const int edge_direction + ) const { + check_edge(edge); + check_edge_direction(edge_direction); + const auto strides = make_strides(); + const auto begin = offsets_[edge]; + const auto end = offsets_[edge + 1]; + + std::vector coordinates; + for (auto index = begin; index < end; ++index) { + if (!selected(index - begin, edge_direction)) { + continue; + } + decode(points_[index], strides, coordinates); + } + return coordinates; + } + + // Paint per-edge values into `out` (shape == this->shape()). Every pixel is + // first set to `ignore_value`; then each selected boundary point receives + // its edge's value. Painting is sequential in ascending edge id, so where + // the boundaries of several edges coincide on the same pixel the highest + // edge id wins — a deterministic, race-free rule. + template + void edges_to_volume( + const std::vector &edge_values, + ArrayView &out, + const int edge_direction, + const V ignore_value + ) const { + const auto n_edges = static_cast(number_of_edges()); + if (edge_values.size() != n_edges) { + throw std::invalid_argument("edge_values length must match number_of_edges"); + } + check_edge_direction(edge_direction); + if (static_cast(out.ndim()) != shape_.size()) { + throw std::invalid_argument("out ndim must match rag shape"); + } + std::size_t total = 1; + for (std::size_t axis = 0; axis < shape_.size(); ++axis) { + if (out.shape[axis] != static_cast(shape_[axis])) { + throw std::invalid_argument("out shape must match rag shape"); + } + total *= shape_[axis]; + } + + auto *data = out.data; + for (std::size_t index = 0; index < total; ++index) { + data[index] = ignore_value; + } + + for (std::size_t edge = 0; edge < n_edges; ++edge) { + const auto value = edge_values[edge]; + const auto slice_begin = offsets_[edge]; + const auto slice_end = offsets_[edge + 1]; + for (auto index = slice_begin; index < slice_end; ++index) { + if (selected(index - slice_begin, edge_direction)) { + data[points_[index]] = value; + } + } + } + } + +private: + template + void run_scan( + const ConstArrayView &labels, + const RegionAdjacencyGraph &rag, + const std::size_t n_threads, + MakeCallback &&make_callback + ) const { + namespace drc = detail_rag_coordinates; + const auto work_items = static_cast(labels.shape[0]); + bioimage_cpp::detail::parallel_for_chunks( + n_threads, + work_items, + [&](std::size_t thread_id, std::size_t begin, std::size_t end) { + auto on_contact = make_callback(thread_id); + if (labels.ndim() == 2) { + drc::scan_contacts_2d( + labels.data, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, + end, + rag, + on_contact + ); + } else { + drc::scan_contacts_3d( + labels.data, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, + end, + rag, + on_contact + ); + } + } + ); + } + + [[nodiscard]] std::vector make_strides() const { + std::vector shape(shape_.size()); + for (std::size_t axis = 0; axis < shape_.size(); ++axis) { + shape[axis] = static_cast(shape_[axis]); + } + return bioimage_cpp::detail::c_order_strides(shape); + } + + void decode( + const std::uint64_t flat, + const std::vector &strides, + std::vector &out + ) const { + auto remainder = flat; + for (std::size_t axis = 0; axis < shape_.size(); ++axis) { + const auto stride = static_cast(strides[axis]); + out.push_back(remainder / stride); + remainder %= stride; + } + } + + static bool selected(const std::uint64_t local_index, const int edge_direction) { + if (edge_direction == 0) { + return true; + } + if (edge_direction == 1) { + return (local_index % 2) == 0; + } + return (local_index % 2) == 1; + } + + static void check_edge_direction(const int edge_direction) { + if (edge_direction < 0 || edge_direction > 2) { + throw std::invalid_argument("edge_direction must be 0, 1, or 2"); + } + } + + void check_edge(const std::uint64_t edge) const { + if (edge >= number_of_edges()) { + throw std::invalid_argument("edge id out of range"); + } + } + + std::vector shape_; + std::vector offsets_; + std::vector points_; +}; + +} // namespace bioimage_cpp::graph diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 7109879..0500b5b 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -24,6 +24,7 @@ #include "bioimage_cpp/graph/proposal_generator.hxx" #include "bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx" #include "bioimage_cpp/graph/proposal_generators/watershed.hxx" +#include "bioimage_cpp/graph/rag_coordinates.hxx" #include "bioimage_cpp/graph/region_adjacency_graph.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" @@ -1336,6 +1337,75 @@ Rag region_adjacency_graph_t( return graph::region_adjacency_graph(labels_view, number_of_threads); } +template +graph::RagCoordinates rag_coordinates_t( + const Rag &rag, + LabelArray labels, + const std::size_t number_of_threads +) { + std::vector shape(labels.ndim()); + for (std::size_t axis = 0; axis < labels.ndim(); ++axis) { + shape[axis] = static_cast(labels.shape(axis)); + } + + ConstArrayView labels_view{ + labels.data(), + shape, + {}, + }; + + nb::gil_scoped_release release; + return graph::RagCoordinates(rag, labels_view, number_of_threads); +} + +UInt64Array rag_coordinates_storage_lengths(const graph::RagCoordinates &coords) { + return vector_to_uint64_array(coords.storage_lengths()); +} + +UInt64Array rag_coordinates_edge_coordinates( + const graph::RagCoordinates &coords, + const std::uint64_t edge, + const int edge_direction +) { + const auto flat = coords.edge_coordinates(edge, edge_direction); + const auto ndim = coords.ndim(); + const std::size_t n_points = ndim == 0 ? 0 : flat.size() / ndim; + auto result = make_uint64_array({n_points, ndim}); + std::copy(flat.begin(), flat.end(), result.data()); + return result; +} + +template +TypedArray rag_coordinates_edges_to_volume( + const graph::RagCoordinates &coords, + ConstFloatingArray edge_values, + const int edge_direction, + const V ignore_value +) { + if (edge_values.ndim() != 1) { + throw std::invalid_argument("edge_values must be a 1D array"); + } + std::vector values( + edge_values.data(), edge_values.data() + edge_values.shape(0) + ); + + const auto &shape = coords.shape(); + std::vector out_shape(shape.size()); + std::vector view_shape(shape.size()); + for (std::size_t axis = 0; axis < shape.size(); ++axis) { + out_shape[axis] = static_cast(shape[axis]); + view_shape[axis] = static_cast(shape[axis]); + } + auto result = make_typed_array(out_shape); + ArrayView view{result.data(), view_shape, {}}; + + { + nb::gil_scoped_release release; + coords.edges_to_volume(values, view, edge_direction, ignore_value); + } + return result; +} + template std::vector ndarray_shape(LabelArray array) { std::vector shape(array.ndim()); @@ -1810,6 +1880,46 @@ void bind_graph(nb::module_ &m) { nb::class_(m, "RegionAdjacencyGraph") .def_prop_ro("shape", &Rag::shape); + nb::class_(m, "RagCoordinates") + .def_prop_ro("ndim", &graph::RagCoordinates::ndim) + .def_prop_ro("shape", &graph::RagCoordinates::shape) + .def_prop_ro("number_of_edges", &graph::RagCoordinates::number_of_edges) + .def("storage_lengths", &rag_coordinates_storage_lengths) + .def( + "edge_coordinates", + &rag_coordinates_edge_coordinates, + nb::arg("edge"), + nb::arg("edge_direction") + ) + .def( + "_edges_to_volume_float32", + &rag_coordinates_edges_to_volume, + nb::arg("edge_values"), + nb::arg("edge_direction"), + nb::arg("ignore_value") + ) + .def( + "_edges_to_volume_float64", + &rag_coordinates_edges_to_volume, + nb::arg("edge_values"), + nb::arg("edge_direction"), + nb::arg("ignore_value") + ) + .def( + "_edges_to_volume_uint32", + &rag_coordinates_edges_to_volume, + nb::arg("edge_values"), + nb::arg("edge_direction"), + nb::arg("ignore_value") + ) + .def( + "_edges_to_volume_uint64", + &rag_coordinates_edges_to_volume, + nb::arg("edge_values"), + nb::arg("edge_direction"), + nb::arg("ignore_value") + ); + const auto register_grid_boundary = [&m](const char *name) { m.def( name, @@ -2229,6 +2339,39 @@ void bind_graph(nb::module_ &m) { "Build a region adjacency graph for int64 labels." ); + m.def( + "_rag_coordinates_uint32", + &rag_coordinates_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("number_of_threads"), + "Build RAG edge coordinates for uint32 labels." + ); + m.def( + "_rag_coordinates_uint64", + &rag_coordinates_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("number_of_threads"), + "Build RAG edge coordinates for uint64 labels." + ); + m.def( + "_rag_coordinates_int32", + &rag_coordinates_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("number_of_threads"), + "Build RAG edge coordinates for int32 labels." + ); + m.def( + "_rag_coordinates_int64", + &rag_coordinates_t, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("number_of_threads"), + "Build RAG edge coordinates for int64 labels." + ); + m.def( "_accumulate_edge_map_features_uint32", &accumulate_edge_map_features_t, diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index cae232e..7a3ffad 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -272,6 +272,94 @@ def project_edge_ids_to_pixels_with_offsets( RegionAdjacencyGraph = _core.RegionAdjacencyGraph +class RagCoordinates: + """Mapping from RAG edges to the pixel coordinates of region boundaries. + + Created via :func:`rag_coordinates`. The label volume is scanned once at + construction and the per-edge boundary coordinates are cached, so the same + object can be reused across many :meth:`edges_to_volume` calls. + + Each boundary "contact" — a pair of directly adjacent pixels with different + labels — contributes two coordinates to its edge: the lower-coordinate pixel + and its ``+axis`` neighbor. The ``edge_direction`` argument selects which + side(s) are reported: ``0`` = both (default), ``1`` = lower-side only, + ``2`` = higher-side only. + """ + + def __init__(self, core): + self._core = core + + @property + def ndim(self) -> int: + return self._core.ndim + + @property + def shape(self): + return self._core.shape + + @property + def number_of_edges(self) -> int: + return self._core.number_of_edges + + def storage_lengths(self) -> np.ndarray: + """Number of stored boundary points per edge (``2 * n_contacts``).""" + return self._core.storage_lengths() + + def edge_coordinates(self, edge: int, *, edge_direction: int = 0) -> np.ndarray: + """Boundary coordinates of one edge as an ``(n_points, ndim)`` array.""" + if edge_direction not in (0, 1, 2): + raise ValueError("edge_direction must be 0, 1, or 2") + return self._core.edge_coordinates(int(edge), int(edge_direction)) + + def edges_to_volume( + self, + edge_values, + *, + edge_direction: int = 0, + ignore_value=0, + ) -> np.ndarray: + """Paint per-edge values onto a volume of the label shape. + + Every pixel is set to ``ignore_value`` and then each selected boundary + point receives its edge's value. ``edge_values`` is a 1D array of length + ``number_of_edges``; supported dtypes are ``float32``, ``float64``, + ``uint32`` and ``uint64``. The returned volume has the same dtype. + + Painting is in ascending edge id, so where several edges' boundaries + share a pixel the highest edge id wins (deterministic). + """ + if edge_direction not in (0, 1, 2): + raise ValueError("edge_direction must be 0, 1, or 2") + values = np.asarray(edge_values) + try: + method_name = _EDGES_TO_VOLUME_BY_DTYPE[values.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _EDGES_TO_VOLUME_BY_DTYPE) + raise TypeError( + f"edge_values must have one of dtypes ({supported}), " + f"got dtype={values.dtype}" + ) from error + if values.ndim != 1: + raise ValueError("edge_values must be a 1D array") + if values.shape[0] != self.number_of_edges: + raise ValueError("edge_values length must match number_of_edges") + method = getattr(self._core, method_name) + return method( + np.ascontiguousarray(values), + int(edge_direction), + values.dtype.type(ignore_value), + ) + + def storageLengths(self) -> np.ndarray: + return self.storage_lengths() + + def edgeCoordinates(self, edge: int, *, edge_direction: int = 0) -> np.ndarray: + return self.edge_coordinates(edge, edge_direction=edge_direction) + + def edgesToVolume(self, edge_values, **kwargs) -> np.ndarray: + return self.edges_to_volume(edge_values, **kwargs) + + def undirected_graph(number_of_nodes: int) -> UndirectedGraph: """Create an :class:`UndirectedGraph`.""" return UndirectedGraph(number_of_nodes) @@ -310,6 +398,22 @@ def grid_graph(shape): } +_RAG_COORDINATES_BY_DTYPE = { + np.dtype("uint32"): _core._rag_coordinates_uint32, + np.dtype("uint64"): _core._rag_coordinates_uint64, + np.dtype("int32"): _core._rag_coordinates_int32, + np.dtype("int64"): _core._rag_coordinates_int64, +} + + +_EDGES_TO_VOLUME_BY_DTYPE = { + np.dtype("float32"): "_edges_to_volume_float32", + np.dtype("float64"): "_edges_to_volume_float64", + np.dtype("uint32"): "_edges_to_volume_uint32", + np.dtype("uint64"): "_edges_to_volume_uint64", +} + + _PROJECT_NODE_LABELS_TO_PIXELS_BY_DTYPE = { np.dtype("uint32"): _core._project_node_labels_to_pixels_uint32, np.dtype("uint64"): _core._project_node_labels_to_pixels_uint64, @@ -485,6 +589,43 @@ def region_adjacency_graph( return run(np.ascontiguousarray(array), number_of_threads) +def rag_coordinates( + rag: RegionAdjacencyGraph, + labels: np.ndarray, + *, + number_of_threads: int = 0, +) -> RagCoordinates: + """Map RAG edges to the pixel coordinates of the region boundaries. + + Scans ``labels`` (the over-segmentation used to construct ``rag``) once and + caches, per edge, the boundary coordinates between the two adjacent regions. + The returned :class:`RagCoordinates` exposes :meth:`~RagCoordinates.storage_lengths`, + :meth:`~RagCoordinates.edge_coordinates`, and + :meth:`~RagCoordinates.edges_to_volume`. + """ + array = np.asarray(labels) + if array.ndim not in (2, 3): + raise ValueError(f"labels must be a 2D or 3D array, got ndim={array.ndim}") + if tuple(int(size) for size in rag.shape) != array.shape: + raise ValueError( + "rag shape must match labels shape, got " + f"rag shape={tuple(rag.shape)}, labels shape={array.shape}" + ) + + dtype = array.dtype + try: + run = _RAG_COORDINATES_BY_DTYPE[dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _RAG_COORDINATES_BY_DTYPE) + raise TypeError( + f"labels must have one of dtypes ({supported}), got dtype={dtype}" + ) from error + + number_of_threads = _normalize_number_of_threads(number_of_threads) + core = run(rag, np.ascontiguousarray(array), number_of_threads) + return RagCoordinates(core) + + def project_node_labels_to_pixels( rag: RegionAdjacencyGraph, labels: np.ndarray, @@ -531,6 +672,7 @@ def project_node_labels_to_pixels( __all__ = [ "GridGraph2D", "GridGraph3D", + "RagCoordinates", "RegionAdjacencyGraph", "UndirectedGraph", "agglomeration", @@ -543,6 +685,7 @@ def project_node_labels_to_pixels( "multicut", "mutex_watershed", "project_node_labels_to_pixels", + "rag_coordinates", "region_adjacency_graph", "undirected_graph", ] diff --git a/tests/graph/test_rag_coordinates.py b/tests/graph/test_rag_coordinates.py new file mode 100644 index 0000000..5128c88 --- /dev/null +++ b/tests/graph/test_rag_coordinates.py @@ -0,0 +1,231 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _labels_2d(): + return np.array( + [ + [1, 1, 2], + [1, 3, 2], + [4, 3, 3], + ], + dtype=np.uint64, + ) + + +def test_rag_coordinates_basic_2d(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + assert isinstance(coords, bic.graph.RagCoordinates) + assert coords.ndim == 2 + assert coords.shape == [3, 3] + assert coords.number_of_edges == rag.number_of_edges + + # uv_ids: [[1,2],[1,3],[1,4],[2,3],[3,4]] + # Each contact contributes 2 points (low + high). + np.testing.assert_array_equal( + coords.storage_lengths(), + np.array([2, 4, 2, 4, 2], dtype=np.uint64), + ) + + +def test_rag_coordinates_edge_coordinates_and_directions(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + # edge 0 is (1, 2): single contact at (0,1)-(0,2). + np.testing.assert_array_equal( + coords.edge_coordinates(0), np.array([[0, 1], [0, 2]]) + ) + np.testing.assert_array_equal( + coords.edge_coordinates(0, edge_direction=1), np.array([[0, 1]]) + ) + np.testing.assert_array_equal( + coords.edge_coordinates(0, edge_direction=2), np.array([[0, 2]]) + ) + + # edge 1 is (1, 3): two contacts (0,1)-(1,1) and (1,0)-(1,1), in scan order. + np.testing.assert_array_equal( + coords.edge_coordinates(1), + np.array([[0, 1], [1, 1], [1, 0], [1, 1]]), + ) + np.testing.assert_array_equal( + coords.edge_coordinates(1, edge_direction=1), + np.array([[0, 1], [1, 0]]), + ) + np.testing.assert_array_equal( + coords.edge_coordinates(1, edge_direction=2), + np.array([[1, 1], [1, 1]]), + ) + + +def test_rag_coordinates_edges_to_volume_both(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + values = np.arange(1, rag.number_of_edges + 1, dtype=np.float32) + volume = coords.edges_to_volume(values, ignore_value=-1.0) + + assert volume.dtype == np.float32 + assert volume.shape == (3, 3) + # Higher edge id wins where boundaries coincide. + np.testing.assert_array_equal( + volume, + np.array( + [ + [-1.0, 2.0, 1.0], + [3.0, 4.0, 4.0], + [5.0, 5.0, 4.0], + ], + dtype=np.float32, + ), + ) + + +def test_rag_coordinates_edges_to_volume_low_side(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + values = np.arange(1, rag.number_of_edges + 1, dtype=np.float64) + volume = coords.edges_to_volume(values, edge_direction=1, ignore_value=0.0) + + assert volume.dtype == np.float64 + # Only the lower-coordinate pixel of each contact is painted; (2,1) is only + # ever a higher-side pixel (for edge 4) and so stays at the ignore value. + np.testing.assert_array_equal( + volume, + np.array( + [ + [0.0, 2.0, 0.0], + [3.0, 4.0, 4.0], + [5.0, 0.0, 0.0], + ] + ), + ) + + +def test_rag_coordinates_edges_to_volume_uint(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + values = (np.arange(rag.number_of_edges) + 10).astype(np.uint64) + volume = coords.edges_to_volume(values, ignore_value=0) + assert volume.dtype == np.uint64 + # Every boundary pixel got some non-zero edge value. + assert volume.max() > 0 + + +def test_rag_coordinates_3d(): + labels = np.array( + [ + [[1, 1], [2, 2]], + [[1, 3], [2, 3]], + ], + dtype=np.uint32, + ) + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels, number_of_threads=1) + + assert coords.ndim == 3 + assert coords.shape == [2, 2, 2] + assert coords.number_of_edges == rag.number_of_edges + + # Every stored coordinate must point to one of the two endpoint labels. + for edge in range(rag.number_of_edges): + u, v = rag.uv(edge) + pts = coords.edge_coordinates(edge) + seen = {int(labels[tuple(p)]) for p in pts} + assert seen <= {int(u), int(v)} + assert seen # not empty + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_rag_coordinates_label_dtypes(dtype): + labels = _labels_2d().astype(dtype) + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + np.testing.assert_array_equal( + coords.storage_lengths(), + np.array([2, 4, 2, 4, 2], dtype=np.uint64), + ) + + +def test_rag_coordinates_parallel_matches_single_thread(): + rng = np.random.default_rng(0) + labels = rng.integers(0, 8, size=(16, 17), dtype=np.int64) + rag = bic.graph.region_adjacency_graph(labels) + + single = bic.graph.rag_coordinates(rag, labels, number_of_threads=1) + parallel = bic.graph.rag_coordinates(rag, labels, number_of_threads=4) + + np.testing.assert_array_equal( + single.storage_lengths(), parallel.storage_lengths() + ) + for edge in range(rag.number_of_edges): + np.testing.assert_array_equal( + single.edge_coordinates(edge), parallel.edge_coordinates(edge) + ) + + +def test_rag_coordinates_storage_lengths_sum(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + # Each inter-region adjacent pixel pair contributes exactly 2 points. + expected_contacts = 0 + for axis in range(labels.ndim): + a = np.take(labels, range(labels.shape[axis] - 1), axis=axis) + b = np.take(labels, range(1, labels.shape[axis]), axis=axis) + expected_contacts += int(np.count_nonzero(a != b)) + assert int(coords.storage_lengths().sum()) == 2 * expected_contacts + + +def test_rag_coordinates_rejects_invalid_input(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + with pytest.raises(ValueError, match="shape"): + bic.graph.rag_coordinates(rag, labels[:, :2]) + + with pytest.raises(TypeError, match="one of dtypes"): + bic.graph.rag_coordinates(rag, labels.astype(np.uint8)) + + with pytest.raises(ValueError, match="length must match"): + coords.edges_to_volume(np.ones(rag.number_of_edges + 1, dtype=np.float32)) + + with pytest.raises(TypeError, match="one of dtypes"): + coords.edges_to_volume(np.ones(rag.number_of_edges, dtype=np.int8)) + + with pytest.raises(ValueError, match="edge_direction"): + coords.edge_coordinates(0, edge_direction=3) + + with pytest.raises(ValueError, match="edge_direction"): + coords.edges_to_volume( + np.ones(rag.number_of_edges, dtype=np.float32), edge_direction=5 + ) + + +def test_rag_coordinates_camelcase_aliases(): + labels = _labels_2d() + rag = bic.graph.region_adjacency_graph(labels) + coords = bic.graph.rag_coordinates(rag, labels) + + np.testing.assert_array_equal(coords.storageLengths(), coords.storage_lengths()) + np.testing.assert_array_equal( + coords.edgeCoordinates(1), coords.edge_coordinates(1) + ) + values = np.arange(1, rag.number_of_edges + 1, dtype=np.float32) + np.testing.assert_array_equal( + coords.edgesToVolume(values, ignore_value=-1.0), + coords.edges_to_volume(values, ignore_value=-1.0), + )