diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 6478bbb..d952472 100644 --- a/MIGRATION_GUIDE.md +++ b/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. diff --git a/examples/segmentation/multicut_from_affinities.py b/examples/segmentation/multicut_from_affinities.py new file mode 100644 index 0000000..50b6a57 --- /dev/null +++ b/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() diff --git a/include/bioimage_cpp/graph/node_label_projection.hxx b/include/bioimage_cpp/graph/node_label_projection.hxx new file mode 100644 index 0000000..afc8439 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +namespace detail_projection { + +inline std::size_t number_of_pixels(const std::vector &shape) { + return static_cast(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 &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(labels_shape[axis])) { + throw std::invalid_argument("rag shape must match labels shape"); + } + } +} + +} // namespace detail_projection + +template +void project_node_labels_to_pixels( + const RegionAdjacencyGraph &rag, + const ConstArrayView &labels, + const ConstArrayView &node_labels, + const std::size_t number_of_threads, + const ArrayView &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(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(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(labels.data[index]); + out.data[index] = node_labels.data[static_cast(node)]; + } + } + ); +} + +} // namespace bioimage_cpp::graph diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index b4979db..5055e4f 100644 --- a/src/bindings/graph.cxx +++ b/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 +UInt64Array project_node_labels_to_pixels_t( + const Rag &rag, + LabelArray 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 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 labels_view{ + labels.data(), + ndarray_shape(labels), + {}, + }; + ConstArrayView node_labels_view{ + node_labels.data(), + {static_cast(node_labels.shape(0))}, + {}, + }; + ArrayView out_view{ + result.data(), + ndarray_shape(labels), + {}, + }; + + nb::gil_scoped_release release; + graph::project_node_labels_to_pixels( + 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, + 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, + 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, + 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, + nb::arg("rag"), + nb::arg("labels"), + nb::arg("node_labels"), + nb::arg("number_of_threads") + ); } } // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 98d01b9..a2f05fc 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -36,6 +36,13 @@ np.dtype("int64"): _core._accumulate_affinity_features_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, + np.dtype("int32"): _core._project_node_labels_to_pixels_int32, + np.dtype("int64"): _core._project_node_labels_to_pixels_int64, +} + SIMPLE_EDGE_FEATURE_NAMES = ("mean", "size") COMPLEX_EDGE_FEATURE_NAMES = ( "mean", @@ -523,6 +530,42 @@ def affinity_features_complex( ) +def project_node_labels_to_pixels( + rag: RegionAdjacencyGraph, + labels: np.ndarray, + node_labels, + *, + number_of_threads: int = 0, +) -> np.ndarray: + """Map RAG node labels back to a pixel-wise segmentation. + + ``labels`` is the over-segmentation used to construct ``rag``. Each pixel + value is interpreted as a RAG node id and replaced by the corresponding + entry in the 1D ``node_labels`` array. The returned segmentation has the + same shape as ``labels`` and dtype ``uint64``. + """ + label_array = _normalize_labels(labels) + if tuple(int(size) for size in rag.shape) != label_array.shape: + raise ValueError( + "rag shape must match labels shape, got " + f"rag shape={tuple(rag.shape)}, labels shape={label_array.shape}" + ) + + node_label_array = np.asarray(node_labels, dtype=np.uint64) + if node_label_array.ndim != 1: + raise ValueError("node_labels must be a 1D array") + if node_label_array.shape[0] != rag.number_of_nodes: + raise ValueError("node_labels length must match rag number_of_nodes") + + run = _PROJECT_NODE_LABELS_TO_PIXELS_BY_DTYPE[label_array.dtype] + return run( + rag, + label_array, + np.ascontiguousarray(node_label_array), + _normalize_number_of_threads(number_of_threads), + ) + + def _accumulate_edge_map_features( rag: RegionAdjacencyGraph, labels: np.ndarray, @@ -632,6 +675,7 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "external_multicut_problem_path", "load_external_multicut_problem", "load_external_multicut_problem_data", + "project_node_labels_to_pixels", "region_adjacency_graph", "undirected_graph", ] diff --git a/tests/graph/test_node_label_projection.py b/tests/graph/test_node_label_projection.py new file mode 100644 index 0000000..12148d9 --- /dev/null +++ b/tests/graph/test_node_label_projection.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_project_node_labels_to_pixels_matches_nifty(): + nrag = pytest.importorskip("nifty.graph.rag") + + labels = np.array([[1, 1, 2], [3, 2, 2]], dtype=np.uint64) + node_labels = np.array([10, 11, 12, 13], dtype=np.uint64) + rag = bic.graph.region_adjacency_graph(labels) + + actual = bic.graph.project_node_labels_to_pixels(rag, labels, node_labels) + expected = nrag.projectScalarNodeDataToPixels(nrag.gridRag(labels), node_labels) + + np.testing.assert_array_equal(actual, expected) + + +def test_project_node_labels_to_pixels_parallel_matches_single_thread(): + labels = np.arange(6 * 7, dtype=np.uint32).reshape(6, 7) + rag = bic.graph.region_adjacency_graph(labels, number_of_threads=2) + node_labels = (np.arange(rag.number_of_nodes, dtype=np.uint64) // 3) + 17 + + single_threaded = bic.graph.project_node_labels_to_pixels( + rag, labels, node_labels, number_of_threads=1 + ) + parallel = bic.graph.project_node_labels_to_pixels( + rag, labels, node_labels, number_of_threads=4 + ) + + np.testing.assert_array_equal(parallel, single_threaded) + + +def test_project_node_labels_to_pixels_validation(): + labels = np.array([[0, 1]], dtype=np.uint64) + rag = bic.graph.region_adjacency_graph(labels) + + with pytest.raises(ValueError, match="node_labels must be a 1D array"): + bic.graph.project_node_labels_to_pixels(rag, labels, [[0, 1]]) + + with pytest.raises(ValueError, match="node_labels length"): + bic.graph.project_node_labels_to_pixels(rag, labels, [0]) + + other_labels = np.array([[0, 1], [0, 1]], dtype=np.uint64) + with pytest.raises(ValueError, match="rag shape"): + bic.graph.project_node_labels_to_pixels(rag, other_labels, [0, 1]) diff --git a/tests/graph/test_rag_multicut_integration.py b/tests/graph/test_rag_multicut_integration.py new file mode 100644 index 0000000..73f2389 --- /dev/null +++ b/tests/graph/test_rag_multicut_integration.py @@ -0,0 +1,64 @@ +import numpy as np + +import bioimage_cpp as bic + + +def _dense_by_first_occurrence(labels): + labels = np.asarray(labels) + dense = np.empty(labels.shape, dtype=np.uint64) + mapping = {} + next_label = 0 + for index in np.ndindex(labels.shape): + label = int(labels[index]) + if label not in mapping: + mapping[label] = next_label + next_label += 1 + dense[index] = mapping[label] + return dense + + +def test_rag_feature_multicut_projection_on_realistic_oversegmentation(): + shape = (48, 64) + block_shape = (4, 4) + true_segmentation = np.zeros(shape, dtype=np.uint64) + true_segmentation[8:40, 8:32] = 1 + true_segmentation[12:44, 36:60] = 2 + + oversegmentation = np.empty(shape, dtype=np.uint64) + next_label = 0 + for y in range(0, shape[0], block_shape[0]): + for x in range(0, shape[1], block_shape[1]): + yy = slice(y, y + block_shape[0]) + xx = slice(x, x + block_shape[1]) + oversegmentation[yy, xx] = next_label + next_label += 1 + + boundary_map = np.full(shape, 0.08, dtype=np.float64) + boundary_pixels = np.zeros(shape, dtype=bool) + boundary_pixels[:, :-1] |= true_segmentation[:, :-1] != true_segmentation[:, 1:] + boundary_pixels[:, 1:] |= true_segmentation[:, :-1] != true_segmentation[:, 1:] + boundary_pixels[:-1, :] |= true_segmentation[:-1, :] != true_segmentation[1:, :] + boundary_pixels[1:, :] |= true_segmentation[:-1, :] != true_segmentation[1:, :] + boundary_map[boundary_pixels] = 0.92 + + rag = bic.graph.region_adjacency_graph(oversegmentation, number_of_threads=3) + features = bic.graph.edge_map_features( + rag, oversegmentation, boundary_map, number_of_threads=3 + ) + edge_costs = 0.6 - features[:, 0] + + objective = bic.graph.MulticutObjective(rag, edge_costs) + node_labels = bic.graph.ChainedMulticutSolvers( + [ + bic.graph.GreedyAdditiveMulticut(), + bic.graph.KernighanLinMulticut(number_of_outer_iterations=5), + ] + ).optimize(objective) + projected = bic.graph.project_node_labels_to_pixels( + rag, oversegmentation, node_labels, number_of_threads=3 + ) + + expected = _dense_by_first_occurrence(true_segmentation) + actual = _dense_by_first_occurrence(projected) + np.testing.assert_array_equal(actual, expected) + assert np.unique(projected).size == 3