Add watershed from affinities

Author: constantinpape

MIGRATION_GUIDE.md

@@ -841,6 +841,44 @@ Notes:
 - Signed integer inputs must not contain negative labels.
 - Inputs are converted to contiguous `uint64` arrays before entering C++.
 
+## Marker-controlled Watershed
+
+`bioimage-cpp` ships two marker-controlled watershed entry points: one that
+consumes a node-valued heightmap and one that consumes an edge-valued
+nearest-neighbour affinity map. Both share the same flooding scaffolding
+internally — a 65536-bucket Meyer-style monotone queue — and have the same
+public-facing semantics: markers are mandatory, connectivity is 1
+(4-neighbour in 2D, 6-neighbour in 3D), an optional foreground mask is
+supported, and tie-breaking on equal heights / equal affinities is
+unspecified.
+
+```python
+# Heightmap-driven (analogous to skimage.segmentation.watershed)
+labels = bic.segmentation.watershed(image, markers, mask=optional_mask)
+
+# Affinity-driven: edge priorities, no heightmap derivation needed
+labels = bic.segmentation.watershed_from_affinities(
+    affinities,                       # (C, *spatial), C == spatial_ndim
+    offsets=[(-1, 0), (0, -1)],       # one NN offset per channel, same sign
+    markers=markers,
+    mask=optional_mask,
+)
+```
+
+Notes for `watershed_from_affinities`:
+
+- Each channel must encode a single nearest-neighbour edge (exactly one
+  ±1 entry, the rest zero). All offsets must have the same sign — mixing
+  positive and negative directions is rejected. The function dispatches
+  to a positive-direction or negative-direction specialisation at the C++
+  layer so the inner loop has no per-channel sign branches.
+- Offsets may be passed in any axis order; the channel ↔ axis mapping is
+  rebuilt internally.
+- Higher affinity is processed first (high affinity = strong bond).
+- Compared to `affogato.segmentation.compute_mws_segmentation`, there are
+  no mutex (repulsive) channels and no long-range offsets — use
+  `bic.segmentation.mutex_watershed` for that.
+
 ## Mutex Watershed
 
 `bioimage-cpp` ships two mutex-watershed entry points, mirroring the two

development/segmentation/_watershed_from_affinities_equivalence.py

