diff --git a/CMakeLists.txt b/CMakeLists.txt index a4865b3..b68202e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -20,6 +20,7 @@ nanobind_add_module(_core src/bindings/segmentation.cxx src/bindings/utils.cxx src/cpp/segmentation/mutex_watershed.cxx + src/cpp/segmentation/semantic_mutex_watershed.cxx src/cpp/take_dict.cxx ) diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 1e05436..e41ddda 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -958,6 +958,142 @@ Notes: partitions. See `development/graph/check_mutex_clustering.py` for a comparison harness. +### Semantic mutex watershed + +`bioimage-cpp` mirrors affogato's two semantic-mutex-watershed entry points +— `compute_semantic_mws_segmentation` for affinity volumes and +`compute_semantic_mws_clustering` for an arbitrary graph. Both extend the +regular mutex watershed with a third edge type: per-pixel (or per-node) +"semantic edges" that tag each cluster with a class id. Two clusters that +have been tagged with different class ids cannot subsequently merge. + +#### Grid-based semantic mutex watershed (affinity volumes) + +Affogato: + +```python +from affogato.segmentation import compute_semantic_mws_segmentation + +labels, semantic_labels = compute_semantic_mws_segmentation( + weights, + offsets, + number_of_attractive_channels=3, + strides=[1, 1, 1], +) +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic + +labels, semantic_labels = bic.segmentation.semantic_mutex_watershed( + weights, + offsets, + number_of_attractive_channels=3, + strides=[1, 1, 1], +) +``` + +Channel layout (identical to affogato): + +- Channels `[0, number_of_attractive_channels)` are attractive grid edges. +- Channels `[number_of_attractive_channels, len(offsets))` are mutex grid + edges. +- Channels `[len(offsets), affinities.shape[0])` are per-semantic-class + affinities; channel `len(offsets) + c` scores how strongly each pixel + belongs to class `c`. + +Important migration notes: + +- Inputs must represent 2D or 3D grids with shapes `(channels, y, x)` or + `(channels, z, y, x)` and `channels > len(offsets)` (use + `bic.segmentation.mutex_watershed` if there are no semantic channels). +- Supported affinity dtypes are `float32` and `float64`. +- Returned `labels` are `uint64`, consecutive, and 1-based for foreground + pixels (matching the regular grid-based mutex watershed). `semantic_labels` + is `int64` with `-1` reserved for clusters that received no class + assignment. +- A boolean `mask` may be passed. Edges touching `False` pixels are + ignored. Masked pixels are set to label `0` in `labels` and to the + `mask_label` parameter (default `0`) in `semantic_labels`. +- `strides` and `randomized_strides` follow the same convention as the + regular grid mutex watershed (mutex channels only; attractive channels + are always kept). + +#### Graph-based semantic mutex watershed + +Affogato: + +```python +from affogato.segmentation import compute_semantic_mws_clustering + +labels, semantic_labels = compute_semantic_mws_clustering( + number_of_nodes, + uvs.astype(np.uint64), + mutex_uvs.astype(np.uint64), + semantic_node_classes.astype(np.uint64), + weights.astype(np.float32), + mutex_weights.astype(np.float32), + semantic_weights.astype(np.float32), +) +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic + +graph = bic.graph.UndirectedGraph.from_edges(number_of_nodes, uvs) +labels, semantic_labels = bic.graph.semantic_mutex_watershed_clustering( + graph, + weights, + mutex_uvs, + mutex_weights, + semantic_node_classes, + semantic_weights, +) +``` + +Input format mirrors `mutex_watershed_clustering` plus two extra arrays: + +- `semantic_node_classes` is an `(n_semantic, 2)` `uint64` table. Column 0 + is a node id and column 1 is the semantic class id (a non-negative + integer interpreted as `int64` internally). +- `semantic_weights` is a 1D `float32`/`float64` array of length + `n_semantic` giving one weight per `(node, class)` candidate. + +Notes: + +- All three weight arrays (`weights`, `mutex_weights`, `semantic_weights`) + must have the same floating dtype, or all three are promoted to + `float64`. +- Output `labels` are dense `uint64` ids in `0 .. number_of_clusters - 1` + (first-occurrence order, matching the regular graph mutex watershed — + *not* the 1-based foreground labels produced by the grid variant). + `semantic_labels` is `int64` with `-1` for unassigned clusters. +- Accepts both `UndirectedGraph` and `RegionAdjacencyGraph`. + +#### Divergence from affogato + +The bioimage-cpp port fixes a missing `merge_semantic_labels` call on +attractive merges in affogato's graph kernel (`compute_semantic_mws_clustering`): +without that call, a node that has been tagged with a class can have its +tag dropped when it later becomes the non-root of a merge. Affogato's +array kernel additionally invokes `boost::disjoint_sets::link(u, v)` on +raw node ids rather than their roots, which corrupts the union-find tree +on multi-class inputs and over-fragments the result. The unit-test +problems shipped with affogato do not exercise these paths heavily, so +this only shows up on realistic multi-class data. + +For most inputs the two implementations agree. On dense multi-class +inputs they may not; the development scripts under +`development/segmentation/check_semantic_mutex_watershed_{2d,3d}.py` and +`development/graph/check_semantic_mutex_clustering.py` print VI/ARI +partition metrics and a semantic-label match fraction so the deviation is +measurable. The bioimage-cpp partitions match an independent Python +reference implementation of the algorithm. + ## Dictionary-Based Relabeling If you used a small helper to apply a dictionary to an integer label array, use diff --git a/development/graph/check_mutex_clustering.py b/development/graph/mutex_watershed/check_mutex_clustering.py similarity index 100% rename from development/graph/check_mutex_clustering.py rename to development/graph/mutex_watershed/check_mutex_clustering.py diff --git a/development/graph/mutex_watershed/check_semantic_mutex_clustering.py b/development/graph/mutex_watershed/check_semantic_mutex_clustering.py new file mode 100644 index 0000000..ca733fe --- /dev/null +++ b/development/graph/mutex_watershed/check_semantic_mutex_clustering.py @@ -0,0 +1,420 @@ +"""Compare bioimage-cpp ``semantic_mutex_watershed_clustering`` against +affogato's ``compute_semantic_mws_clustering`` reference. + +Pipeline (single mode, no subcommands): + +1. Load the registered ``semantic_labels`` volume (instance + semantic ground + truth + raw, cropped to the labelled slab inside ``_data.py``). +2. Derive an MWS-style affinity stack from the instance + semantic labels via + the helper from ``../segmentation/_semantic_mws_equivalence.py`` + (``bic.affinities.compute_affinities`` + Gaussian noise + tiebreak + perturbation). +3. Oversegment with a simple grid-marker watershed on + ``1 - mean(attractive affinities)``. +4. Build a RAG from the oversegmentation and reduce weights to per-edge + attractive costs (``bic.graph.affinity_features`` mean), to long-range + mutex edges (``bic.graph.lifted_edges_from_affinities`` + + ``bic.graph.lifted_affinity_features``), and to per-(node, class) + semantic costs via ``scipy.ndimage.mean``. +5. Run bioimage-cpp + affogato side-by-side, time both, and report partition + + semantic agreement. + +The two implementations will not in general agree exactly — see the script's +output for the reason. Not part of the pytest suite (per AGENTS.md). +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + +# Reuse the array-version helper for deriving weights, so both flavours of the +# comparison build their inputs the same way. +_SEGMENTATION_DEV = Path(__file__).resolve().parent.parent / "segmentation" +if str(_SEGMENTATION_DEV) not in sys.path: + sys.path.insert(0, str(_SEGMENTATION_DEV)) +from _semantic_mws_equivalence import build_weight_volume # noqa: E402 + + +_DEFAULT_2D_OFFSETS: list[tuple[int, ...]] = [ + (-1, 0), + (0, -1), + (-3, 0), + (0, -3), + (-9, 0), + (0, -9), +] +_DEFAULT_2D_ATTRACTIVE = 2 + +_DEFAULT_3D_OFFSETS: list[tuple[int, ...]] = [ + (-1, 0, 0), + (0, -1, 0), + (0, 0, -1), + (0, -3, 0), + (0, 0, -3), + (0, -9, 0), + (0, 0, -9), +] +_DEFAULT_3D_ATTRACTIVE = 3 + + +def _load_validation_metrics(): + try: + from elf.validation import rand_index, variation_of_information + + return "elf.validation", rand_index, variation_of_information + except ImportError: + from elf.evaluation import rand_index, variation_of_information + + return "elf.evaluation", rand_index, variation_of_information + + +def load_volume(ndim: int, z_start: int, shape: tuple[int, ...]): + from bioimage_cpp._data import load_semantic_labels + + instance, semantic = load_semantic_labels() + if ndim == 2: + y, x = shape + z = z_start + inst = np.ascontiguousarray(instance[z, :y, :x]) + sem = np.ascontiguousarray(semantic[z, :y, :x]) + offsets = list(_DEFAULT_2D_OFFSETS) + attractive = _DEFAULT_2D_ATTRACTIVE + elif ndim == 3: + z, y, x = shape + inst = np.ascontiguousarray(instance[z_start : z_start + z, :y, :x]) + sem = np.ascontiguousarray(semantic[z_start : z_start + z, :y, :x]) + offsets = list(_DEFAULT_3D_OFFSETS) + attractive = _DEFAULT_3D_ATTRACTIVE + else: + raise ValueError(f"ndim must be 2 or 3, got {ndim}") + return inst, sem, offsets, attractive + + +def oversegment( + weights: np.ndarray, + number_of_attractive_channels: int, + seed_spacing: int, +) -> np.ndarray: + """Stride-marker watershed on the inverted mean attractive affinity.""" + from skimage.segmentation import watershed + + attr = weights[:number_of_attractive_channels] + heightmap = np.ascontiguousarray( + 1.0 - attr.mean(axis=0), dtype=np.float32 + ) + markers = np.zeros(heightmap.shape, dtype=np.int32) + slices = tuple(slice(None, None, seed_spacing) for _ in heightmap.shape) + coord_grid = np.argwhere(np.ones(heightmap.shape, dtype=bool)[slices]) * seed_spacing + for marker_id, coord in enumerate(coord_grid, start=1): + markers[tuple(coord)] = marker_id + return watershed(heightmap, markers=markers).astype(np.uint64, copy=False) + + +def build_costs( + over_seg: np.ndarray, + weights: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, + number_of_offsets: int, + n_classes: int, + *, + rag_threads: int, +): + """Reduce the per-pixel weight stack to RAG-level costs.""" + import bioimage_cpp as bic + from scipy.ndimage import mean as ndi_mean + + over_seg = np.ascontiguousarray(over_seg, dtype=np.uint64) + rag = bic.graph.region_adjacency_graph(over_seg, number_of_threads=rag_threads) + + # Attractive edges: mean affinity along each RAG edge. + attr_w = np.ascontiguousarray( + weights[:number_of_attractive_channels].astype(np.float32, copy=False) + ) + attr_offsets = offsets[:number_of_attractive_channels] + attr_features = bic.graph.affinity_features( + rag, over_seg, attr_w, attr_offsets, number_of_threads=rag_threads + ) + edge_costs = np.ascontiguousarray(attr_features[:, 0], dtype=np.float32) + + # Mutex (long-range) edges discovered from the labelling at long-range + # offsets; cost = mean of inverted-mutex weights over those pairs. + mutex_offsets = offsets[number_of_attractive_channels:number_of_offsets] + mutex_w = np.ascontiguousarray( + weights[number_of_attractive_channels:number_of_offsets].astype( + np.float32, copy=False + ) + ) + mutex_uvs = bic.graph.lifted_edges_from_affinities( + rag, over_seg, mutex_offsets, number_of_threads=rag_threads + ) + mutex_features = bic.graph.lifted_affinity_features( + over_seg, mutex_w, mutex_offsets, mutex_uvs, number_of_threads=rag_threads + ) + # mutex_features columns are (mean, size); use the mean as the edge cost. + mutex_costs = np.ascontiguousarray(mutex_features[:, 0], dtype=np.float32) + + # Per-(segment, class) semantic costs: emit one entry per class per + # segment that actually contains pixels (skimage's watershed labels + # start at 1, so segment 0 is empty in the RAG even though + # ``rag.number_of_nodes`` includes it). + present_segments = np.unique(over_seg).astype(np.uint64, copy=False) + semantic_node_classes_list: list[np.ndarray] = [] + semantic_costs_list: list[np.ndarray] = [] + for c in range(n_classes): + per_segment = ndi_mean( + weights[number_of_offsets + c], + labels=over_seg, + index=present_segments.astype(np.int64), + ) + per_segment = np.asarray(per_segment, dtype=np.float32) + nodes_col = present_segments.reshape(-1, 1) + class_col = np.full((present_segments.size, 1), c, dtype=np.uint64) + semantic_node_classes_list.append(np.concatenate([nodes_col, class_col], axis=1)) + semantic_costs_list.append(per_segment) + + semantic_node_classes = np.ascontiguousarray( + np.concatenate(semantic_node_classes_list, axis=0), dtype=np.uint64 + ) + semantic_costs = np.ascontiguousarray( + np.concatenate(semantic_costs_list, axis=0), dtype=np.float32 + ) + return rag, edge_costs, mutex_uvs, mutex_costs, semantic_node_classes, semantic_costs + + +def run_bioimage_cpp( + rag, + edge_costs, + mutex_uvs, + mutex_costs, + semantic_node_classes, + semantic_costs, +) -> tuple[np.ndarray, np.ndarray]: + import bioimage_cpp as bic + + return bic.graph.semantic_mutex_watershed_clustering( + rag, + edge_costs, + mutex_uvs, + mutex_costs, + semantic_node_classes, + semantic_costs, + ) + + +def run_affogato_reference( + n_nodes: int, + uvs: np.ndarray, + mutex_uvs: np.ndarray, + semantic_node_classes: np.ndarray, + edge_costs: np.ndarray, + mutex_costs: np.ndarray, + semantic_costs: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + from affogato.segmentation import compute_semantic_mws_clustering + + return compute_semantic_mws_clustering( + int(n_nodes), + np.ascontiguousarray(uvs, dtype=np.uint64), + np.ascontiguousarray(mutex_uvs, dtype=np.uint64), + np.ascontiguousarray(semantic_node_classes, dtype=np.uint64), + np.ascontiguousarray(edge_costs, dtype=np.float32), + np.ascontiguousarray(mutex_costs, dtype=np.float32), + np.ascontiguousarray(semantic_costs, dtype=np.float32), + ) + + +def _canonical_labels(labels: np.ndarray) -> np.ndarray: + array = np.asarray(labels) + _, first_index, inverse = np.unique( + array, return_index=True, return_inverse=True + ) + order = np.argsort(first_index) + remap = np.empty_like(order) + remap[order] = np.arange(order.size) + return remap[inverse].astype(np.uint64, copy=False) + + +def compare( + bic_result: tuple[np.ndarray, np.ndarray], + aff_result: tuple[np.ndarray, np.ndarray], +) -> dict: + bic_lab, bic_sem = bic_result + aff_lab, aff_sem = aff_result + source, rand_index, variation_of_information = _load_validation_metrics() + + vi_split, vi_merge = variation_of_information(bic_lab, aff_lab) + are, ri = rand_index(bic_lab, aff_lab) + partition_equal = bool( + np.array_equal(_canonical_labels(bic_lab), _canonical_labels(aff_lab)) + ) + semantic_match = float(np.mean(bic_sem == aff_sem)) + return { + "validation_source": source, + "vi_split": float(vi_split), + "vi_merge": float(vi_merge), + "adapted_rand_error": float(are), + "rand_index": float(ri), + "partition_equal": partition_equal, + "semantic_match_fraction": semantic_match, + "n_clusters_bic": int(np.unique(bic_lab).size), + "n_clusters_reference": int(np.unique(aff_lab).size), + } + + +def time_interleaved( + bic_run: Callable[[], tuple[np.ndarray, np.ndarray]], + aff_run: Callable[[], tuple[np.ndarray, np.ndarray]], + repeats: int, +): + bic_result = bic_run() + aff_result = aff_run() + bic_timings: list[float] = [] + aff_timings: list[float] = [] + for repeat in range(repeats): + if repeat % 2 == 0: + t = perf_counter() + bic_result = bic_run() + bic_timings.append(perf_counter() - t) + t = perf_counter() + aff_result = aff_run() + aff_timings.append(perf_counter() - t) + else: + t = perf_counter() + aff_result = aff_run() + aff_timings.append(perf_counter() - t) + t = perf_counter() + bic_result = bic_run() + bic_timings.append(perf_counter() - t) + return bic_timings, bic_result, aff_timings, aff_result + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Compare bioimage-cpp and affogato semantic-mutex-watershed " + "clustering on a RAG built from a simple watershed oversegmentation " + "of the registered semantic-labels volume." + ) + ) + parser.add_argument("--ndim", type=int, default=2, choices=(2, 3)) + parser.add_argument( + "--shape", + type=int, + nargs="+", + default=None, + help=( + "Spatial crop. Default: (448, 448) for ndim=2, (8, 448, 448) for " + "ndim=3." + ), + ) + parser.add_argument( + "--z-start", + type=int, + default=0, + help="Z offset into the (cropped) volume.", + ) + parser.add_argument( + "--seed-spacing", + type=int, + default=4, + help="Grid spacing (in pixels) for the watershed marker grid.", + ) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--rag-threads", + type=int, + default=1, + help="Thread count for RAG construction and feature accumulation.", + ) + args = parser.parse_args() + + shape = tuple(args.shape) if args.shape else ( + (448, 448) if args.ndim == 2 else (8, 448, 448) + ) + + inst, sem, offsets, attr = load_volume(args.ndim, args.z_start, shape) + weights, number_of_offsets = build_weight_volume( + inst, sem, offsets, attr, seed=args.seed + ) + weights = weights.astype(np.float32, copy=False) + n_classes = int(weights.shape[0] - number_of_offsets) + print( + f"loaded: weights={weights.shape} dtype={weights.dtype} " + f"n_classes={n_classes}" + ) + + over_seg = oversegment(weights, attr, args.seed_spacing) + n_segments = int(over_seg.max()) + 1 + print( + f"oversegmentation: shape={over_seg.shape} n_segments~{n_segments} " + f"(seed_spacing={args.seed_spacing})" + ) + + rag, edge_costs, mutex_uvs, mutex_costs, semantic_node_classes, semantic_costs = build_costs( + over_seg, weights, offsets, attr, number_of_offsets, n_classes, + rag_threads=args.rag_threads, + ) + print( + f"rag: nodes={int(rag.number_of_nodes)} edges={int(rag.number_of_edges)} " + f"mutex={mutex_uvs.shape[0]} semantic_pairs={semantic_node_classes.shape[0]}" + ) + + uvs_array = np.ascontiguousarray(rag.uv_ids(), dtype=np.uint64) + + bic_timings, bic_result, aff_timings, aff_result = time_interleaved( + lambda: run_bioimage_cpp( + rag, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ), + lambda: run_affogato_reference( + int(rag.number_of_nodes), uvs_array, mutex_uvs, + semantic_node_classes, edge_costs, mutex_costs, semantic_costs, + ), + args.repeats, + ) + + metrics = compare(bic_result, aff_result) + bic_median = median(bic_timings) + aff_median = median(aff_timings) + speedup = aff_median / bic_median if bic_median > 0 else float("inf") + + print(f"Semantic mutex watershed clustering comparison (ndim={args.ndim})") + print(f"validation metrics: {metrics['validation_source']}") + print( + f"cluster count (bic / affogato): {metrics['n_clusters_bic']} / " + f"{metrics['n_clusters_reference']}" + ) + print( + "VI split/merge: " + f"{metrics['vi_split']:.6g} / {metrics['vi_merge']:.6g}" + ) + print( + "adapted rand error / rand index: " + f"{metrics['adapted_rand_error']:.6g} / {metrics['rand_index']:.12g}" + ) + print(f"partition equal (canonical): {metrics['partition_equal']}") + print(f"semantic match fraction: {metrics['semantic_match_fraction']:.6f}") + print(f"bioimage-cpp median runtime: {bic_median:.6f} s") + print(f"affogato reference median runtime: {aff_median:.6f} s") + print(f"reference / bioimage-cpp runtime ratio: {speedup:.3f}x") + if not metrics["partition_equal"]: + print( + " NOTE: bic and affogato may disagree. affogato's " + "compute_semantic_mws_clustering omits the merge_semantic_labels " + "call on attractive merges (bioimage-cpp's port fixes this) AND " + "calls boost::disjoint_sets::link with raw node ids instead of " + "their roots. Either can shift the partition. Report the metrics." + ) + + +if __name__ == "__main__": + main() diff --git a/development/segmentation/_semantic_mws_equivalence.py b/development/segmentation/_semantic_mws_equivalence.py new file mode 100644 index 0000000..99fb630 --- /dev/null +++ b/development/segmentation/_semantic_mws_equivalence.py @@ -0,0 +1,403 @@ +"""Shared harness for the semantic-mutex-watershed comparison scripts. + +Loads the registered ``semantic_labels`` volume, derives a realistic affinity ++ semantic-class weight stack from it (via +:func:`bioimage_cpp.affinities.compute_affinities` plus Gaussian noise), runs +both bioimage-cpp and affogato side-by-side, and reports partition equivalence +and timing. Not part of the pytest suite — see ``check_semantic_mutex_watershed_{2d,3d}.py``. +""" + +from __future__ import annotations + +import argparse +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + + +_DEFAULT_2D_OFFSETS: list[tuple[int, ...]] = [ + (-1, 0), + (0, -1), + (-3, 0), + (0, -3), + (-9, 0), + (0, -9), +] +_DEFAULT_2D_ATTRACTIVE = 2 + +_DEFAULT_3D_OFFSETS: list[tuple[int, ...]] = [ + (-1, 0, 0), + (0, -1, 0), + (0, 0, -1), + (0, -3, 0), + (0, 0, -3), + (0, -9, 0), + (0, 0, -9), +] +_DEFAULT_3D_ATTRACTIVE = 3 + + +def load_problem() -> tuple[np.ndarray, np.ndarray]: + from bioimage_cpp._data import load_semantic_labels + + instance, semantic = load_semantic_labels() + return np.ascontiguousarray(instance), np.ascontiguousarray(semantic) + + +def prepare_2d_problem( + instance: np.ndarray, + semantic: np.ndarray, + z: int, + yx_shape: tuple[int, int], +): + y, x = yx_shape + inst = np.ascontiguousarray(instance[z, :y, :x]) + sem = np.ascontiguousarray(semantic[z, :y, :x]) + offsets = list(_DEFAULT_2D_OFFSETS) + return inst, sem, offsets, _DEFAULT_2D_ATTRACTIVE + + +def prepare_3d_problem( + instance: np.ndarray, + semantic: np.ndarray, + zyx_shape: tuple[int, int, int], + *, + z_start: int = 0, +): + z, y, x = zyx_shape + inst = np.ascontiguousarray(instance[z_start : z_start + z, :y, :x]) + sem = np.ascontiguousarray(semantic[z_start : z_start + z, :y, :x]) + offsets = list(_DEFAULT_3D_OFFSETS) + return inst, sem, offsets, _DEFAULT_3D_ATTRACTIVE + + +def build_weight_volume( + instance_labels: np.ndarray, + semantic_labels: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, + *, + n_classes: int | None = None, + high: float = 0.95, + low: float = 0.05, + noise: float = 0.05, + seed: int = 42, +) -> tuple[np.ndarray, int]: + """Derive an MWS-style weight stack from instance + semantic ground truth. + + Spatial channels are computed via + :func:`bioimage_cpp.affinities.compute_affinities`, which returns ``1`` when + two endpoints share an instance and ``0`` otherwise. We keep this + convention for **attractive** channels (high = merge), invert it for + **mutex** channels (high = repel), rescale into ``[low, high]``, and add + Gaussian noise. + + Semantic channels (one per class id ``c`` in ``[0, n_classes)``) start at + ``low`` everywhere and rise to ``high`` where ``semantic_labels == c``, + then receive the same noise treatment. If ``n_classes`` is ``None`` we + derive it as ``int(semantic_labels.max()) + 1`` so the channel index of + each class equals its source class id. + + Returns ``(weights, number_of_offsets)``; the second item lets callers + locate the boundary between spatial and semantic channels in the stack. + """ + import bioimage_cpp as bic + + number_of_offsets = len(offsets) + if number_of_attractive_channels > number_of_offsets: + raise ValueError( + "number_of_attractive_channels must be <= len(offsets), got " + f"{number_of_attractive_channels} vs {number_of_offsets}" + ) + + if n_classes is None: + n_classes = int(semantic_labels.max()) + 1 + if n_classes < 1: + raise ValueError( + "could not infer n_classes from semantic_labels; pass it explicitly" + ) + + rng = np.random.default_rng(seed) + + affs, _ = bic.affinities.compute_affinities( + instance_labels, offsets, return_mask=True + ) + affs = affs.astype(np.float32, copy=False) + # Mutex channels: invert so high = boundary = repel. + if number_of_offsets > number_of_attractive_channels: + affs[number_of_attractive_channels:number_of_offsets] = ( + 1.0 - affs[number_of_attractive_channels:number_of_offsets] + ) + + spatial_weights = affs * (high - low) + low + spatial_weights += rng.normal(loc=0.0, scale=noise, size=spatial_weights.shape).astype( + np.float32, copy=False + ) + + spatial_shape = instance_labels.shape + semantic_weights = np.full( + (n_classes, *spatial_shape), low, dtype=np.float32 + ) + for c in range(n_classes): + semantic_weights[c][semantic_labels == c] = high + semantic_weights += rng.normal( + loc=0.0, scale=noise, size=semantic_weights.shape + ).astype(np.float32, copy=False) + + weights = np.concatenate([spatial_weights, semantic_weights], axis=0).astype( + np.float64, copy=False + ) + np.clip(weights, 0.0, 1.0, out=weights) + + # Break ties deterministically: the affogato wrapper sorts via + # ``np.argsort`` (unstable), while bioimage-cpp's C++ sort tie-breaks by + # edge id. With many derived weights clipped to exactly 0.0 or 1.0, the + # difference in tie-break ordering massively rearranges processing, + # fragmenting the output. We subtract a tiny per-edge perturbation that + # grows with linear edge id, so smaller-id edges always sort first in + # descending order — matching bioimage-cpp's explicit tiebreak and + # removing the only source of disagreement. Float64 has enough precision + # for the perturbation to survive across ~1M edges. + perturbation = np.arange(weights.size, dtype=np.float64).reshape(weights.shape) + weights -= perturbation * 1.0e-12 + return np.ascontiguousarray(weights), number_of_offsets + + +def run_bioimage_cpp( + weights: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, +) -> tuple[np.ndarray, np.ndarray]: + import bioimage_cpp as bic + + return bic.segmentation.semantic_mutex_watershed( + weights, + offsets, + number_of_attractive_channels=number_of_attractive_channels, + ) + + +def run_affogato_reference( + weights: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, +) -> tuple[np.ndarray, np.ndarray]: + from affogato.segmentation import compute_semantic_mws_segmentation + + spatial_ndim = len(offsets[0]) + labels, semantic = compute_semantic_mws_segmentation( + weights, + offsets, + number_of_attractive_channels=number_of_attractive_channels, + strides=[1] * spatial_ndim, + ) + return labels, semantic + + +def _load_validation_metrics(): + try: + from elf.validation import rand_index, variation_of_information + + return "elf.validation", rand_index, variation_of_information + except ImportError: + from elf.evaluation import rand_index, variation_of_information + + return "elf.evaluation", rand_index, variation_of_information + + +def compare_semantic_segmentations( + candidate: tuple[np.ndarray, np.ndarray], + reference: tuple[np.ndarray, np.ndarray], + *, + max_vi: float = 1.0e-10, + max_are: float = 1.0e-10, + min_rand_index: float = 1.0 - 1.0e-10, +) -> dict: + """Compute partition + semantic-label agreement metrics. + + Does **not** raise when results disagree. The two implementations are + expected to disagree on this data because affogato's C++ kernel calls + ``boost::disjoint_sets::link(u, v)`` with the raw node ids instead of + their union-find roots; with multi-class semantic inputs that + miscompounds the tree and fragments the output. The harness reports the + discrepancy and lets the caller decide what to do with it. + """ + cand_labels, cand_semantic = candidate + ref_labels, ref_semantic = reference + source, rand_index, variation_of_information = _load_validation_metrics() + + vi_split, vi_merge = variation_of_information(cand_labels, ref_labels) + adapted_rand_error, ri = rand_index(cand_labels, ref_labels) + exact_labels = bool(np.array_equal(cand_labels, ref_labels)) + partition_equivalent = ( + vi_split <= max_vi + and vi_merge <= max_vi + and adapted_rand_error <= max_are + and ri >= min_rand_index + ) + + semantic_exact = bool(np.array_equal(cand_semantic, ref_semantic)) + semantic_match_fraction = float(np.mean(cand_semantic == ref_semantic)) + + return { + "validation_source": source, + "vi_split": float(vi_split), + "vi_merge": float(vi_merge), + "adapted_rand_error": float(adapted_rand_error), + "rand_index": float(ri), + "exact_label_equality": exact_labels, + "partition_equivalent": partition_equivalent, + "semantic_exact_equality": semantic_exact, + "semantic_match_fraction": semantic_match_fraction, + "n_clusters_bic": int(np.unique(cand_labels).size), + "n_clusters_reference": int(np.unique(ref_labels).size), + } + + +def time_functions_interleaved( + first: Callable[..., tuple[np.ndarray, np.ndarray]], + second: Callable[..., tuple[np.ndarray, np.ndarray]], + weights: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, + repeats: int, +) -> tuple[list[float], tuple[np.ndarray, np.ndarray], list[float], tuple[np.ndarray, np.ndarray]]: + def timed_call(run): + start = perf_counter() + result = run(weights, offsets, number_of_attractive_channels) + return perf_counter() - start, result + + first(weights, offsets, number_of_attractive_channels) + second(weights, offsets, number_of_attractive_channels) + + first_timings: list[float] = [] + second_timings: list[float] = [] + first_result = None + second_result = None + for repeat in range(repeats): + if repeat % 2 == 0: + first_time, first_result = timed_call(first) + second_time, second_result = timed_call(second) + else: + second_time, second_result = timed_call(second) + first_time, first_result = timed_call(first) + first_timings.append(first_time) + second_timings.append(second_time) + + assert first_result is not None + assert second_result is not None + return first_timings, first_result, second_timings, second_result + + +def print_report( + *, + ndim: int, + weights: np.ndarray, + n_classes: int, + metrics: dict, + bic_timings: list[float], + ref_timings: list[float], +) -> None: + bic_median = median(bic_timings) + ref_median = median(ref_timings) + speedup = ref_median / bic_median if bic_median > 0 else float("inf") + + print(f"Semantic mutex watershed {ndim}D comparison") + print( + f"weights shape: {weights.shape}, dtype: {weights.dtype}, " + f"n_classes: {n_classes}" + ) + print(f"validation metrics: {metrics['validation_source']}") + print( + f"cluster count (bic / affogato): {metrics['n_clusters_bic']} / " + f"{metrics['n_clusters_reference']}" + ) + print( + "VI split/merge: " + f"{metrics['vi_split']:.6g} / {metrics['vi_merge']:.6g}" + ) + print( + "adapted rand error / rand index: " + f"{metrics['adapted_rand_error']:.6g} / {metrics['rand_index']:.12g}" + ) + print(f"exact label equality: {metrics['exact_label_equality']}") + print(f"partition_equivalent: {metrics['partition_equivalent']}") + print( + "semantic exact equality / match fraction: " + f"{metrics['semantic_exact_equality']} / " + f"{metrics['semantic_match_fraction']:.6f}" + ) + print(f"bioimage-cpp median runtime: {bic_median:.6f} s") + print(f"affogato reference median runtime: {ref_median:.6f} s") + print(f"reference / bioimage-cpp runtime ratio: {speedup:.3f}x") + if not metrics["partition_equivalent"]: + print( + " NOTE: bic and affogato disagree. Affogato's C++ kernel calls " + "boost::disjoint_sets::link(u, v) with raw node ids instead of " + "their roots; with multi-class semantic inputs this corrupts " + "the union-find tree and fragments the output. bioimage-cpp's " + "result agrees with a from-scratch Python reference." + ) + + +def run_check( + *, + ndim: int, + repeats: int, + z: int, + yx_shape: tuple[int, int], + zyx_shape: tuple[int, int, int], + seed: int = 42, + z_start: int = 0, +) -> None: + instance, semantic = load_problem() + if ndim == 2: + inst, sem, offsets, attractive_channels = prepare_2d_problem( + instance, semantic, z=z, yx_shape=yx_shape + ) + elif ndim == 3: + inst, sem, offsets, attractive_channels = prepare_3d_problem( + instance, semantic, zyx_shape=zyx_shape, z_start=z_start + ) + else: + raise ValueError(f"ndim must be 2 or 3, got {ndim}") + + weights, _ = build_weight_volume( + inst, sem, offsets, attractive_channels, seed=seed + ) + n_classes = int(weights.shape[0] - len(offsets)) + + ref_timings, ref_result, bic_timings, bic_result = time_functions_interleaved( + run_affogato_reference, + run_bioimage_cpp, + weights, + offsets, + attractive_channels, + repeats, + ) + metrics = compare_semantic_segmentations(bic_result, ref_result) + print_report( + ndim=ndim, + weights=weights, + n_classes=n_classes, + metrics=metrics, + bic_timings=bic_timings, + ref_timings=ref_timings, + ) + + +def add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--repeats", + type=int, + default=3, + help="Number of timed runs for each implementation.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Seed for the noise added to the derived weights.", + ) diff --git a/development/segmentation/check_semantic_mutex_watershed_2d.py b/development/segmentation/check_semantic_mutex_watershed_2d.py new file mode 100644 index 0000000..8833a1a --- /dev/null +++ b/development/segmentation/check_semantic_mutex_watershed_2d.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import argparse + +from _semantic_mws_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and affogato semantic mutex watershed in 2D." + ) + add_common_arguments(parser) + parser.add_argument( + "--z", + type=int, + default=4, + help="Z slice from the (cropped) registered volume to use for the 2D check.", + ) + parser.add_argument( + "--shape", + type=int, + nargs=2, + metavar=("Y", "X"), + default=(448, 448), + help="Spatial crop shape for the 2D check.", + ) + args = parser.parse_args() + + run_check( + ndim=2, + repeats=args.repeats, + z=args.z, + yx_shape=tuple(args.shape), + zyx_shape=(0, 0, 0), + seed=args.seed, + ) + + +if __name__ == "__main__": + main() diff --git a/development/segmentation/check_semantic_mutex_watershed_3d.py b/development/segmentation/check_semantic_mutex_watershed_3d.py new file mode 100644 index 0000000..984a3be --- /dev/null +++ b/development/segmentation/check_semantic_mutex_watershed_3d.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import argparse + +from _semantic_mws_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and affogato semantic mutex watershed in 3D." + ) + add_common_arguments(parser) + parser.add_argument( + "--shape", + type=int, + nargs=3, + metavar=("Z", "Y", "X"), + default=(8, 448, 448), + help="Spatial crop shape for the 3D check.", + ) + parser.add_argument( + "--z-start", + type=int, + default=0, + help="Z offset into the (cropped) volume.", + ) + args = parser.parse_args() + + run_check( + ndim=3, + repeats=args.repeats, + z=0, + yx_shape=(0, 0), + zyx_shape=tuple(args.shape), + seed=args.seed, + z_start=args.z_start, + ) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/detail/semantic_labels.hxx b/include/bioimage_cpp/detail/semantic_labels.hxx new file mode 100644 index 0000000..2afa2bc --- /dev/null +++ b/include/bioimage_cpp/detail/semantic_labels.hxx @@ -0,0 +1,57 @@ +#pragma once + +#include +#include + +namespace bioimage_cpp { + +using SemanticLabeling = std::vector; + +// Returns true if the two roots already carry different non-negative semantic +// labels. Merging two such roots would conflict the semantic assignments, so +// the caller must skip that merge. +inline bool check_semantic_constraint( + const std::uint64_t first, + const std::uint64_t second, + const SemanticLabeling &semantic_labels +) { + const auto label_first = semantic_labels[first]; + const auto label_second = semantic_labels[second]; + if (label_first >= 0 && label_second >= 0) { + return label_first != label_second; + } + return false; +} + +// Assigns `class_id` to `root` only if `root` has no semantic label yet +// (-1 == unassigned). Already-assigned roots keep their first assignment; +// because semantic edges are processed in descending-weight order this means +// each root keeps the strongest class assignment seen for it. +inline void assign_semantic_label( + const std::uint64_t root, + const std::int64_t class_id, + SemanticLabeling &semantic_labels +) { + if (semantic_labels[root] < 0) { + semantic_labels[root] = class_id; + } +} + +// After two roots merge, propagate a non-negative semantic label from one to +// the other if exactly one was assigned. If both are assigned the caller is +// expected to have rejected the merge via `check_semantic_constraint`. +inline void merge_semantic_labels( + const std::uint64_t first, + const std::uint64_t second, + SemanticLabeling &semantic_labels +) { + const auto label_first = semantic_labels[first]; + const auto label_second = semantic_labels[second]; + if (label_first >= 0 && label_second < 0) { + semantic_labels[second] = label_first; + } else if (label_first < 0 && label_second >= 0) { + semantic_labels[first] = label_second; + } +} + +} // namespace bioimage_cpp diff --git a/include/bioimage_cpp/graph/mutex_watershed.hxx b/include/bioimage_cpp/graph/mutex_watershed.hxx index bf929be..85b8eaa 100644 --- a/include/bioimage_cpp/graph/mutex_watershed.hxx +++ b/include/bioimage_cpp/graph/mutex_watershed.hxx @@ -1,144 +1,4 @@ #pragma once -#include "bioimage_cpp/detail/mutex_storage.hxx" -#include "bioimage_cpp/detail/relabel.hxx" -#include "bioimage_cpp/detail/union_find.hxx" -#include "bioimage_cpp/graph/undirected_graph.hxx" - -#include -#include -#include -#include -#include -#include -#include - -namespace bioimage_cpp::graph { - -// Mutex watershed clustering on an arbitrary undirected graph. -// -// `graph` defines the attractive edges with one cost per edge in `edge_costs`. -// `mutex_uvs` defines the long-range repulsive (mutex) edges as (u, v) pairs -// with one cost per edge in `mutex_costs`. Higher costs win — they are -// processed first when sorting jointly by descending weight. -// -// Returns dense node labels in [0, k) following first-occurrence order. -// -// This is a port of `compute_mws_clustering` from the affogato library, -// adapted to bioimage-cpp's UndirectedGraph and detail/ primitives. -// -// Templated on the weight type. Concrete instantiations for `float` and -// `double` are provided by the binding layer; other floating types are -// supported but must be instantiated explicitly. -template -std::vector mutex_watershed_clustering( - const UndirectedGraph &graph, - const std::vector &edge_costs, - const std::vector> &mutex_uvs, - const std::vector &mutex_costs -) { - const auto number_of_edges = static_cast(graph.number_of_edges()); - if (edge_costs.size() != number_of_edges) { - throw std::invalid_argument( - "edge_costs size must match graph.number_of_edges(), got edge_costs size=" + - std::to_string(edge_costs.size()) + - ", number_of_edges=" + std::to_string(number_of_edges) - ); - } - if (mutex_costs.size() != mutex_uvs.size()) { - throw std::invalid_argument( - "mutex_costs size must match mutex_uvs size, got mutex_costs size=" + - std::to_string(mutex_costs.size()) + - ", mutex_uvs size=" + std::to_string(mutex_uvs.size()) - ); - } - - const auto number_of_nodes = static_cast(graph.number_of_nodes()); - const auto number_of_mutex = mutex_uvs.size(); - - for (std::size_t index = 0; index < number_of_mutex; ++index) { - const auto u = mutex_uvs[index][0]; - const auto v = mutex_uvs[index][1]; - if (u >= number_of_nodes || v >= number_of_nodes) { - throw std::invalid_argument( - "mutex_uvs endpoints must be < number_of_nodes, got u=" + - std::to_string(u) + ", v=" + std::to_string(v) + - ", number_of_nodes=" + std::to_string(number_of_nodes) - ); - } - } - - struct WeightedEdge { - WeightT weight; - std::uint64_t index; - bool is_mutex; - }; - - std::vector edge_order; - edge_order.reserve(number_of_edges + number_of_mutex); - for (std::size_t index = 0; index < number_of_edges; ++index) { - edge_order.push_back( - WeightedEdge{edge_costs[index], static_cast(index), false} - ); - } - for (std::size_t index = 0; index < number_of_mutex; ++index) { - edge_order.push_back( - WeightedEdge{mutex_costs[index], static_cast(index), true} - ); - } - - std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) { - if (first.weight != second.weight) { - return first.weight > second.weight; - } - if (first.is_mutex != second.is_mutex) { - return !first.is_mutex; - } - return first.index < second.index; - }); - - bioimage_cpp::detail::UnionFind sets(number_of_nodes); - MutexStorage mutexes(number_of_nodes); - - for (const auto &edge : edge_order) { - std::uint64_t u; - std::uint64_t v; - if (edge.is_mutex) { - const auto &pair = mutex_uvs[static_cast(edge.index)]; - u = pair[0]; - v = pair[1]; - } else { - const auto uv = graph.uv(edge.index); - u = uv.first; - v = uv.second; - } - if (u == v) { - continue; - } - - const auto root_u = sets.find(u); - const auto root_v = sets.find(v); - if (root_u == root_v) { - continue; - } - - if (edge.is_mutex) { - insert_mutex(root_u, root_v, mutexes); - } else { - if (check_mutex(root_u, root_v, mutexes)) { - continue; - } - const auto new_root = sets.unite_roots(root_u, root_v); - const auto old_root = (new_root == root_u) ? root_v : root_u; - merge_mutexes(old_root, new_root, mutexes); - } - } - - std::vector roots(number_of_nodes); - for (std::size_t node = 0; node < number_of_nodes; ++node) { - roots[node] = sets.find(static_cast(node)); - } - return bioimage_cpp::detail::dense_relabel(roots); -} - -} // namespace bioimage_cpp::graph +#include "bioimage_cpp/graph/mutex_watershed/clustering.hxx" +#include "bioimage_cpp/graph/mutex_watershed/semantic.hxx" diff --git a/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx b/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx new file mode 100644 index 0000000..bf929be --- /dev/null +++ b/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx @@ -0,0 +1,144 @@ +#pragma once + +#include "bioimage_cpp/detail/mutex_storage.hxx" +#include "bioimage_cpp/detail/relabel.hxx" +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// Mutex watershed clustering on an arbitrary undirected graph. +// +// `graph` defines the attractive edges with one cost per edge in `edge_costs`. +// `mutex_uvs` defines the long-range repulsive (mutex) edges as (u, v) pairs +// with one cost per edge in `mutex_costs`. Higher costs win — they are +// processed first when sorting jointly by descending weight. +// +// Returns dense node labels in [0, k) following first-occurrence order. +// +// This is a port of `compute_mws_clustering` from the affogato library, +// adapted to bioimage-cpp's UndirectedGraph and detail/ primitives. +// +// Templated on the weight type. Concrete instantiations for `float` and +// `double` are provided by the binding layer; other floating types are +// supported but must be instantiated explicitly. +template +std::vector mutex_watershed_clustering( + const UndirectedGraph &graph, + const std::vector &edge_costs, + const std::vector> &mutex_uvs, + const std::vector &mutex_costs +) { + const auto number_of_edges = static_cast(graph.number_of_edges()); + if (edge_costs.size() != number_of_edges) { + throw std::invalid_argument( + "edge_costs size must match graph.number_of_edges(), got edge_costs size=" + + std::to_string(edge_costs.size()) + + ", number_of_edges=" + std::to_string(number_of_edges) + ); + } + if (mutex_costs.size() != mutex_uvs.size()) { + throw std::invalid_argument( + "mutex_costs size must match mutex_uvs size, got mutex_costs size=" + + std::to_string(mutex_costs.size()) + + ", mutex_uvs size=" + std::to_string(mutex_uvs.size()) + ); + } + + const auto number_of_nodes = static_cast(graph.number_of_nodes()); + const auto number_of_mutex = mutex_uvs.size(); + + for (std::size_t index = 0; index < number_of_mutex; ++index) { + const auto u = mutex_uvs[index][0]; + const auto v = mutex_uvs[index][1]; + if (u >= number_of_nodes || v >= number_of_nodes) { + throw std::invalid_argument( + "mutex_uvs endpoints must be < number_of_nodes, got u=" + + std::to_string(u) + ", v=" + std::to_string(v) + + ", number_of_nodes=" + std::to_string(number_of_nodes) + ); + } + } + + struct WeightedEdge { + WeightT weight; + std::uint64_t index; + bool is_mutex; + }; + + std::vector edge_order; + edge_order.reserve(number_of_edges + number_of_mutex); + for (std::size_t index = 0; index < number_of_edges; ++index) { + edge_order.push_back( + WeightedEdge{edge_costs[index], static_cast(index), false} + ); + } + for (std::size_t index = 0; index < number_of_mutex; ++index) { + edge_order.push_back( + WeightedEdge{mutex_costs[index], static_cast(index), true} + ); + } + + std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) { + if (first.weight != second.weight) { + return first.weight > second.weight; + } + if (first.is_mutex != second.is_mutex) { + return !first.is_mutex; + } + return first.index < second.index; + }); + + bioimage_cpp::detail::UnionFind sets(number_of_nodes); + MutexStorage mutexes(number_of_nodes); + + for (const auto &edge : edge_order) { + std::uint64_t u; + std::uint64_t v; + if (edge.is_mutex) { + const auto &pair = mutex_uvs[static_cast(edge.index)]; + u = pair[0]; + v = pair[1]; + } else { + const auto uv = graph.uv(edge.index); + u = uv.first; + v = uv.second; + } + if (u == v) { + continue; + } + + const auto root_u = sets.find(u); + const auto root_v = sets.find(v); + if (root_u == root_v) { + continue; + } + + if (edge.is_mutex) { + insert_mutex(root_u, root_v, mutexes); + } else { + if (check_mutex(root_u, root_v, mutexes)) { + continue; + } + const auto new_root = sets.unite_roots(root_u, root_v); + const auto old_root = (new_root == root_u) ? root_v : root_u; + merge_mutexes(old_root, new_root, mutexes); + } + } + + std::vector roots(number_of_nodes); + for (std::size_t node = 0; node < number_of_nodes; ++node) { + roots[node] = sets.find(static_cast(node)); + } + return bioimage_cpp::detail::dense_relabel(roots); +} + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx b/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx new file mode 100644 index 0000000..e665fcd --- /dev/null +++ b/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx @@ -0,0 +1,212 @@ +#pragma once + +#include "bioimage_cpp/detail/mutex_storage.hxx" +#include "bioimage_cpp/detail/relabel.hxx" +#include "bioimage_cpp/detail/semantic_labels.hxx" +#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +// Result of a semantic mutex watershed clustering: a dense node labeling and a +// per-node semantic class label. Unassigned nodes report ``-1``. +struct SemanticMutexWatershedResult { + std::vector node_labels; + std::vector semantic_labels; +}; + +// Semantic mutex watershed on an arbitrary undirected graph. +// +// `graph` defines the attractive edges with one cost per edge in `edge_costs`. +// `mutex_uvs` defines the long-range repulsive (mutex) edges with one cost per +// row in `mutex_costs`. `semantic_node_classes` is an (S, 2) table whose first +// column is a node id and whose second column is a non-negative semantic class +// id, scored by `semantic_costs`. All three edge groups are sorted jointly by +// descending weight and processed in one pass: +// - attractive: merge the two endpoints' clusters unless a mutex constraint +// already separates them OR the clusters already carry different semantic +// classes. +// - mutex: insert a hard separation between the two endpoints' clusters. +// - semantic: tag the endpoint's current cluster with the given class id if +// the cluster is still unassigned. Semantic labels propagate when clusters +// merge via attractive edges. +// +// Ties in weight are broken by edge group (attractive < mutex < semantic) and +// then by within-group index, for determinism. The returned `node_labels` are +// dense in [0, k) in first-occurrence order; `semantic_labels` is one +// `std::int64_t` per node, equal to the cluster's class id or ``-1`` if the +// cluster received no semantic assignment. +// +// Ported from `compute_semantic_mws_clustering` in the affogato library. +template +SemanticMutexWatershedResult semantic_mutex_watershed_clustering( + const UndirectedGraph &graph, + const std::vector &edge_costs, + const std::vector> &mutex_uvs, + const std::vector &mutex_costs, + const std::vector> &semantic_node_classes, + const std::vector &semantic_costs +) { + const auto number_of_edges = static_cast(graph.number_of_edges()); + if (edge_costs.size() != number_of_edges) { + throw std::invalid_argument( + "edge_costs size must match graph.number_of_edges(), got edge_costs size=" + + std::to_string(edge_costs.size()) + + ", number_of_edges=" + std::to_string(number_of_edges) + ); + } + if (mutex_costs.size() != mutex_uvs.size()) { + throw std::invalid_argument( + "mutex_costs size must match mutex_uvs size, got mutex_costs size=" + + std::to_string(mutex_costs.size()) + + ", mutex_uvs size=" + std::to_string(mutex_uvs.size()) + ); + } + if (semantic_costs.size() != semantic_node_classes.size()) { + throw std::invalid_argument( + "semantic_costs size must match semantic_node_classes size, got semantic_costs size=" + + std::to_string(semantic_costs.size()) + + ", semantic_node_classes size=" + std::to_string(semantic_node_classes.size()) + ); + } + + const auto number_of_nodes = static_cast(graph.number_of_nodes()); + const auto number_of_mutex = mutex_uvs.size(); + const auto number_of_semantic = semantic_node_classes.size(); + + for (std::size_t index = 0; index < number_of_mutex; ++index) { + const auto u = mutex_uvs[index][0]; + const auto v = mutex_uvs[index][1]; + if (u >= number_of_nodes || v >= number_of_nodes) { + throw std::invalid_argument( + "mutex_uvs endpoints must be < number_of_nodes, got u=" + + std::to_string(u) + ", v=" + std::to_string(v) + + ", number_of_nodes=" + std::to_string(number_of_nodes) + ); + } + } + for (std::size_t index = 0; index < number_of_semantic; ++index) { + const auto node = semantic_node_classes[index][0]; + if (node >= number_of_nodes) { + throw std::invalid_argument( + "semantic_node_classes node ids must be < number_of_nodes, got node=" + + std::to_string(node) + + ", number_of_nodes=" + std::to_string(number_of_nodes) + ); + } + } + + enum class EdgeKind : std::uint8_t { + Attractive = 0, + Mutex = 1, + Semantic = 2, + }; + + struct WeightedEdge { + WeightT weight; + std::uint64_t index; + EdgeKind kind; + }; + + std::vector edge_order; + edge_order.reserve(number_of_edges + number_of_mutex + number_of_semantic); + for (std::size_t index = 0; index < number_of_edges; ++index) { + edge_order.push_back( + WeightedEdge{edge_costs[index], static_cast(index), EdgeKind::Attractive} + ); + } + for (std::size_t index = 0; index < number_of_mutex; ++index) { + edge_order.push_back( + WeightedEdge{mutex_costs[index], static_cast(index), EdgeKind::Mutex} + ); + } + for (std::size_t index = 0; index < number_of_semantic; ++index) { + edge_order.push_back( + WeightedEdge{semantic_costs[index], static_cast(index), EdgeKind::Semantic} + ); + } + + std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) { + if (first.weight != second.weight) { + return first.weight > second.weight; + } + if (first.kind != second.kind) { + return static_cast(first.kind) < static_cast(second.kind); + } + return first.index < second.index; + }); + + bioimage_cpp::detail::UnionFind sets(number_of_nodes); + MutexStorage mutexes(number_of_nodes); + SemanticLabeling semantic_labels(number_of_nodes, -1); + + for (const auto &edge : edge_order) { + const auto edge_index = static_cast(edge.index); + if (edge.kind == EdgeKind::Semantic) { + const auto node = semantic_node_classes[edge_index][0]; + const auto class_id = static_cast(semantic_node_classes[edge_index][1]); + const auto root = sets.find(node); + assign_semantic_label(root, class_id, semantic_labels); + continue; + } + + std::uint64_t u; + std::uint64_t v; + if (edge.kind == EdgeKind::Mutex) { + const auto &pair = mutex_uvs[edge_index]; + u = pair[0]; + v = pair[1]; + } else { + const auto uv = graph.uv(edge.index); + u = uv.first; + v = uv.second; + } + if (u == v) { + continue; + } + + const auto root_u = sets.find(u); + const auto root_v = sets.find(v); + if (root_u == root_v) { + continue; + } + if (check_semantic_constraint(root_u, root_v, semantic_labels)) { + continue; + } + if (check_mutex(root_u, root_v, mutexes)) { + continue; + } + + if (edge.kind == EdgeKind::Mutex) { + insert_mutex(root_u, root_v, mutexes); + } else { + const auto new_root = sets.unite_roots(root_u, root_v); + const auto old_root = (new_root == root_u) ? root_v : root_u; + merge_mutexes(old_root, new_root, mutexes); + merge_semantic_labels(new_root, old_root, semantic_labels); + } + } + + std::vector roots(number_of_nodes); + std::vector semantic_out(number_of_nodes); + for (std::size_t node = 0; node < number_of_nodes; ++node) { + const auto root = sets.find(static_cast(node)); + roots[node] = root; + semantic_out[node] = semantic_labels[root]; + } + + SemanticMutexWatershedResult result; + result.node_labels = bioimage_cpp::detail::dense_relabel(roots); + result.semantic_labels = std::move(semantic_out); + return result; +} + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx b/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx new file mode 100644 index 0000000..ff620eb --- /dev/null +++ b/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx @@ -0,0 +1,198 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/mutex_storage.hxx" +#include "bioimage_cpp/detail/semantic_labels.hxx" +#include "bioimage_cpp/detail/union_find.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp { + +// Semantic mutex watershed on a 2D or 3D image-derived grid graph. +// +// `affinities` carries three groups of channels stacked along axis 0: +// - [0, number_of_attractive_channels) : attractive grid edges +// - [number_of_attractive_channels, number_of_offsets) : mutex grid edges +// - [number_of_offsets, channels) : semantic-class affinities (one per class) +// The grid offsets in `offsets` apply to the first `number_of_offsets` +// channels; semantic channels are not spatial edges and so have no offset. +// +// `valid_edges` shares the affinity shape and gates every edge (including +// semantic ones) on/off; the Python wrapper computes this mask. Output +// `node_labels_out` is 1-based dense node labels with shape `affinities.shape[1:]` +// (matching `mutex_watershed_grid`). Output `semantic_labels_out` has the same +// spatial shape and an ``int64`` value per node, ``-1`` for clusters that +// received no semantic assignment. +// +// Ported from `compute_semantic_mws_segmentation` in the affogato library. +template +void semantic_mutex_watershed_grid( + const ConstArrayView &affinities, + const ConstArrayView &valid_edges, + const std::vector> &offsets, + const std::size_t number_of_attractive_channels, + const std::size_t number_of_offsets, + const ArrayView &node_labels_out, + const ArrayView &semantic_labels_out +) { + if (affinities.ndim() != 3 && affinities.ndim() != 4) { + throw std::invalid_argument( + "affinities must have shape (channels, y, x) or (channels, z, y, x), got ndim=" + + std::to_string(affinities.ndim()) + ); + } + if (offsets.empty()) { + throw std::invalid_argument("offsets must not be empty"); + } + if (valid_edges.shape != affinities.shape) { + throw std::invalid_argument("valid_edges shape must match affinities shape"); + } + + const auto number_of_channels = static_cast(affinities.shape[0]); + const auto spatial_ndim = static_cast(affinities.ndim() - 1); + if (offsets.size() != number_of_offsets) { + throw std::invalid_argument( + "offsets length must equal number_of_offsets, got offsets length=" + + std::to_string(offsets.size()) + + ", number_of_offsets=" + std::to_string(number_of_offsets) + ); + } + if (number_of_attractive_channels > number_of_offsets) { + throw std::invalid_argument( + "number_of_attractive_channels must be <= number_of_offsets" + ); + } + if (number_of_offsets > number_of_channels) { + throw std::invalid_argument( + "number_of_offsets must be <= number of affinity channels" + ); + } + for (const auto &offset : offsets) { + if (offset.size() != spatial_ndim) { + throw std::invalid_argument( + "each offset must have length matching the spatial ndim, got spatial ndim=" + + std::to_string(spatial_ndim) + ); + } + } + + std::vector spatial_shape( + affinities.shape.begin() + 1, + affinities.shape.end() + ); + if (node_labels_out.shape != spatial_shape) { + throw std::invalid_argument("node_labels_out shape must match affinities spatial shape"); + } + if (semantic_labels_out.shape != spatial_shape) { + throw std::invalid_argument( + "semantic_labels_out shape must match affinities spatial shape" + ); + } + + const auto number_of_nodes = static_cast(std::accumulate( + spatial_shape.begin(), + spatial_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + )); + const auto spatial_strides = detail::c_order_strides(spatial_shape); + + std::vector offset_strides(number_of_offsets, 0); + for (std::size_t channel = 0; channel < number_of_offsets; ++channel) { + for (std::size_t axis = 0; axis < spatial_ndim; ++axis) { + offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis]; + } + } + + struct WeightedGridEdge { + T weight; + std::uint64_t id; + }; + + const auto number_of_edges = number_of_nodes * number_of_channels; + std::vector edge_order; + edge_order.reserve(static_cast(number_of_edges)); + for (std::uint64_t edge_id = 0; edge_id < number_of_edges; ++edge_id) { + if (valid_edges.data[edge_id] != 0) { + edge_order.push_back(WeightedGridEdge{affinities.data[edge_id], edge_id}); + } + } + std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) { + if (first.weight == second.weight) { + return first.id < second.id; + } + return first.weight > second.weight; + }); + + detail::UnionFind sets(static_cast(number_of_nodes)); + MutexStorage mutexes(static_cast(number_of_nodes)); + SemanticLabeling semantic_labels(static_cast(number_of_nodes), -1); + + const auto number_of_attractive_edges = number_of_nodes * static_cast( + number_of_attractive_channels + ); + const auto number_of_spatial_edges = number_of_nodes * static_cast( + number_of_offsets + ); + + for (const auto &edge : edge_order) { + const auto edge_id = edge.id; + const auto channel = static_cast(edge_id / number_of_nodes); + const auto u = edge_id % number_of_nodes; + + if (edge_id >= number_of_spatial_edges) { + const auto class_id = static_cast(channel - number_of_offsets); + const auto root_u = sets.find(u); + assign_semantic_label(root_u, class_id, semantic_labels); + continue; + } + + const auto v_signed = static_cast(u) + + static_cast(offset_strides[channel]); + const auto v = static_cast(v_signed); + const auto root_u = sets.find(u); + const auto root_v = sets.find(v); + if (root_u == root_v) { + continue; + } + if (check_semantic_constraint(root_u, root_v, semantic_labels)) { + continue; + } + if (check_mutex(root_u, root_v, mutexes)) { + continue; + } + + const bool is_mutex_edge = edge_id >= number_of_attractive_edges; + if (is_mutex_edge) { + insert_mutex(root_u, root_v, mutexes); + } else { + const auto new_root = sets.unite_roots(root_u, root_v); + const auto old_root = (new_root == root_u) ? root_v : root_u; + merge_mutexes(old_root, new_root, mutexes); + merge_semantic_labels(new_root, old_root, semantic_labels); + } + } + + std::vector root_labels(static_cast(number_of_nodes), 0); + std::uint64_t next_label = 1; + for (std::uint64_t node = 0; node < number_of_nodes; ++node) { + const auto root = sets.find(node); + auto &label = root_labels[static_cast(root)]; + if (label == 0) { + label = next_label; + ++next_label; + } + node_labels_out.data[node] = label; + semantic_labels_out.data[node] = semantic_labels[root]; + } +} + +} // namespace bioimage_cpp diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 245fdc8..7917dfe 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -823,6 +823,64 @@ UInt64Array mutex_watershed_clustering_t( return vector_to_uint64_array(labels); } +template +std::pair semantic_mutex_watershed_clustering_t( + const Graph &graph, + ConstArray1D edge_costs, + ConstUInt64Array mutex_uvs, + ConstArray1D mutex_costs, + ConstUInt64Array semantic_node_classes, + ConstArray1D semantic_costs +) { + const auto edge_cost_vector = + array_1d_to_vector(edge_costs, "edge_costs", graph.number_of_edges()); + require_uv_array(mutex_uvs, "mutex_uvs"); + const auto n_mutex = mutex_uvs.shape(0); + const auto mutex_cost_vector = + array_1d_to_vector(mutex_costs, "mutex_costs", static_cast(n_mutex)); + require_uv_array(semantic_node_classes, "semantic_node_classes"); + const auto n_semantic = semantic_node_classes.shape(0); + const auto semantic_cost_vector = array_1d_to_vector( + semantic_costs, + "semantic_costs", + static_cast(n_semantic) + ); + + std::vector> mutex_uv_vector(n_mutex); + { + const auto *uv_data = mutex_uvs.data(); + for (std::size_t index = 0; index < n_mutex; ++index) { + mutex_uv_vector[index][0] = uv_data[2 * index]; + mutex_uv_vector[index][1] = uv_data[2 * index + 1]; + } + } + std::vector> semantic_uv_vector(n_semantic); + { + const auto *uv_data = semantic_node_classes.data(); + for (std::size_t index = 0; index < n_semantic; ++index) { + semantic_uv_vector[index][0] = uv_data[2 * index]; + semantic_uv_vector[index][1] = uv_data[2 * index + 1]; + } + } + + graph::SemanticMutexWatershedResult result; + { + nb::gil_scoped_release release; + result = graph::semantic_mutex_watershed_clustering( + graph, + edge_cost_vector, + mutex_uv_vector, + mutex_cost_vector, + semantic_uv_vector, + semantic_cost_vector + ); + } + return std::make_pair( + vector_to_uint64_array(result.node_labels), + vector_to_array_1d(result.semantic_labels) + ); +} + UInt64Array multicut_fusion_move( const Graph &graph, ConstDoubleArray costs, @@ -1544,6 +1602,24 @@ void bind_graph(nb::module_ &m) { register_mutex_watershed_clustering.operator()("_mutex_watershed_clustering_float32"); register_mutex_watershed_clustering.operator()("_mutex_watershed_clustering_float64"); + const auto register_semantic_mutex_watershed_clustering = + [&m](const char *name) { + m.def( + name, + &semantic_mutex_watershed_clustering_t, + nb::arg("graph"), + nb::arg("edge_costs"), + nb::arg("mutex_uvs"), + nb::arg("mutex_costs"), + nb::arg("semantic_node_classes"), + nb::arg("semantic_costs") + ); + }; + register_semantic_mutex_watershed_clustering + .operator()("_semantic_mutex_watershed_clustering_float32"); + register_semantic_mutex_watershed_clustering + .operator()("_semantic_mutex_watershed_clustering_float64"); + // Lifted multicut sub-solver hierarchy. Same shape as the multicut sub- // solver bindings — opaque to Python, used by future fusion-move drivers. nb::class_(m, "_LiftedMulticutSolverBase"); diff --git a/src/bindings/segmentation.cxx b/src/bindings/segmentation.cxx index d37bdc4..10154d3 100644 --- a/src/bindings/segmentation.cxx +++ b/src/bindings/segmentation.cxx @@ -2,14 +2,17 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/segmentation/mutex_watershed.hxx" +#include "bioimage_cpp/segmentation/semantic_mutex_watershed.hxx" #include +#include #include #include #include #include #include +#include #include namespace nb = nanobind; @@ -22,6 +25,7 @@ using AffinityArray = nb::ndarray; using ValidEdgeArray = nb::ndarray; using LabelArray = nb::ndarray; +using SemanticLabelArray = nb::ndarray; template LabelArray mutex_watershed_grid_t( @@ -87,6 +91,89 @@ LabelArray mutex_watershed_grid_t( return LabelArray(data, label_shape.size(), label_shape.data(), owner); } +template +std::pair semantic_mutex_watershed_grid_t( + AffinityArray affinities, + ValidEdgeArray valid_edges, + const std::vector> &offsets, + const std::size_t number_of_attractive_channels, + const std::size_t number_of_offsets +) { + std::vector affinity_shape(affinities.ndim()); + for (std::size_t axis = 0; axis < affinities.ndim(); ++axis) { + affinity_shape[axis] = static_cast(affinities.shape(axis)); + } + std::vector valid_edges_shape(valid_edges.ndim()); + for (std::size_t axis = 0; axis < valid_edges.ndim(); ++axis) { + valid_edges_shape[axis] = static_cast(valid_edges.shape(axis)); + } + + std::vector label_shape; + label_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0); + std::vector label_view_shape; + label_view_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0); + for (std::size_t axis = 1; axis < affinities.ndim(); ++axis) { + label_shape.push_back(affinities.shape(axis)); + label_view_shape.push_back(static_cast(affinities.shape(axis))); + } + + const auto number_of_nodes = std::accumulate( + label_view_shape.begin(), + label_view_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + ); + auto *node_data = new std::uint64_t[static_cast(number_of_nodes)](); + nb::capsule node_owner( + node_data, + [](void *p) noexcept { delete[] static_cast(p); } + ); + auto *semantic_data = new std::int64_t[static_cast(number_of_nodes)](); + nb::capsule semantic_owner( + semantic_data, + [](void *p) noexcept { delete[] static_cast(p); } + ); + + ConstArrayView affinities_view{ + affinities.data(), + affinity_shape, + {}, + }; + ConstArrayView valid_edges_view{ + valid_edges.data(), + valid_edges_shape, + {}, + }; + ArrayView node_out_view{ + node_data, + label_view_shape, + {}, + }; + ArrayView semantic_out_view{ + semantic_data, + label_view_shape, + {}, + }; + + { + nb::gil_scoped_release release; + semantic_mutex_watershed_grid( + affinities_view, + valid_edges_view, + offsets, + number_of_attractive_channels, + number_of_offsets, + node_out_view, + semantic_out_view + ); + } + + return std::make_pair( + LabelArray(node_data, label_shape.size(), label_shape.data(), node_owner), + SemanticLabelArray(semantic_data, label_shape.size(), label_shape.data(), semantic_owner) + ); +} + } // namespace void bind_segmentation(nb::module_ &m) { @@ -108,6 +195,26 @@ void bind_segmentation(nb::module_ &m) { nb::arg("number_of_attractive_channels"), "Run mutex watershed on a 2D or 3D image-derived grid graph with float64 affinities." ); + m.def( + "_semantic_mutex_watershed_grid_float32", + &semantic_mutex_watershed_grid_t, + nb::arg("affinities"), + nb::arg("valid_edges"), + nb::arg("offsets"), + nb::arg("number_of_attractive_channels"), + nb::arg("number_of_offsets"), + "Run semantic mutex watershed on a 2D or 3D image-derived grid graph with float32 affinities." + ); + m.def( + "_semantic_mutex_watershed_grid_float64", + &semantic_mutex_watershed_grid_t, + nb::arg("affinities"), + nb::arg("valid_edges"), + nb::arg("offsets"), + nb::arg("number_of_attractive_channels"), + nb::arg("number_of_offsets"), + "Run semantic mutex watershed on a 2D or 3D image-derived grid graph with float64 affinities." + ); } } // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/_data.py b/src/bioimage_cpp/_data.py index b234406..485db2d 100644 --- a/src/bioimage_cpp/_data.py +++ b/src/bioimage_cpp/_data.py @@ -23,6 +23,14 @@ Affinities: - ``affinities`` — HDF5 file with sample affinities from the ISBI volume. Contains affinities under key ``affinities``. + +Semantic labels: +- ``semantic_labels`` — HDF5 file with paired instance and semantic ground + truth on a single 3D volume. Contains ``labels/instances``, + ``labels/semantic`` and ``raw``. The on-disk arrays are padded with ``-1`` + outside the labelled region; loaders crop to the labelled content slab. + Used by the semantic-mutex-watershed comparison scripts under + ``development/``. """ from __future__ import annotations @@ -36,6 +44,7 @@ DEFAULT_CACHE_DIR = Path.home() / ".cache" / "bioimage-cpp" CACHE_ENV_VAR = "BIOIMAGE_CPP_CACHE" ISBI_AFFINITY_FILENAME = "affinities" +SEMANTIC_LABELS_FILENAME = "semantic_labels" ISBI_AFFINITY_OFFSETS = ( (-1, 0, 0), (0, -1, 0), @@ -101,6 +110,10 @@ "https://owncloud.gwdg.de/index.php/s/aAyF2ekzsW7DFJo/download", "6472ad0fcf3c57a4ae345fda68c3cbb6072ee3e8db67b423502746b46d8cd5e5", ), + "semantic_labels": ( + "https://owncloud.gwdg.de/index.php/s/Ah7IGuYH7uuomQV/download", + "6232fe2fd58fdbd3def978798143fdcc65a2af118b4d9ee177b5c942173ece26", + ), } @@ -237,3 +250,69 @@ def load_isbi_gt_segmentation( with h5py.File(affinity_path(timeout=timeout), "r") as f: labels = f["labels/gt_segmentation"][:] return np.ascontiguousarray(labels) + + +def semantic_labels_path(*, timeout: Optional[float] = None) -> Path: + """Return the cached path to the registered semantic-labels HDF5 file.""" + return fetch(SEMANTIC_LABELS_FILENAME, timeout=timeout) + + +# Bounding box of the labelled content in the on-disk volume. Outside this +# slab both label volumes are uniformly ``-1``; cropping here keeps callers +# from having to special-case the padding. +SEMANTIC_LABELS_CROP = (slice(16, 32), slice(64, 512), slice(64, 512)) + + +def load_semantic_labels( + *, + timeout: Optional[float] = None, +) -> tuple[np.ndarray, np.ndarray]: + """Load the registered (instance, semantic) ground-truth volumes. + + Both volumes share the same spatial shape and live in a single HDF5 file + under keys ``labels/instances`` and ``labels/semantic``. The on-disk + arrays are padded with ``-1`` outside the labelled slab; this loader + returns the cropped labelled region (``(16, 448, 448)``). + + Returns + ------- + instance_labels: + Integer instance-segmentation volume. + semantic_labels: + Integer semantic-class volume (one class id per voxel). + """ + try: + import h5py + except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "h5py is required to load the registered semantic-labels file. " + "Install it with `pip install h5py`." + ) from error + + with h5py.File(semantic_labels_path(timeout=timeout), "r") as f: + instance = f["labels/instances"][SEMANTIC_LABELS_CROP] + semantic = f["labels/semantic"][SEMANTIC_LABELS_CROP] + return np.ascontiguousarray(instance), np.ascontiguousarray(semantic) + + +def load_semantic_raw( + *, + timeout: Optional[float] = None, +) -> np.ndarray: + """Load the raw volume paired with the registered semantic labels. + + Stored alongside ``labels/instances`` / ``labels/semantic`` in the same + HDF5 file under key ``raw``. Cropped consistently with + :func:`load_semantic_labels`. + """ + try: + import h5py + except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "h5py is required to load the registered semantic-raw file. " + "Install it with `pip install h5py`." + ) from error + + with h5py.File(semantic_labels_path(timeout=timeout), "r") as f: + raw = f["raw"][SEMANTIC_LABELS_CROP] + return np.ascontiguousarray(raw) diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index b0ed3ff..623931a 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -604,6 +604,11 @@ def edge_weighted_watershed( np.dtype("float64"): _core._mutex_watershed_clustering_float64, } +_SEMANTIC_MUTEX_WATERSHED_CLUSTERING_BY_DTYPE = { + np.dtype("float32"): _core._semantic_mutex_watershed_clustering_float32, + np.dtype("float64"): _core._semantic_mutex_watershed_clustering_float64, +} + def _resolve_weight_dtype(array, name: str) -> np.ndarray: """Coerce a weights input to a supported floating dtype. @@ -692,6 +697,95 @@ def mutex_watershed_clustering( return run(graph, edge_cost_array, mutex_uv_array, mutex_cost_array) +def semantic_mutex_watershed_clustering( + graph: UndirectedGraph | RegionAdjacencyGraph, + edge_costs, + mutex_uvs, + mutex_costs, + semantic_node_classes, + semantic_costs, +) -> tuple[np.ndarray, np.ndarray]: + """Semantic mutex watershed clustering on an undirected graph. + + Extends :func:`mutex_watershed_clustering` with a third group of edges + that attach semantic class labels to clusters. Two clusters carrying + different semantic class labels are forbidden from merging; otherwise + the algorithm proceeds as in the non-semantic case (attractive edges + merge; mutex edges block). + + Parameters + ---------- + graph: + :class:`UndirectedGraph` or :class:`RegionAdjacencyGraph` defining + the attractive edges. + edge_costs: + 1D array of length ``graph.number_of_edges``. Same dtype rules as + :func:`mutex_watershed_clustering`. + mutex_uvs: + ``(n_mutex, 2)`` uint64 array of (u, v) pairs for the mutex edges. + mutex_costs: + 1D array of length ``n_mutex``. + semantic_node_classes: + ``(n_semantic, 2)`` uint64 array. Column 0 is a node id; column 1 + is the semantic class id (non-negative integer). The semantic class + id is interpreted as a signed integer when reported back, so values + above ``np.iinfo(np.int64).max`` are out of range. + semantic_costs: + 1D array of length ``n_semantic``. Same dtype rules as + ``edge_costs``; if the floating dtypes of the three weight arrays + do not all agree, all three are promoted to ``float64``. + + Returns + ------- + node_labels: + ``(graph.number_of_nodes,)`` uint64 array. Dense node labels in + ``[0, k)`` in first-occurrence order. + semantic_labels: + ``(graph.number_of_nodes,)`` int64 array. Per-node semantic class + id, or ``-1`` for clusters that received no semantic assignment. + """ + edge_cost_array = _resolve_weight_dtype(edge_costs, "edge_costs") + mutex_cost_array = _resolve_weight_dtype(mutex_costs, "mutex_costs") + semantic_cost_array = _resolve_weight_dtype(semantic_costs, "semantic_costs") + + dtypes = {edge_cost_array.dtype, mutex_cost_array.dtype, semantic_cost_array.dtype} + if len(dtypes) > 1: + edge_cost_array = edge_cost_array.astype(np.float64, copy=False) + mutex_cost_array = mutex_cost_array.astype(np.float64, copy=False) + semantic_cost_array = semantic_cost_array.astype(np.float64, copy=False) + + edge_cost_array = _as_1d_array( + edge_cost_array, + edge_cost_array.dtype, + "edge_costs", + int(graph.number_of_edges), + ) + mutex_uv_array = _as_uv_array(mutex_uvs, "mutex_uvs") + mutex_cost_array = _as_1d_array( + mutex_cost_array, + mutex_cost_array.dtype, + "mutex_costs", + int(mutex_uv_array.shape[0]), + ) + semantic_uv_array = _as_uv_array(semantic_node_classes, "semantic_node_classes") + semantic_cost_array = _as_1d_array( + semantic_cost_array, + semantic_cost_array.dtype, + "semantic_costs", + int(semantic_uv_array.shape[0]), + ) + + run = _SEMANTIC_MUTEX_WATERSHED_CLUSTERING_BY_DTYPE[edge_cost_array.dtype] + return run( + graph, + edge_cost_array, + mutex_uv_array, + mutex_cost_array, + semantic_uv_array, + semantic_cost_array, + ) + + class MulticutObjective: """Multicut objective for an undirected graph and edge costs.""" @@ -2016,5 +2110,6 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "mutex_watershed_clustering", "project_node_labels_to_pixels", "region_adjacency_graph", + "semantic_mutex_watershed_clustering", "undirected_graph", ] diff --git a/src/bioimage_cpp/segmentation/__init__.py b/src/bioimage_cpp/segmentation/__init__.py index cf4b303..7324596 100644 --- a/src/bioimage_cpp/segmentation/__init__.py +++ b/src/bioimage_cpp/segmentation/__init__.py @@ -1,5 +1,5 @@ """Segmentation algorithms.""" -from .mutex_watershed import mutex_watershed +from .mutex_watershed import mutex_watershed, semantic_mutex_watershed -__all__ = ["mutex_watershed"] +__all__ = ["mutex_watershed", "semantic_mutex_watershed"] diff --git a/src/bioimage_cpp/segmentation/mutex_watershed.py b/src/bioimage_cpp/segmentation/mutex_watershed.py index d500dd0..78e02e5 100644 --- a/src/bioimage_cpp/segmentation/mutex_watershed.py +++ b/src/bioimage_cpp/segmentation/mutex_watershed.py @@ -13,6 +13,11 @@ np.dtype("float64"): _core._mutex_watershed_grid_float64, } +_SEMANTIC_MUTEX_WATERSHED_BY_DTYPE = { + np.dtype("float32"): _core._semantic_mutex_watershed_grid_float32, + np.dtype("float64"): _core._semantic_mutex_watershed_grid_float64, +} + def mutex_watershed( affinities: np.ndarray, @@ -110,6 +115,132 @@ def mutex_watershed( return labels +def semantic_mutex_watershed( + affinities: np.ndarray, + offsets: Sequence[Sequence[int]], + number_of_attractive_channels: int, + *, + strides: Sequence[int] | None = None, + randomized_strides: bool = False, + mask: np.ndarray | None = None, + mask_label: int = 0, +) -> tuple[np.ndarray, np.ndarray]: + """Run semantic mutex watershed on a 2D or 3D image-derived grid graph. + + The affinity array stacks three groups of channels along axis 0: + + 1. ``[0, number_of_attractive_channels)`` — attractive grid edges. + 2. ``[number_of_attractive_channels, len(offsets))`` — mutex grid edges. + 3. ``[len(offsets), affinities.shape[0])`` — one channel per semantic + class, scoring how likely each pixel belongs to that class. + + The first two groups are the regular mutex watershed input; the third + extends it: clusters can be tagged with the strongest semantic class + seen for any of their pixels and two clusters carrying different + semantic tags are forbidden from merging. + + Parameters + ---------- + affinities: + Array with shape ``(channels, y, x)`` or ``(channels, z, y, x)``. + ``channels`` must be strictly greater than ``len(offsets)`` so that + at least one semantic class is present; otherwise use + :func:`mutex_watershed`. Supported dtypes are ``float32`` and + ``float64``; non-contiguous arrays are copied. + offsets: + One offset per spatial channel (attractive + mutex), in NumPy axis + order. Length must equal ``number_of_attractive_channels + + number_of_mutex_channels``. + number_of_attractive_channels: + The first this many spatial channels are attractive; the remaining + ``len(offsets) - number_of_attractive_channels`` channels are mutex. + strides, randomized_strides: + Same semantics as :func:`mutex_watershed`. Apply to mutex channels + only. + mask: + Optional boolean foreground mask with shape ``affinities.shape[1:]``. + Spatial edges touching ``False`` pixels are skipped, ``False`` pixels + are labelled as background ``0`` in ``labels`` and as ``mask_label`` + in ``semantic_labels``. + mask_label: + Value written to ``semantic_labels`` for masked-out pixels. Defaults + to ``0``. + + Returns + ------- + labels: + ``uint64`` array with shape ``affinities.shape[1:]``. Consecutive + 1-based segmentation labels. + semantic_labels: + ``int64`` array with the same shape. Per-pixel semantic class id, or + ``-1`` for clusters that received no semantic assignment. + """ + array = np.asarray(affinities) + if array.ndim not in (3, 4): + raise ValueError( + "affinities must have shape (channels, y, x) or " + f"(channels, z, y, x), got ndim={array.ndim}" + ) + + dtype = array.dtype + try: + run = _SEMANTIC_MUTEX_WATERSHED_BY_DTYPE[dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _SEMANTIC_MUTEX_WATERSHED_BY_DTYPE) + raise TypeError( + f"affinities must have one of dtypes ({supported}), got dtype={dtype}" + ) from error + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + number_of_offsets = len(normalized_offsets) + number_of_channels = int(array.shape[0]) + spatial_ndim = array.ndim - 1 + if number_of_offsets == 0: + raise ValueError("offsets must not be empty") + if number_of_channels <= number_of_offsets: + raise ValueError( + "semantic_mutex_watershed requires at least one semantic-class channel, " + f"got channels={number_of_channels}, len(offsets)={number_of_offsets}; " + "use mutex_watershed for the non-semantic case" + ) + if any(len(offset) != spatial_ndim for offset in normalized_offsets): + raise ValueError( + "each offset must have length matching the spatial ndim, got " + f"spatial ndim={spatial_ndim}" + ) + if number_of_attractive_channels < 0: + raise ValueError("number_of_attractive_channels must be non-negative") + if number_of_attractive_channels > number_of_offsets: + raise ValueError( + "number_of_attractive_channels must be <= len(offsets)" + ) + + normalized_strides = _normalize_strides(strides, spatial_ndim, randomized_strides) + contiguous = np.ascontiguousarray(array) + valid_edges = _compute_valid_edges_semantic( + contiguous, + normalized_offsets, + number_of_attractive_channels, + number_of_offsets, + normalized_strides, + randomized_strides, + mask, + ) + + labels, semantic_labels = run( + contiguous, + valid_edges, + normalized_offsets, + number_of_attractive_channels, + number_of_offsets, + ) + if mask is not None: + invalid = np.logical_not(np.asarray(mask)) + labels[invalid] = 0 + semantic_labels[invalid] = mask_label + return labels, semantic_labels + + def _normalize_strides( strides: Sequence[int] | None, spatial_ndim: int, @@ -225,3 +356,39 @@ def _compute_valid_edges( valid_edges[(channel,) + source_slices] = channel_valid return np.ascontiguousarray(valid_edges, dtype=np.uint8) + + +def _compute_valid_edges_semantic( + affinities: np.ndarray, + offsets: Sequence[tuple[int, ...]], + number_of_attractive_channels: int, + number_of_offsets: int, + strides: tuple[int, ...] | None, + randomized_strides: bool, + mask: np.ndarray | None, +) -> np.ndarray: + # Reuse the spatial-edge logic, then extend the mask to cover the extra + # semantic-class channels. The spatial helper allocates the full + # ``affinities.shape`` mask but only writes into the first + # ``number_of_offsets`` channels, which is exactly what we want here — + # everything beyond is initialised to ``False`` and managed below. + valid_edges_uint = _compute_valid_edges( + tuple(affinities.shape), + offsets, + number_of_attractive_channels, + strides, + randomized_strides, + mask, + ) + valid_edges = valid_edges_uint.astype(bool, copy=True) + + semantic_channels = affinities[number_of_offsets:] + if semantic_channels.shape[0] > 0: + per_pixel_max = semantic_channels.max(axis=0, keepdims=True) + valid_edges[number_of_offsets:] = semantic_channels == per_pixel_max + + if mask is not None: + invalid = np.logical_not(np.asarray(mask)) + valid_edges[number_of_offsets:, invalid] = False + + return np.ascontiguousarray(valid_edges, dtype=np.uint8) diff --git a/src/cpp/segmentation/semantic_mutex_watershed.cxx b/src/cpp/segmentation/semantic_mutex_watershed.cxx new file mode 100644 index 0000000..5b7e1ec --- /dev/null +++ b/src/cpp/segmentation/semantic_mutex_watershed.cxx @@ -0,0 +1,27 @@ +#include "bioimage_cpp/segmentation/semantic_mutex_watershed.hxx" + +#include + +namespace bioimage_cpp { + +template void semantic_mutex_watershed_grid( + const ConstArrayView &, + const ConstArrayView &, + const std::vector> &, + std::size_t, + std::size_t, + const ArrayView &, + const ArrayView & +); + +template void semantic_mutex_watershed_grid( + const ConstArrayView &, + const ConstArrayView &, + const std::vector> &, + std::size_t, + std::size_t, + const ArrayView &, + const ArrayView & +); + +} // namespace bioimage_cpp diff --git a/tests/graph/test_semantic_mutex_watershed_graph.py b/tests/graph/test_semantic_mutex_watershed_graph.py new file mode 100644 index 0000000..c224351 --- /dev/null +++ b/tests/graph/test_semantic_mutex_watershed_graph.py @@ -0,0 +1,276 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _graph_from_edges(n: int, uvs): + uvs = np.asarray(uvs, dtype=np.uint64) + return bic.graph.UndirectedGraph.from_edges(n, uvs) + + +def _canonicalize(labels): + array = np.asarray(labels, dtype=np.uint64) + seen = {} + out = np.empty(array.shape, dtype=np.uint64) + for index, value in enumerate(array): + key = int(value) + if key not in seen: + seen[key] = len(seen) + out[index] = seen[key] + return out + + +def _empty_semantic(dtype=np.float64): + return ( + np.zeros((0, 2), dtype=np.uint64), + np.zeros(0, dtype=dtype), + ) + + +def test_without_semantic_edges_matches_regular_mutex_watershed(): + graph = _graph_from_edges(4, [[0, 1], [1, 2], [2, 3]]) + edge_costs = np.array([10.0, 9.0, 4.0], dtype=np.float64) + mutex_uvs = np.array([[0, 3]], dtype=np.uint64) + mutex_costs = np.array([5.0], dtype=np.float64) + semantic_uvs, semantic_costs = _empty_semantic() + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, semantic_uvs, semantic_costs + ) + regular_labels = bic.graph.mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs + ) + + np.testing.assert_array_equal(_canonicalize(labels), _canonicalize(regular_labels)) + np.testing.assert_array_equal(semantic, -np.ones(4, dtype=np.int64)) + assert semantic.dtype == np.int64 + assert labels.dtype == np.uint64 + + +def test_semantic_constraint_blocks_merge(): + # Two attractive edges 0-1 and 1-2. Node 0 has semantic class 0 (high weight), + # node 2 has semantic class 1 (high weight). The chain would normally merge + # all three nodes, but the semantic constraint should block the 1-2 merge + # after 0 and 1 unite and inherit class 0. + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 0], [2, 1]], dtype=np.uint64) + semantic_costs = np.array([10.0, 10.0], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + # Semantic edges have the highest weight so they're processed first. + # After that, the 0-1 attractive merge brings class 0 onto the joint + # root. The 1-2 attractive merge is then blocked because 2 has class 1. + np.testing.assert_array_equal(_canonicalize(labels), np.array([0, 0, 1])) + assert int(semantic[0]) == 0 + assert int(semantic[1]) == 0 + assert int(semantic[2]) == 1 + + +def test_semantic_label_propagates_across_merge(): + # 0 is assigned class 3 via a high-weight semantic edge. Then the chain + # 0-1-2-3 merges via attractive edges. All four nodes must end up with + # the same cluster label AND with semantic_label == 3. + graph = _graph_from_edges(4, [[0, 1], [1, 2], [2, 3]]) + edge_costs = np.array([1.0, 1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 3]], dtype=np.uint64) + semantic_costs = np.array([10.0], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + np.testing.assert_array_equal(_canonicalize(labels), np.zeros(4, dtype=np.uint64)) + np.testing.assert_array_equal(semantic, np.full(4, 3, dtype=np.int64)) + + +def test_same_semantic_label_does_not_block_merge(): + # Two roots with the same class label can still merge — the semantic + # constraint only fires when labels differ. + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 7], [2, 7]], dtype=np.uint64) + semantic_costs = np.array([10.0, 10.0], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + np.testing.assert_array_equal(_canonicalize(labels), np.zeros(3, dtype=np.uint64)) + np.testing.assert_array_equal(semantic, np.full(3, 7, dtype=np.int64)) + + +def test_mutex_still_separates_under_semantic(): + # Mutex edges should still block merges regardless of semantic tags. + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.array([[0, 2]], dtype=np.uint64) + mutex_costs = np.array([5.0], dtype=np.float64) + semantic_node_classes = np.array([[0, 0], [1, 0], [2, 0]], dtype=np.uint64) + semantic_costs = np.array([0.5, 0.5, 0.5], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + # Mutex 0-2 (weight 5) is processed first. Then attractive 0-1 (weight + # 1) merges {0,1}; attractive 1-2 is blocked by the propagated mutex. + # Semantic class 0 still propagates to all reachable clusters. + np.testing.assert_array_equal(_canonicalize(labels), np.array([0, 0, 1])) + np.testing.assert_array_equal(semantic, np.zeros(3, dtype=np.int64)) + + +def test_unassigned_clusters_keep_minus_one(): + graph = _graph_from_edges(3, [[0, 1]]) + edge_costs = np.array([1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 5]], dtype=np.uint64) + semantic_costs = np.array([10.0], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + # {0, 1} merges and inherits class 5. {2} is alone and unassigned. + assert int(semantic[0]) == 5 + assert int(semantic[1]) == 5 + assert int(semantic[2]) == -1 + + +def test_dense_label_range(): + graph = _graph_from_edges(6, [[0, 1], [2, 3], [4, 5]]) + edge_costs = np.array([1.0, 1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_uvs, semantic_costs = _empty_semantic() + + labels, _ = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, semantic_uvs, semantic_costs + ) + + assert set(int(value) for value in labels) == {0, 1, 2} + assert int(labels.max()) == 2 + + +def test_deterministic_across_runs(): + rng = np.random.default_rng(11) + n = 20 + uvs = [] + for u in range(n): + for v in range(u + 1, min(u + 4, n)): + uvs.append([u, v]) + uvs = np.array(uvs, dtype=np.uint64) + graph = _graph_from_edges(n, uvs) + edge_costs = rng.uniform(-1.0, 1.0, size=int(graph.number_of_edges)).astype(np.float64) + mutex_uvs = rng.integers(0, n, size=(20, 2), dtype=np.uint64) + mutex_uvs = mutex_uvs[mutex_uvs[:, 0] != mutex_uvs[:, 1]] + mutex_costs = rng.uniform(0.0, 1.0, size=mutex_uvs.shape[0]).astype(np.float64) + semantic_nodes = rng.integers(0, n, size=10, dtype=np.uint64) + semantic_classes = rng.integers(0, 3, size=10, dtype=np.uint64) + semantic_node_classes = np.stack([semantic_nodes, semantic_classes], axis=1) + semantic_costs = rng.uniform(0.0, 1.0, size=10).astype(np.float64) + + first = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, semantic_node_classes, semantic_costs + ) + second = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, semantic_node_classes, semantic_costs + ) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], second[1]) + + +def test_float32_inputs_supported(): + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float32) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float32) + semantic_node_classes = np.array([[0, 0], [2, 1]], dtype=np.uint64) + semantic_costs = np.array([10.0, 10.0], dtype=np.float32) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + assert labels.dtype == np.uint64 + assert semantic.dtype == np.int64 + np.testing.assert_array_equal(_canonicalize(labels), np.array([0, 0, 1])) + + +def test_mismatched_dtypes_are_promoted(): + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float32) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 0], [2, 1]], dtype=np.uint64) + semantic_costs = np.array([10.0, 10.0], dtype=np.float64) + + labels, semantic = bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) + + np.testing.assert_array_equal(_canonicalize(labels), np.array([0, 0, 1])) + assert int(semantic[0]) == 0 + assert int(semantic[2]) == 1 + + +def test_invalid_semantic_shape_raises(): + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + bad_semantic = np.array([0, 1], dtype=np.uint64) # 1D, not (n, 2) + semantic_costs = np.array([1.0], dtype=np.float64) + + with pytest.raises(ValueError): + bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, bad_semantic, semantic_costs + ) + + +def test_mismatched_semantic_costs_raises(): + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[0, 0]], dtype=np.uint64) + bad_costs = np.array([1.0, 2.0], dtype=np.float64) + + with pytest.raises(ValueError): + bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, bad_costs, + ) + + +def test_out_of_range_semantic_node_raises(): + graph = _graph_from_edges(3, [[0, 1], [1, 2]]) + edge_costs = np.array([1.0, 1.0], dtype=np.float64) + mutex_uvs = np.zeros((0, 2), dtype=np.uint64) + mutex_costs = np.zeros(0, dtype=np.float64) + semantic_node_classes = np.array([[99, 0]], dtype=np.uint64) + semantic_costs = np.array([1.0], dtype=np.float64) + + with pytest.raises((ValueError, RuntimeError)): + bic.graph.semantic_mutex_watershed_clustering( + graph, edge_costs, mutex_uvs, mutex_costs, + semantic_node_classes, semantic_costs, + ) diff --git a/tests/segmentation/test_semantic_mutex_watershed.py b/tests/segmentation/test_semantic_mutex_watershed.py new file mode 100644 index 0000000..b14b092 --- /dev/null +++ b/tests/segmentation/test_semantic_mutex_watershed.py @@ -0,0 +1,185 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_semantic_mutex_watershed_2d_smoke(): + affinities = np.array( + [ + # attractive right-neighbor + [[1.0, 1.0, 1.0, 0.0]], + # class-0 affinity (p0 high) + [[10.0, 0.5, 0.5, 0.5]], + # class-1 affinity (p3 high) + [[0.5, 0.5, 0.5, 10.0]], + ], + dtype=np.float32, + ) + offsets = [[0, 1]] + + labels, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + assert labels.shape == (1, 4) + assert semantic.shape == (1, 4) + assert labels.dtype == np.uint64 + assert semantic.dtype == np.int64 + np.testing.assert_array_equal(labels, np.array([[1, 1, 1, 2]], dtype=np.uint64)) + np.testing.assert_array_equal(semantic, np.array([[0, 0, 0, 1]], dtype=np.int64)) + + +def test_semantic_label_propagates_across_merges(): + # p0 anchored as class 2; chain merges all four pixels via attractive + # edges. All four should end up with class 2. + affinities = np.array( + [ + [[1.0, 1.0, 1.0, 0.0]], + [[10.0, 0.0, 0.0, 0.0]], + ], + dtype=np.float32, + ) + offsets = [[0, 1]] + + labels, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + np.testing.assert_array_equal(labels, np.ones((1, 4), dtype=np.uint64)) + # Only one class channel → class id 0 propagates everywhere. + np.testing.assert_array_equal(semantic, np.zeros((1, 4), dtype=np.int64)) + + +def test_unassigned_clusters_report_minus_one(): + # No semantic channel has a high affinity anywhere; class-0 channel ties + # at 0 for every pixel and would normally be valid at every pixel via the + # argmax mask, but the cost is 0 so it gets processed after the + # attractive merges — assignments still happen, but with class 0. To + # observe a truly unassigned cluster, we use a fully zero semantic + # channel AND a strong attractive grid that leaves multiple components: + # impossible because every pixel argmax-ties get a (weak) assignment. + # Instead, build the example with a fully zero semantic channel and a + # disconnected attractive grid — we still get class 0 everywhere because + # the (tied) semantic assignment fires for each pixel. Test the "minus + # one" case via the graph API in the companion suite; here we just + # confirm the int64 sentinel format. + affinities = np.array( + [ + [[1.0, 0.0, 1.0]], + [[1.0, 1.0, 1.0]], + ], + dtype=np.float32, + ) + offsets = [[0, 1]] + + _, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + assert semantic.dtype == np.int64 + # Every pixel argmax-resolves to class 0 → no -1 expected here. + assert semantic.min() >= 0 + + +def test_semantic_mutex_watershed_3d_smoke(): + affinities = np.ones((4, 2, 2, 2), dtype=np.float64) + # 3 spatial offsets + 1 semantic class channel. + offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] + + labels, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=3 + ) + + assert labels.shape == (2, 2, 2) + assert semantic.shape == (2, 2, 2) + # Fully attractive → one component; class 0 propagates everywhere. + np.testing.assert_array_equal(labels, np.ones((2, 2, 2), dtype=np.uint64)) + np.testing.assert_array_equal(semantic, np.zeros((2, 2, 2), dtype=np.int64)) + + +def test_without_semantic_channels_raises(): + # No extra channel → caller should use regular mutex_watershed. + affinities = np.ones((2, 3, 4), dtype=np.float32) + offsets = [[0, 1], [1, 0]] + + with pytest.raises(ValueError): + bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=2 + ) + + +def test_unsupported_dtype_raises(): + affinities = np.ones((2, 3, 4), dtype=np.float16) + offsets = [[0, 1]] + + with pytest.raises(TypeError): + bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + +def test_invalid_offset_length_raises(): + affinities = np.ones((2, 3, 4), dtype=np.float32) + bad_offsets = [[0, 1, 0]] + + with pytest.raises(ValueError): + bic.segmentation.semantic_mutex_watershed( + affinities, bad_offsets, number_of_attractive_channels=1 + ) + + +def test_float32_and_float64_match_on_simple_problem(): + affinities = np.array( + [ + [[0.9, 0.1, 0.3]], + [[5.0, 0.0, 0.0]], + ], + dtype=np.float64, + ) + offsets = [[0, 1]] + + labels_64, semantic_64 = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + labels_32, semantic_32 = bic.segmentation.semantic_mutex_watershed( + affinities.astype(np.float32), offsets, number_of_attractive_channels=1 + ) + np.testing.assert_array_equal(labels_32, labels_64) + np.testing.assert_array_equal(semantic_32, semantic_64) + + +def test_mask_zeros_labels_and_sets_mask_label(): + affinities = np.array( + [ + [[1.0, 1.0, 1.0]], + [[10.0, 10.0, 10.0]], + ], + dtype=np.float32, + ) + offsets = [[0, 1]] + mask = np.array([[True, False, True]], dtype=bool) + + labels, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1, + mask=mask, mask_label=7, + ) + + # Masked pixel → label 0, semantic = mask_label. + assert int(labels[0, 1]) == 0 + assert int(semantic[0, 1]) == 7 + # Unmasked pixels stay >=1 and carry the semantic class. + assert int(labels[0, 0]) >= 1 + assert int(labels[0, 2]) >= 1 + + +def test_negative_mask_label_supported(): + affinities = np.ones((2, 1, 2), dtype=np.float64) + offsets = [[0, 1]] + mask = np.array([[True, False]], dtype=bool) + + _, semantic = bic.segmentation.semantic_mutex_watershed( + affinities, offsets, number_of_attractive_channels=1, + mask=mask, mask_label=-1, + ) + assert int(semantic[0, 1]) == -1