Add rag label accumulation (#43)

Author: constantinpape

MIGRATION_GUIDE.md

@@ -1385,6 +1385,48 @@ Notes:
 - `number_of_threads=0` uses the library default; pass a positive integer for a
   fixed thread count.
 
+### Accumulating Labels on a RAG
+
+Nifty's `gridRagAccumulateLabels` projects a second label volume onto a RAG
+by taking a per-node majority vote (commonly used to project a ground-truth
+segmentation onto an over-segmentation).
+
+Nifty:
+
+```python
+import nifty.graph.rag as nrag
+
+rag = nrag.gridRag(labels)
+node_labels = nrag.gridRagAccumulateLabels(rag, gt)
+# ignore label 0 in the ground truth (nifty's "ignoreBackground"):
+node_labels = nrag.gridRagAccumulateLabels(rag, gt, ignoreBackground=True)
+```
+
+bioimage-cpp:
+
+```python
+rag = bic.graph.region_adjacency_graph(labels)
+node_labels = bic.graph.features.accumulate_labels(rag, labels, gt)
+# arbitrary ignore value (covers nifty's ignoreBackground=True by passing 0):
+node_labels = bic.graph.features.accumulate_labels(
+    rag, labels, gt, ignore_value=0
+)
+```
+
+Notes:
+
+- `labels` must be the over-segmentation used to construct `rag`.
+- `other_labels` must have the same shape as `labels`. Supported dtypes for
+  both arrays: `uint32`, `uint64`, `int32`, `int64`; they may differ.
+- The output has length `rag.number_of_nodes` and the same dtype as
+  `other_labels`. Nodes whose pixels are all ignored receive `0`.
+- Ties in the majority vote are broken by smaller label id (deterministic).
+  Nifty's tie-breaking depends on `std::unordered_map` iteration order and
+  is therefore platform-dependent; `bic` resolves ties deterministically.
+- `ignore_value` is more general than nifty's boolean `ignoreBackground`:
+  pass `0` to reproduce `ignoreBackground=True`, or any other value to skip
+  arbitrary sentinels (e.g. `255` or `-1`).
+
 ### RAG Boundary and Affinity Features
 
 Nifty has RAG feature helpers such as `accumulateEdgeMeanAndLength`,

include/bioimage_cpp/graph/label_accumulation.hxx

@@ -0,0 +1,139 @@
+#pragma once
+
+#include "bioimage_cpp/array_view.hxx"
+#include "bioimage_cpp/detail/threading.hxx"
+#include "bioimage_cpp/graph/node_label_projection.hxx"
+#include "bioimage_cpp/graph/region_adjacency_graph.hxx"
+
+#include <cstddef>
+#include <cstdint>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace bioimage_cpp::graph {
+
+namespace detail_label_accumulation {
+
+inline std::size_t number_of_pixels(const std::vector<std::ptrdiff_t> &shape) {
+    return static_cast<std::size_t>(std::accumulate(
+        shape.begin(),
+        shape.end(),
+        std::ptrdiff_t{1},
+        [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; }
+    ));
+}
+
+template <class LabelT, class OtherT>
+void scan_chunk(
+    const LabelT *labels,
+    const OtherT *other_labels,
+    const std::size_t pixel_begin,
+    const std::size_t pixel_end,
+    const bool has_ignore_value,
+    const OtherT ignore_value,
+    std::vector<std::unordered_map<OtherT, std::uint64_t>> &histograms
+) {
+    for (std::size_t index = pixel_begin; index < pixel_end; ++index) {
+        const auto other = other_labels[index];
+        if (has_ignore_value && other == ignore_value) {
+            continue;
+        }
+        const auto node = detail::checked_label_to_node(labels[index]);
+        ++histograms[static_cast<std::size_t>(node)][other];
+    }
+}
+
+} // namespace detail_label_accumulation
+
+template <class LabelT, class OtherT>
+void accumulate_labels(
+    const RegionAdjacencyGraph &rag,
+    const ConstArrayView<LabelT> &labels,
+    const ConstArrayView<OtherT> &other_labels,
+    const bool has_ignore_value,
+    const OtherT ignore_value,
+    const std::size_t number_of_threads,
+    const ArrayView<OtherT> &out
+) {
+    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())
+        );
+    }
+    detail_projection::require_rag_shape_matches_labels(rag, labels.shape);
+    if (other_labels.shape != labels.shape) {
+        throw std::invalid_argument("other_labels shape must match labels shape");
+    }
+    const auto number_of_nodes = static_cast<std::size_t>(rag.number_of_nodes());
+    if (out.shape != std::vector<std::ptrdiff_t>{static_cast<std::ptrdiff_t>(number_of_nodes)}) {
+        throw std::invalid_argument(
+            "out shape must be (number_of_nodes,)"
+        );
+    }
+    if (detail::max_label(labels) >= static_cast<std::uint64_t>(number_of_nodes)) {
+        throw std::invalid_argument("labels contain a node id outside the rag");
+    }
+
+    const auto n_pixels = detail_label_accumulation::number_of_pixels(labels.shape);
+    const auto n_threads = detail::normalize_thread_count(number_of_threads, n_pixels);
+
+    std::vector<std::vector<std::unordered_map<OtherT, std::uint64_t>>> per_thread(
+        n_threads,
+        std::vector<std::unordered_map<OtherT, std::uint64_t>>(number_of_nodes)
+    );
+
+    bioimage_cpp::detail::parallel_for_chunks(
+        n_threads,
+        n_pixels,
+        [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) {
+            detail_label_accumulation::scan_chunk(
+                labels.data,
+                other_labels.data,
+                begin,
+                end,
+                has_ignore_value,
+                ignore_value,
+                per_thread[thread_id]
+            );
+        }
+    );
+
+    // Merge per-thread histograms and pick majority per node in one pass over
+    // nodes. This is embarrassingly parallel across nodes.
+    const auto node_threads = detail::normalize_thread_count(number_of_threads, number_of_nodes);
+    bioimage_cpp::detail::parallel_for_chunks(
+        node_threads,
+        number_of_nodes,
+        [&](const std::size_t, const std::size_t node_begin, const std::size_t node_end) {
+            for (std::size_t node = node_begin; node < node_end; ++node) {
+                std::unordered_map<OtherT, std::uint64_t> merged;
+                for (auto &thread_histograms : per_thread) {
+                    auto &node_hist = thread_histograms[node];
+                    for (const auto &entry : node_hist) {
+                        merged[entry.first] += entry.second;
+                    }
+                    node_hist.clear();
+                }
+                OtherT best_label = OtherT{0};
+                std::uint64_t best_count = 0;
+                bool has_best = false;
+                for (const auto &entry : merged) {
+                    if (!has_best ||
+                        entry.second > best_count ||
+                        (entry.second == best_count && entry.first < best_label)) {
+                        best_label = entry.first;
+                        best_count = entry.second;
+                        has_best = true;
+                    }
+                }
+                out.data[node] = has_best ? best_label : OtherT{0};
+            }
+        }
+    );
+}
+
+} // namespace bioimage_cpp::graph

