Implement edge-weighted graph watershed

Author: constantinpape

MIGRATION_GUIDE.md

@@ -254,6 +254,50 @@ Notes:
 - `number_of_threads=0` uses the library default; pass a positive integer for a
   fixed thread count.
 
+## Edge-Weighted Watershed
+
+Nifty:
+
+```python
+import nifty.graph as ng
+
+labels = ng.edgeWeightedWatershedsSegmentation(graph, edge_weights, seeds)
+```
+
+bioimage-cpp:
+
+```python
+import bioimage_cpp as bic
+
+labels = bic.graph.edge_weighted_watershed(graph, edge_weights, seeds)
+```
+
+Notes:
+
+- Only the Kruskal variant of nifty's algorithm is provided. Edges are visited
+  in ascending weight order; two distinct components merge iff at least one is
+  unlabeled (seed label `0`), so seed boundaries are preserved.
+- `graph` may be an `UndirectedGraph` or a `RegionAdjacencyGraph`.
+- `edge_weights` must be 1D with length `graph.number_of_edges`. Supported
+  dtypes are `float32` and `float64`; other floating dtypes are promoted to
+  `float32` (matching nifty, whose Python binding is `float32`-only). Non-float
+  dtypes raise `TypeError`.
+- `seeds` must be 1D with length `graph.number_of_nodes`. Supported dtypes are
+  `uint32`, `uint64`, `int32`, `int64`. The value `0` marks unlabeled nodes;
+  non-zero ids are propagated along low-weight paths. Signed seed arrays must
+  not contain negative values.
+- The output is 1D with length `graph.number_of_nodes` and the same dtype as
+  `seeds`. Seed label values are preserved (no dense relabeling). Nodes that
+  no seed can reach remain `0`.
+
+Intentional differences vs. nifty:
+
+- No priority-queue variant — only the simpler sort + union-find Kruskal flow.
+  For the same input it matches nifty's default behavior (which also dispatches
+  to the Kruskal implementation).
+- No carving / background-bias variant. Build a carving prior into the edge
+  weights before calling the function if needed.
+
 ## Multicut
 
 Nifty exposes multicut through an objective + factory-style solver hierarchy.

include/bioimage_cpp/graph/edge_weighted_watershed.hxx