@@ -0,0 +1,225 @@
+"""Shared logic for the affinity-watershed comparison scripts.
+
+Runs ``bioimage_cpp.segmentation.watershed_from_affinities`` directly on the
+nearest-neighbour ISBI affinity channels and compares against
+``bioimage_cpp.segmentation.watershed`` running on the heightmap derived
+from those same channels (``1 − mean(NN affinities)``). Both algorithms see
+identical markers (local minima of the smoothed heightmap). Partition
+agreement is reported but expected to be partial — the two algorithms have
+deliberately different priority semantics (edge-keyed vs node-keyed) and the
+purpose of this script is to surface that difference, not to gate it.
+"""
+
+from __future__ import annotations
+
+import argparse
+from statistics import median
+from time import perf_counter
+from typing import Callable
+
+import numpy as np
+
+# Reuse the existing equivalence helpers — heightmap + marker generation are
+# identical to the node-watershed comparison.
+from _watershed_equivalence import (
+    _load_validation_metrics,
+    load_problem,
+    make_heightmap,
+    make_markers,
+)
+
+
+def _nearest_neighbour_channels(offsets):
+    return [
+        i
+        for i, offset in enumerate(offsets)
+        if sum(1 for v in offset if v != 0) == 1
+        and all(abs(v) <= 1 for v in offset)
+    ]
+
+
+def prepare_2d_problem(
+    affinities: np.ndarray,
+    offsets: list[tuple[int, ...]],
+    *,
+    z: int,
+    yx_shape: tuple[int, int],
+    smoothing_sigma: float,
+) -> tuple[np.ndarray, list[tuple[int, ...]], np.ndarray, np.ndarray]:
+    """Return (nn_affinities_2d, nn_offsets_2d, heightmap, markers)."""
+    channels_2d = [i for i, offset in enumerate(offsets) if offset[0] == 0]
+    offsets_2d = [offsets[i][1:] for i in channels_2d]
+    # Keep only nearest-neighbour channels for the affinity watershed input.
+    nn_idx_within_2d = _nearest_neighbour_channels(offsets_2d)
+    nn_channels_in_full = [channels_2d[i] for i in nn_idx_within_2d]
+    nn_offsets_2d = [offsets_2d[i] for i in nn_idx_within_2d]
+    y, x = yx_shape
+
+    nn_affinities = np.ascontiguousarray(
+        affinities[nn_channels_in_full, z, :y, :x]
+    )
+
+    # Reuse the existing heightmap/marker pipeline so the markers are
+    # identical to the node-watershed scripts'.
+    full_2d_affs = np.ascontiguousarray(affinities[channels_2d, z, :y, :x])
+    heightmap = make_heightmap(full_2d_affs, offsets_2d)
+    markers = make_markers(heightmap, smoothing_sigma=smoothing_sigma)
+    return nn_affinities, nn_offsets_2d, heightmap, markers
+
+
+def prepare_3d_problem(
+    affinities: np.ndarray,
+    offsets: list[tuple[int, ...]],
+    *,
+    zyx_shape: tuple[int, int, int],
+    smoothing_sigma: float,
+) -> tuple[np.ndarray, list[tuple[int, ...]], np.ndarray, np.ndarray]:
+    z, y, x = zyx_shape
+    nn_channels = _nearest_neighbour_channels(offsets)
+    nn_affinities = np.ascontiguousarray(
+        affinities[nn_channels, :z, :y, :x]
+    )
+    nn_offsets = [offsets[i] for i in nn_channels]
+    full_3d_affs = np.ascontiguousarray(affinities[:, :z, :y, :x])
+    heightmap = make_heightmap(full_3d_affs, offsets)
+    markers = make_markers(heightmap, smoothing_sigma=smoothing_sigma)
+    return nn_affinities, nn_offsets, heightmap, markers
+
+
+def run_from_affinities(
+    nn_affinities: np.ndarray,
+    nn_offsets: list[tuple[int, ...]],
+    markers: np.ndarray,
+) -> np.ndarray:
+    import bioimage_cpp as bic
+
+    return bic.segmentation.watershed_from_affinities(
+        nn_affinities, nn_offsets, markers,
+    )
+
+
+def run_node_watershed(heightmap: np.ndarray, markers: np.ndarray) -> np.ndarray:
+    import bioimage_cpp as bic
+
+    return bic.segmentation.watershed(heightmap, markers)
+
+
+def compare_segmentations(
+    candidate: np.ndarray,
+    reference: np.ndarray,
+) -> dict[str, float | str | bool]:
+    source, rand_index, variation_of_information = _load_validation_metrics()
+
+    vi_split, vi_merge = variation_of_information(candidate, reference)
+    adapted_rand_error, ri = rand_index(candidate, reference)
+    exact_equal = bool(np.array_equal(candidate, reference))
+    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_equal,
+    }
+
+
+def time_runs(
+    fn: Callable[[], np.ndarray],
+    repeats: int,
+) -> tuple[list[float], np.ndarray]:
+    fn()  # warmup
+    timings: list[float] = []
+    last_result: np.ndarray | None = None
+    for _ in range(repeats):
+        start = perf_counter()
+        last_result = fn()
+        timings.append(perf_counter() - start)
+    assert last_result is not None
+    return timings, last_result
+
+
+def print_report(
+    *,
+    ndim: int,
+    nn_affinities: np.ndarray,
+    markers: np.ndarray,
+    metrics: dict[str, float | str | bool],
+    aff_timings: list[float],
+    node_timings: list[float],
+):
+    aff_median = median(aff_timings)
+    node_median = median(node_timings)
+    ratio = node_median / aff_median if aff_median > 0 else float("inf")
+    n_markers = int(markers.max())
+
+    print(f"Watershed-from-affinities {ndim}D comparison")
+    print(
+        f"NN-affinity shape: {nn_affinities.shape}, dtype: {nn_affinities.dtype}"
+    )
+    print(f"markers: {n_markers} seeds")
+    print(f"validation metrics: {metrics['validation_source']}")
+    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']:.6g}"
+    )
+    print(f"watershed_from_affinities median runtime: {aff_median:.6f} s")
+    print(f"watershed (node-based) median runtime: {node_median:.6f} s")
+    print(f"node-based / affinity-based runtime ratio: {ratio:.3f}x")
+
+
+def run_check(
+    *,
+    ndim: int,
+    repeats: int,
+    z: int,
+    yx_shape: tuple[int, int],
+    zyx_shape: tuple[int, int, int],
+    smoothing_sigma: float,
+):
+    affinities, offsets = load_problem()
+    if ndim == 2:
+        nn_affinities, nn_offsets, heightmap, markers = prepare_2d_problem(
+            affinities, offsets, z=z, yx_shape=yx_shape,
+            smoothing_sigma=smoothing_sigma,
+        )
+    elif ndim == 3:
+        nn_affinities, nn_offsets, heightmap, markers = prepare_3d_problem(
+            affinities, offsets, zyx_shape=zyx_shape,
+            smoothing_sigma=smoothing_sigma,
+        )
+    else:
+        raise ValueError(f"ndim must be 2 or 3, got {ndim}")
+
+    aff_timings, aff_labels = time_runs(
+        lambda: run_from_affinities(nn_affinities, nn_offsets, markers),
+        repeats,
+    )
+    node_timings, node_labels = time_runs(
+        lambda: run_node_watershed(heightmap, markers),
+        repeats,
+    )
+
+    metrics = compare_segmentations(aff_labels, node_labels)
+    print_report(
+        ndim=ndim,
+        nn_affinities=nn_affinities,
+        markers=markers,
+        metrics=metrics,
+        aff_timings=aff_timings,
+        node_timings=node_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(
+        "--smoothing-sigma", type=float, default=1.5,
+        help="Gaussian sigma applied to the heightmap before finding local minima.",
+    )

development/segmentation/check_watershed_from_affinities_2d.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import argparse
+
+from _watershed_from_affinities_equivalence import add_common_arguments, run_check
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Compare bioimage-cpp affinity-based watershed against the node-based watershed in 2D."
+    )
+    add_common_arguments(parser)
+    parser.add_argument(
+        "--z", type=int, default=0,
+        help="Z slice from the ISBI affinities used for the 2D check.",
+    )
+    parser.add_argument(
+        "--shape", type=int, nargs=2, metavar=("Y", "X"), default=(512, 512),
+        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),
+        smoothing_sigma=args.smoothing_sigma,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/segmentation/check_watershed_from_affinities_3d.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import argparse
+
+from _watershed_from_affinities_equivalence import add_common_arguments, run_check
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Compare bioimage-cpp affinity-based watershed against the node-based watershed in 3D."
+    )
+    add_common_arguments(parser)
+    parser.add_argument(
+        "--shape", type=int, nargs=3, metavar=("Z", "Y", "X"), default=(6, 512, 512),
+        help="Spatial crop shape for the 3D check.",
+    )
+    args = parser.parse_args()
+
+    run_check(
+        ndim=3,
+        repeats=args.repeats,
+        z=0,
+        yx_shape=(0, 0),
+        zyx_shape=tuple(args.shape),
+        smoothing_sigma=args.smoothing_sigma,
+    )
+
+
+if __name__ == "__main__":
+    main()