src/bindings/graph.cxx

@@ -7,6 +7,7 @@
 #include "bioimage_cpp/graph/feature_accumulation.hxx"
 #include "bioimage_cpp/graph/grid_features.hxx"
 #include "bioimage_cpp/graph/grid_graph.hxx"
+#include "bioimage_cpp/graph/label_accumulation.hxx"
 #include "bioimage_cpp/graph/lifted_from_affinities.hxx"
 #include "bioimage_cpp/graph/lifted_multicut.hxx"
 #include "bioimage_cpp/graph/lifted_multicut/fusion_move.hxx"
@@ -113,6 +114,20 @@ DoubleArray make_double_array(const std::vector<std::size_t> &shape) {
     return DoubleArray(data, shape.size(), shape.data(), owner);
 }
 
+template <class T>
+using TypedArray = nb::ndarray<nb::numpy, T, nb::c_contig>;
+
+template <class T>
+TypedArray<T> make_typed_array(const std::vector<std::size_t> &shape) {
+    std::size_t size = 1;
+    for (const auto axis_size : shape) {
+        size *= axis_size;
+    }
+    auto *data = new T[size]();
+    nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast<T *>(p); });
+    return TypedArray<T>(data, shape.size(), shape.data(), owner);
+}
+
 template <class T>
 FloatingArray<T> make_floating_array(const std::vector<std::size_t> &shape) {
     std::size_t size = 1;
@@ -1245,6 +1260,55 @@ UInt64Array project_node_labels_to_pixels_t(
     return result;
 }
 
+template <class LabelT, class OtherT>
+TypedArray<OtherT> accumulate_labels_t(
+    const Rag &rag,
+    LabelArray<LabelT> labels,
+    LabelArray<OtherT> other_labels,
+    const bool has_ignore_value,
+    const OtherT ignore_value,
+    const std::size_t number_of_threads
+) {
+    if (labels.ndim() != other_labels.ndim()) {
+        throw std::invalid_argument("other_labels shape must match labels shape");
+    }
+    for (std::size_t axis = 0; axis < labels.ndim(); ++axis) {
+        if (labels.shape(axis) != other_labels.shape(axis)) {
+            throw std::invalid_argument("other_labels shape must match labels shape");
+        }
+    }
+
+    auto result = make_typed_array<OtherT>({static_cast<std::size_t>(rag.number_of_nodes())});
+
+    ConstArrayView<LabelT> labels_view{
+        labels.data(),
+        ndarray_shape(labels),
+        {},
+    };
+    ConstArrayView<OtherT> other_labels_view{
+        other_labels.data(),
+        ndarray_shape(other_labels),
+        {},
+    };
+    ArrayView<OtherT> out_view{
+        result.data(),
+        {static_cast<std::ptrdiff_t>(rag.number_of_nodes())},
+        {},
+    };
+
+    nb::gil_scoped_release release;
+    graph::accumulate_labels<LabelT, OtherT>(
+        rag,
+        labels_view,
+        other_labels_view,
+        has_ignore_value,
+        ignore_value,
+        number_of_threads,
+        out_view
+    );
+    return result;
+}
+
 } // namespace
 
 void bind_graph(nb::module_ &m) {
@@ -1889,6 +1953,37 @@ void bind_graph(nb::module_ &m) {
         nb::arg("node_labels"),
         nb::arg("number_of_threads")
     );
+
+#define BIC_BIND_ACCUMULATE_LABELS(LSUF, LT, OSUF, OT)              \
+    m.def(                                                          \
+        "_accumulate_labels_" #LSUF "_" #OSUF,                      \
+        &accumulate_labels_t<LT, OT>,                               \
+        nb::arg("rag"),                                             \
+        nb::arg("labels"),                                          \
+        nb::arg("other_labels"),                                    \
+        nb::arg("has_ignore_value"),                                \
+        nb::arg("ignore_value"),                                    \
+        nb::arg("number_of_threads")                                \
+    )
+
+    BIC_BIND_ACCUMULATE_LABELS(uint32, std::uint32_t, uint32, std::uint32_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint32, std::uint32_t, uint64, std::uint64_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint32, std::uint32_t, int32,  std::int32_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint32, std::uint32_t, int64,  std::int64_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint64, std::uint64_t, uint32, std::uint32_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint64, std::uint64_t, uint64, std::uint64_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint64, std::uint64_t, int32,  std::int32_t);
+    BIC_BIND_ACCUMULATE_LABELS(uint64, std::uint64_t, int64,  std::int64_t);
+    BIC_BIND_ACCUMULATE_LABELS(int32,  std::int32_t,  uint32, std::uint32_t);
+    BIC_BIND_ACCUMULATE_LABELS(int32,  std::int32_t,  uint64, std::uint64_t);
+    BIC_BIND_ACCUMULATE_LABELS(int32,  std::int32_t,  int32,  std::int32_t);
+    BIC_BIND_ACCUMULATE_LABELS(int32,  std::int32_t,  int64,  std::int64_t);
+    BIC_BIND_ACCUMULATE_LABELS(int64,  std::int64_t,  uint32, std::uint32_t);
+    BIC_BIND_ACCUMULATE_LABELS(int64,  std::int64_t,  uint64, std::uint64_t);
+    BIC_BIND_ACCUMULATE_LABELS(int64,  std::int64_t,  int32,  std::int32_t);
+    BIC_BIND_ACCUMULATE_LABELS(int64,  std::int64_t,  int64,  std::int64_t);
+
+#undef BIC_BIND_ACCUMULATE_LABELS
 }
 
 } // namespace bioimage_cpp::bindings