Add node projection functionality and integration test (#9)

Author: constantinpape

MIGRATION_GUIDE.md

@@ -221,6 +221,39 @@ Notes:
 - `number_of_threads=0` uses the library default; pass a positive integer for a
   fixed thread count.
 
+## Projecting RAG Node Labels to Pixels
+
+Nifty projects scalar node data back to pixels with
+`projectScalarNodeDataToPixels`.
+
+Nifty:
+
+```python
+import nifty.graph.rag as nrag
+
+rag = nrag.gridRag(labels)
+pixel_labels = nrag.projectScalarNodeDataToPixels(rag, node_labels)
+```
+
+bioimage-cpp:
+
+```python
+rag = bic.graph.region_adjacency_graph(labels)
+pixel_labels = bic.graph.project_node_labels_to_pixels(
+    rag,
+    labels,
+    node_labels,
+)
+```
+
+Notes:
+
+- `labels` must be the over-segmentation used to construct `rag`.
+- `node_labels` must be a 1D array with length `rag.number_of_nodes`.
+- The output has the same shape as `labels` and dtype `uint64`.
+- `number_of_threads=0` uses the library default; pass a positive integer for a
+  fixed thread count.
+
 ## Multicut
 
 Nifty exposes multicut through an objective + factory-style solver hierarchy.

examples/segmentation/multicut_from_affinities.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import napari
+import numpy as np
+from elf.io import open_file
+from elf.segmentation.utils import load_mutex_watershed_problem
+from skimage.feature import peak_local_max
+from skimage.measure import label as label_components
+from skimage.segmentation import find_boundaries, watershed
+
+import bioimage_cpp as bic
+
+
+THIS_DIR = Path(__file__).resolve().parent
+DEFAULT_DATA_PREFIX = THIS_DIR / "isbi-data-"
+
+
+def load_problem(data_prefix: Path, ndim: int, z_slice: int):
+    affinities, offsets = load_mutex_watershed_problem(prefix=str(data_prefix))
+    offsets = [tuple(int(v) for v in offset) for offset in offsets]
+    if ndim == 2:
+        channels_2d = [index for index, offset in enumerate(offsets) if offset[0] == 0]
+        affinities = affinities[channels_2d, z_slice]
+        offsets = [offsets[index][1:] for index in channels_2d]
+    elif ndim != 3:
+        raise ValueError(f"ndim must be 2 or 3, got {ndim}")
+
+    direct_channels = [
+        index for index, offset in enumerate(offsets) if sum(abs(v) for v in offset) == 1
+    ]
+    direct_affinities = np.ascontiguousarray(affinities[direct_channels], dtype=np.float32)
+    direct_offsets = [offsets[index] for index in direct_channels]
+    return direct_affinities, direct_offsets
+
+
+def load_raw(data_prefix: Path, ndim: int, z_slice: int):
+    data_path = data_prefix.with_name(data_prefix.name + "test.h5")
+    with open_file(data_path, "r") as f:
+        raw = f["raw"][z_slice] if ndim == 2 else f["raw"][:]
+    return raw
+
+
+def make_heightmap(affinities: np.ndarray) -> np.ndarray:
+    return np.ascontiguousarray(np.mean(affinities, axis=0), dtype=np.float32)
+
+
+def make_watershed_oversegmentation(
+    heightmap: np.ndarray,
+    *,
+    min_distance: int,
+    grid_spacing: int,
+    max_markers: int,
+) -> np.ndarray:
+    coordinates = peak_local_max(
+        -heightmap,
+        min_distance=min_distance,
+        exclude_border=False,
+        num_peaks=max_markers,
+    )
+    marker_mask = np.zeros(heightmap.shape, dtype=bool)
+    if len(coordinates) > 0:
+        marker_mask[tuple(coordinates.T)] = True
+    markers = label_components(marker_mask).astype(np.int32, copy=False)
+
+    if int(markers.max()) < 2:
+        markers = np.zeros(heightmap.shape, dtype=np.int32)
+        slices = tuple(slice(None, None, grid_spacing) for _ in heightmap.shape)
+        marker_coordinates = np.argwhere(np.ones(heightmap.shape, dtype=bool)[slices])
+        marker_coordinates *= grid_spacing
+        for marker_id, coord in enumerate(marker_coordinates, start=1):
+            markers[tuple(coord)] = marker_id
+
+    return watershed(heightmap, markers=markers).astype(np.uint32, copy=False)
+
+
+def multicut_from_affinities(
+    affinities: np.ndarray,
+    offsets: list[tuple[int, ...]],
+    *,
+    threshold: float,
+    number_of_threads: int,
+    watershed_min_distance: int,
+    watershed_grid_spacing: int,
+    max_markers: int,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    heightmap = make_heightmap(affinities)
+    oversegmentation = make_watershed_oversegmentation(
+        heightmap,
+        min_distance=watershed_min_distance,
+        grid_spacing=watershed_grid_spacing,
+        max_markers=max_markers,
+    )
+
+    rag = bic.graph.region_adjacency_graph(
+        oversegmentation, number_of_threads=number_of_threads
+    )
+    features = bic.graph.affinity_features(
+        rag,
+        oversegmentation,
+        affinities,
+        offsets,
+        number_of_threads=number_of_threads,
+    )
+    edge_costs = threshold - features[:, 0]
+
+    objective = bic.graph.MulticutObjective(rag, edge_costs)
+    node_labels = bic.graph.ChainedMulticutSolvers(
+        [
+            bic.graph.GreedyAdditiveMulticut(),
+            bic.graph.KernighanLinMulticut(number_of_outer_iterations=10),
+        ]
+    ).optimize(objective)
+    segmentation = bic.graph.project_node_labels_to_pixels(
+        rag,
+        oversegmentation,
+        node_labels,
+        number_of_threads=number_of_threads,
+    )
+    return heightmap, oversegmentation, segmentation
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run watershed oversegmentation + RAG multicut on the ISBI affinity example."
+    )
+    parser.add_argument("--ndim", type=int, choices=(2, 3), default=2)
+    parser.add_argument("--z-slice", type=int, default=0)
+    parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX)
+    parser.add_argument("--threshold", type=float, default=0.1)
+    parser.add_argument("--threads", type=int, default=0)
+    parser.add_argument("--watershed-min-distance", type=int, default=5)
+    parser.add_argument("--watershed-grid-spacing", type=int, default=12)
+    parser.add_argument("--max-markers", type=int, default=2048)
+    args = parser.parse_args()
+
+    affinities, offsets = load_problem(args.data_prefix, args.ndim, args.z_slice)
+    raw = load_raw(args.data_prefix, args.ndim, args.z_slice)
+    heightmap, oversegmentation, segmentation = multicut_from_affinities(
+        affinities,
+        offsets,
+        threshold=args.threshold,
+        number_of_threads=args.threads,
+        watershed_min_distance=args.watershed_min_distance,
+        watershed_grid_spacing=args.watershed_grid_spacing,
+        max_markers=args.max_markers,
+    )
+
+    viewer = napari.Viewer()
+    viewer.add_image(raw, name="raw")
+    viewer.add_image(affinities, name="direct affinities")
+    viewer.add_image(heightmap, name="watershed heightmap")
+    viewer.add_labels(oversegmentation, name="watershed oversegmentation")
+    viewer.add_labels(segmentation, name="multicut segmentation")
+    viewer.add_labels(find_boundaries(segmentation), name="multicut boundaries")
+    napari.run()
+
+
+if __name__ == "__main__":
+    main()