@@ -0,0 +1,105 @@
+#pragma once
+
+#include "bioimage_cpp/detail/union_find.hxx"
+#include "bioimage_cpp/graph/undirected_graph.hxx"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+namespace bioimage_cpp::graph {
+
+// Kruskal-style edge-weighted seeded watershed.
+//
+// Edges are visited in ascending weight order. Two distinct components are
+// merged iff at least one of them is unlabeled (seed label 0); the non-zero
+// seed label then propagates. Two distinct already-labeled components are
+// never merged, so seed boundaries are preserved. Nodes that no seed can
+// reach retain label 0.
+//
+// `WeightT` is any type orderable with `<` (typically float32 or float64).
+// `SeedT` is any integer label type. Signed types must not contain negative
+// values; this is checked at the boundary.
+template <class WeightT, class SeedT>
+inline std::vector<SeedT> edge_weighted_watershed(
+    const UndirectedGraph &graph,
+    const std::vector<WeightT> &edge_weights,
+    const std::vector<SeedT> &seeds
+) {
+    const auto number_of_nodes = graph.number_of_nodes();
+    const auto number_of_edges = graph.number_of_edges();
+
+    if (edge_weights.size() != static_cast<std::size_t>(number_of_edges)) {
+        throw std::invalid_argument(
+            "edge_weights length must equal number_of_edges, got " +
+            std::to_string(edge_weights.size()) + " for number_of_edges=" +
+            std::to_string(number_of_edges)
+        );
+    }
+    if (seeds.size() != static_cast<std::size_t>(number_of_nodes)) {
+        throw std::invalid_argument(
+            "seeds length must equal number_of_nodes, got " +
+            std::to_string(seeds.size()) + " for number_of_nodes=" +
+            std::to_string(number_of_nodes)
+        );
+    }
+    if constexpr (std::is_signed_v<SeedT>) {
+        for (const auto value : seeds) {
+            if (value < 0) {
+                throw std::invalid_argument("seeds must not contain negative values");
+            }
+        }
+    }
+
+    std::vector<SeedT> labels(seeds);
+
+    if (number_of_edges == 0) {
+        return labels;
+    }
+
+    std::vector<std::uint64_t> order(static_cast<std::size_t>(number_of_edges));
+    for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) {
+        order[static_cast<std::size_t>(edge)] = edge;
+    }
+    std::stable_sort(
+        order.begin(),
+        order.end(),
+        [&edge_weights](const std::uint64_t a, const std::uint64_t b) {
+            return edge_weights[static_cast<std::size_t>(a)] <
+                   edge_weights[static_cast<std::size_t>(b)];
+        }
+    );
+
+    detail::UnionFind sets(static_cast<std::size_t>(number_of_nodes));
+
+    constexpr SeedT zero{0};
+    for (const auto edge : order) {
+        const auto uv = graph.uv(edge);
+        const auto ru = sets.find(uv.first);
+        const auto rv = sets.find(uv.second);
+        if (ru == rv) {
+            continue;
+        }
+        const auto lu = labels[static_cast<std::size_t>(ru)];
+        const auto lv = labels[static_cast<std::size_t>(rv)];
+        if (lu != zero && lv != zero) {
+            continue;
+        }
+        const auto new_label = std::max(lu, lv);
+        const auto new_root = sets.unite_roots(ru, rv);
+        labels[static_cast<std::size_t>(new_root)] = new_label;
+    }
+
+    for (std::uint64_t node = 0; node < number_of_nodes; ++node) {
+        const auto root = sets.find(node);
+        labels[static_cast<std::size_t>(node)] = labels[static_cast<std::size_t>(root)];
+    }
+
+    return labels;
+}
+
+} // namespace bioimage_cpp::graph

src/bindings/graph.cxx

@@ -2,6 +2,7 @@
 
 #include "bioimage_cpp/array_view.hxx"
 #include "bioimage_cpp/graph/connected_components.hxx"
+#include "bioimage_cpp/graph/edge_weighted_watershed.hxx"
 #include "bioimage_cpp/graph/feature_accumulation.hxx"
 #include "bioimage_cpp/graph/multicut.hxx"
 #include "bioimage_cpp/graph/node_label_projection.hxx"
@@ -12,6 +13,7 @@
 #include <nanobind/stl/pair.h>
 #include <nanobind/stl/vector.h>
 
+#include <array>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -285,6 +287,55 @@ UInt64Array graph_connected_components_masked(const Graph &graph, ConstUInt8Arra
     return vector_to_uint64_array(labels);
 }
 