include/bioimage_cpp/graph/node_label_projection.hxx

@@ -0,0 +1,86 @@
+#pragma once
+
+#include "bioimage_cpp/array_view.hxx"
+#include "bioimage_cpp/detail/threading.hxx"
+#include "bioimage_cpp/graph/region_adjacency_graph.hxx"
+
+#include <cstddef>
+#include <cstdint>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace bioimage_cpp::graph {
+
+namespace detail_projection {
+
+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; }
+    ));
+}
+
+inline void require_rag_shape_matches_labels(
+    const RegionAdjacencyGraph &rag,
+    const std::vector<std::ptrdiff_t> &labels_shape
+) {
+    const auto &rag_shape = rag.shape();
+    if (rag_shape.size() != labels_shape.size()) {
+        throw std::invalid_argument("rag shape must match labels shape");
+    }
+    for (std::size_t axis = 0; axis < rag_shape.size(); ++axis) {
+        if (rag_shape[axis] != static_cast<std::uint64_t>(labels_shape[axis])) {
+            throw std::invalid_argument("rag shape must match labels shape");
+        }
+    }
+}
+
+} // namespace detail_projection
+
+template <class LabelT>
+void project_node_labels_to_pixels(
+    const RegionAdjacencyGraph &rag,
+    const ConstArrayView<LabelT> &labels,
+    const ConstArrayView<std::uint64_t> &node_labels,
+    const std::size_t number_of_threads,
+    const ArrayView<std::uint64_t> &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 (node_labels.ndim() != 1) {
+        throw std::invalid_argument("node_labels must be a 1D uint64 array");
+    }
+    if (node_labels.shape[0] != static_cast<std::ptrdiff_t>(rag.number_of_nodes())) {
+        throw std::invalid_argument("node_labels length must match rag number_of_nodes");
+    }
+    if (out.shape != labels.shape) {
+        throw std::invalid_argument("out shape must match labels shape");
+    }
+    if (detail::max_label(labels) >= static_cast<std::uint64_t>(node_labels.shape[0])) {
+        throw std::invalid_argument("labels contain a node id outside node_labels");
+    }
+
+    const auto n_pixels = detail_projection::number_of_pixels(labels.shape);
+    const auto n_threads = detail::normalize_thread_count(number_of_threads, n_pixels);
+    bioimage_cpp::detail::parallel_for_chunks(
+        n_threads,
+        n_pixels,
+        [&](const std::size_t, const std::size_t begin, const std::size_t end) {
+            for (std::size_t index = begin; index < end; ++index) {
+                const auto node = static_cast<std::uint64_t>(labels.data[index]);
+                out.data[index] = node_labels.data[static_cast<std::size_t>(node)];
+            }
+        }
+    );
+}
+
+} // namespace bioimage_cpp::graph

src/bindings/graph.cxx

@@ -4,6 +4,7 @@
 #include "bioimage_cpp/graph/connected_components.hxx"
 #include "bioimage_cpp/graph/feature_accumulation.hxx"
 #include "bioimage_cpp/graph/multicut.hxx"
+#include "bioimage_cpp/graph/node_label_projection.hxx"
 #include "bioimage_cpp/graph/region_adjacency_graph.hxx"
 #include "bioimage_cpp/graph/undirected_graph.hxx"
 
@@ -483,6 +484,50 @@ DoubleArray accumulate_affinity_features_t(
     return result;
 }
 
+template <class LabelT>
+UInt64Array project_node_labels_to_pixels_t(
+    const Rag &rag,
+    LabelArray<LabelT> labels,
+    ConstUInt64Array node_labels,
+    const std::size_t number_of_threads
+) {
+    if (node_labels.ndim() != 1) {
+        throw std::invalid_argument("node_labels must be a 1D uint64 array");
+    }
+
+    std::vector<std::size_t> output_shape(labels.ndim());
+    for (std::size_t axis = 0; axis < labels.ndim(); ++axis) {
+        output_shape[axis] = labels.shape(axis);
+    }
+    auto result = make_uint64_array(output_shape);
+
+    ConstArrayView<LabelT> labels_view{
+        labels.data(),
+        ndarray_shape(labels),
+        {},
+    };
+    ConstArrayView<std::uint64_t> node_labels_view{
+        node_labels.data(),
+        {static_cast<std::ptrdiff_t>(node_labels.shape(0))},
+        {},
+    };
+    ArrayView<std::uint64_t> out_view{
+        result.data(),
+        ndarray_shape(labels),
+        {},
+    };
+
+    nb::gil_scoped_release release;
+    graph::project_node_labels_to_pixels<LabelT>(
+        rag,
+        labels_view,
+        node_labels_view,
+        number_of_threads,
+        out_view
+    );
+    return result;
+}
+
 } // namespace
 
 void bind_graph(nb::module_ &m) {
@@ -701,6 +746,38 @@ void bind_graph(nb::module_ &m) {
         nb::arg("compute_complex_features"),
         nb::arg("number_of_threads")
     );
+    m.def(
+        "_project_node_labels_to_pixels_uint32",
+        &project_node_labels_to_pixels_t<std::uint32_t>,
+        nb::arg("rag"),
+        nb::arg("labels"),
+        nb::arg("node_labels"),
+        nb::arg("number_of_threads")
+    );
+    m.def(
+        "_project_node_labels_to_pixels_uint64",
+        &project_node_labels_to_pixels_t<std::uint64_t>,
+        nb::arg("rag"),
+        nb::arg("labels"),
+        nb::arg("node_labels"),
+        nb::arg("number_of_threads")
+    );
+    m.def(
+        "_project_node_labels_to_pixels_int32",
+        &project_node_labels_to_pixels_t<std::int32_t>,
+        nb::arg("rag"),
+        nb::arg("labels"),
+        nb::arg("node_labels"),
+        nb::arg("number_of_threads")
+    );
+    m.def(
+        "_project_node_labels_to_pixels_int64",
+        &project_node_labels_to_pixels_t<std::int64_t>,
+        nb::arg("rag"),
+        nb::arg("labels"),
+        nb::arg("node_labels"),
+        nb::arg("number_of_threads")
+    );
 }
 
 } // namespace bioimage_cpp::bindings