+template <class T>
+using ConstArray1D = nb::ndarray<nb::numpy, const T, nb::c_contig>;
+
+template <class T>
+std::vector<T> array_1d_to_vector(
+    ConstArray1D<T> array,
+    const char *argument_name,
+    const std::uint64_t expected_size
+) {
+    if (array.ndim() != 1) {
+        throw std::invalid_argument(std::string(argument_name) + " must be a 1D array");
+    }
+    if (array.shape(0) != static_cast<std::size_t>(expected_size)) {
+        throw std::invalid_argument(
+            std::string(argument_name) + " length must match expected size"
+        );
+    }
+    const auto *data = array.data();
+    return std::vector<T>(data, data + array.shape(0));
+}
+
+template <class T>
+nb::ndarray<nb::numpy, T, nb::c_contig> vector_to_array_1d(const std::vector<T> &values) {
+    const std::size_t size = values.size();
+    auto *data = new T[size];
+    std::copy(values.begin(), values.end(), data);
+    nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast<T *>(p); });
+    const std::array<std::size_t, 1> shape{size};
+    return nb::ndarray<nb::numpy, T, nb::c_contig>(data, 1, shape.data(), owner);
+}
+
+template <class WeightT, class SeedT>
+nb::ndarray<nb::numpy, SeedT, nb::c_contig> graph_edge_weighted_watershed_t(
+    const Graph &graph,
+    ConstArray1D<WeightT> edge_weights,
+    ConstArray1D<SeedT> seeds
+) {
+    const auto weight_vector =
+        array_1d_to_vector<WeightT>(edge_weights, "edge_weights", graph.number_of_edges());
+    const auto seed_vector =
+        array_1d_to_vector<SeedT>(seeds, "seeds", graph.number_of_nodes());
+    std::vector<SeedT> labels;
+    {
+        nb::gil_scoped_release release;
+        labels = graph::edge_weighted_watershed<WeightT, SeedT>(graph, weight_vector, seed_vector);
+    }
+    return vector_to_array_1d<SeedT>(labels);
+}
+
 double multicut_energy(const Graph &graph, ConstDoubleArray costs, ConstUInt64Array labels) {
     const auto cost_vector = double_array_to_vector(costs, "edge_costs", graph.number_of_edges());
     const auto label_vector = uint64_array_to_vector(labels, "labels", graph.number_of_nodes());
@@ -605,6 +656,25 @@ void bind_graph(nb::module_ &m) {
         nb::arg("graph"),
         nb::arg("edge_mask")
     );
+    const auto register_watershed = [&m]<class WeightT, class SeedT>(
+        const char *name
+    ) {
+        m.def(
+            name,
+            &graph_edge_weighted_watershed_t<WeightT, SeedT>,
+            nb::arg("graph"),
+            nb::arg("edge_weights"),
+            nb::arg("seeds")
+        );
+    };
+    register_watershed.operator()<float, std::uint32_t>("_edge_weighted_watershed_float32_uint32");
+    register_watershed.operator()<float, std::uint64_t>("_edge_weighted_watershed_float32_uint64");
+    register_watershed.operator()<float, std::int32_t>("_edge_weighted_watershed_float32_int32");
+    register_watershed.operator()<float, std::int64_t>("_edge_weighted_watershed_float32_int64");
+    register_watershed.operator()<double, std::uint32_t>("_edge_weighted_watershed_float64_uint32");
+    register_watershed.operator()<double, std::uint64_t>("_edge_weighted_watershed_float64_uint64");
+    register_watershed.operator()<double, std::int32_t>("_edge_weighted_watershed_float64_int32");
+    register_watershed.operator()<double, std::int64_t>("_edge_weighted_watershed_float64_int64");
     m.def(
         "_multicut_energy",
         &multicut_energy,

src/bioimage_cpp/graph/__init__.py

@@ -36,6 +36,17 @@
     np.dtype("int64"): _core._accumulate_affinity_features_int64,
 }
 
+_EDGE_WEIGHTED_WATERSHED_BY_DTYPE = {
+    (np.dtype("float32"), np.dtype("uint32")): _core._edge_weighted_watershed_float32_uint32,
+    (np.dtype("float32"), np.dtype("uint64")): _core._edge_weighted_watershed_float32_uint64,
+    (np.dtype("float32"), np.dtype("int32")): _core._edge_weighted_watershed_float32_int32,
+    (np.dtype("float32"), np.dtype("int64")): _core._edge_weighted_watershed_float32_int64,
+    (np.dtype("float64"), np.dtype("uint32")): _core._edge_weighted_watershed_float64_uint32,
+    (np.dtype("float64"), np.dtype("uint64")): _core._edge_weighted_watershed_float64_uint64,
+    (np.dtype("float64"), np.dtype("int32")): _core._edge_weighted_watershed_float64_int32,
+    (np.dtype("float64"), np.dtype("int64")): _core._edge_weighted_watershed_float64_int64,
+}
+
 _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,
@@ -159,6 +170,17 @@ def _as_node_labels(labels, graph: UndirectedGraph | RegionAdjacencyGraph) -> np
     return np.ascontiguousarray(array)
 
 
+def _as_1d_array(values, dtype, name: str, expected_size: int) -> np.ndarray:
+    array = np.asarray(values, dtype=dtype)
+    if array.ndim != 1:
+        raise ValueError(f"{name} must be a 1D array")
+    if array.shape[0] != expected_size:
+        raise ValueError(
+            f"{name} length must be {expected_size}, got {array.shape[0]}"
+        )
+    return np.ascontiguousarray(array)
+
+
 def _dense_labels(labels) -> np.ndarray:
     labels = np.asarray(labels, dtype=np.uint64)
     _, dense = np.unique(labels, return_inverse=True)
@@ -207,6 +229,73 @@ def connected_components(
     )
 
 
+def edge_weighted_watershed(
+    graph: UndirectedGraph | RegionAdjacencyGraph,
+    edge_weights,
+    seeds,
+) -> np.ndarray:
+    """Kruskal-style edge-weighted seeded watershed on an undirected graph.
+
+    Edges are visited in ascending weight order. Two distinct components are
+    merged iff at least one of them is unlabeled (seed label ``0``); the
+    non-zero seed label then propagates. Two distinct already-labeled
+    components are never merged, so seed boundaries are preserved.
+
+    Parameters
+    ----------
+    graph:
+        :class:`UndirectedGraph` or :class:`RegionAdjacencyGraph`.
+    edge_weights:
+        1D array of length ``graph.number_of_edges``. Supported dtypes are
+        ``float32`` and ``float64``. Other floating dtypes are cast to
+        ``float32`` (matches nifty); other dtypes raise ``TypeError``.
+    seeds:
+        1D array of length ``graph.number_of_nodes``. Supported dtypes are
+        ``uint32``, ``uint64``, ``int32``, ``int64``. ``0`` marks unlabeled
+        nodes; positive ids are seed labels and propagate along low-weight
+        paths. Signed seed arrays must not contain negative values.
+
+    Returns
+    -------
+    np.ndarray
+        1D array of length ``graph.number_of_nodes`` with the same dtype as
+        ``seeds``. Nodes reachable from a seed receive that seed's label;
+        unreachable nodes remain ``0``. Seed label values are preserved (no
+        dense relabeling).
+    """
+    weight_array = np.asarray(edge_weights)
+    if weight_array.dtype not in (np.dtype("float32"), np.dtype("float64")):
+        if np.issubdtype(weight_array.dtype, np.floating):
+            weight_array = weight_array.astype(np.float32, copy=False)
+        else:
+            raise TypeError(
+                "edge_weights must have dtype float32 or float64, got "
+                f"dtype={weight_array.dtype}"
+            )
+
+    seed_array = np.asarray(seeds)
+    if seed_array.dtype not in (
+        np.dtype("uint32"),
+        np.dtype("uint64"),
+        np.dtype("int32"),
+        np.dtype("int64"),
+    ):
+        raise TypeError(
+            "seeds must have dtype uint32, uint64, int32, or int64, got "
+            f"dtype={seed_array.dtype}"
+        )
+
+    weight_array = _as_1d_array(
+        weight_array, weight_array.dtype, "edge_weights", int(graph.number_of_edges)
+    )
+    seed_array = _as_1d_array(
+        seed_array, seed_array.dtype, "seeds", int(graph.number_of_nodes)
+    )
+
+    run = _EDGE_WEIGHTED_WATERSHED_BY_DTYPE[(weight_array.dtype, seed_array.dtype)]
+    return run(graph, weight_array, seed_array)
+
+
 class MulticutObjective:
     """Multicut objective for an undirected graph and edge costs."""
 
@@ -672,6 +761,7 @@ def _normalize_number_of_threads(number_of_threads: int) -> int:
     "connected_components",
     "edge_map_features",
     "edge_map_features_complex",
+    "edge_weighted_watershed",
     "external_multicut_problem_path",
     "load_external_multicut_problem",
     "load_external_multicut_problem_data",