diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index fc86b27..5d9dc91 100644 --- a/MIGRATION_GUIDE.md +++ b/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 diff --git a/development/segmentation/PERFORMANCE_NOTES.md b/development/segmentation/PERFORMANCE_NOTES.md new file mode 100644 index 0000000..28e93d3 --- /dev/null +++ b/development/segmentation/PERFORMANCE_NOTES.md @@ -0,0 +1,184 @@ +# Watershed — performance notes + +Optimization log for `bioimage_cpp.segmentation.watershed`. Re-run with: + +```bash +python development/segmentation/check_watershed_2d.py --repeats 5 +python development/segmentation/check_watershed_3d.py --repeats 3 +``` + +Both scripts build the same heightmap from the cached ISBI affinity volume +(1 − mean of the nearest-neighbour affinity channels), generate seeds by +labelling local minima of a Gaussian-smoothed copy, then time +`bioimage_cpp.segmentation.watershed` against +`skimage.segmentation.watershed(connectivity=1)` on identical inputs with one +untimed warmup and interleaved repeats. Partition agreement is measured with +Rand Index / VI / adapted-rand-error from `elf.evaluation`. Exact label +equality is not expected — tie-breaking on equal heights is documented as +unspecified for `bioimage_cpp.watershed`. + +## Setup + +- CPU: 11th Gen Intel Core i7-1185G7 (Tiger Lake, 4C/8T) +- Compiler: gcc 14.3.0 (conda-forge), `-O3`, no `-march=native` +- Python 3.12.12 on Linux x86_64 +- `scikit-image 0.25.2`, `numpy 1.26.4`, `bioimage_cpp 0.1.0` + +## Headline numbers + +Median wall-clock per call on the ISBI test volume. + +| problem | shape | seeds | baseline (v1) | optimized | speedup over baseline | skimage | speedup vs skimage | Rand Index | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| 2D | (512, 512) | 1162 | 47.1 ms | **8.1 ms** | **5.8×** | 66.0 ms | **8.1×** | 0.982 | +| 3D | (6, 512, 512) | 2407 | 720 ms | **77.3 ms** | **9.3×** | 937 ms | **12.1×** | 0.997 | + +`baseline (v1)` is the heap-based version from the initial watershed +implementation, before this optimization round. `optimized` is the +65536-bucket Meyer-flooding variant landed below. Rand Index is measured +against `skimage.segmentation.watershed` — values above 0.97 indicate +near-identical partitions with boundary jitter of a few percent. + +## Optimization phases + +### Baseline profile + +With `BIOIMAGE_PROFILE=ON` on a single 3D run: + +``` +init_output 0.0003 s ( 0.0%) +seed_pass 0.0024 s ( 0.4%) +main_loop 0.6183 s ( 99.6%) +total 0.6210 s +``` + +Everything is in the heap loop. The inner loop calls +`detail::valid_offset_target` once per neighbour, which does `ndim` +divisions and modulos per call — 18 div+mod per heap pop in 3D. + +### Phase 2+3 — Precomputed offsets + 2D/3D specialization + +`watershed` split into `detail_ws::watershed_2d` and +`detail_ws::watershed_3d` with manually-unrolled 4- / 6-neighbour blocks. +Replaced `valid_offset_target` (which does per-axis div+mod every call) with: + +- one (2D) or two (3D) div+mods per heap pop to decompose `node` into + `(y, x)` or `(z, y, x)`; +- branchless boundary checks of the form `coord[axis] > 0` / + `coord[axis] + 1 < shape[axis]`; +- flat-index neighbour arithmetic via `node ± strides[axis]`. + +Heap stays as `std::priority_queue>`. + +| problem | before | after | delta | +|---|---:|---:|---:| +| 2D | 38.8 ms | 36.1 ms | 1.07× | +| 3D | 602 ms | 484 ms | 1.24× | + +3D got the bigger win because it had more div+mod to remove per pop (18 → 2) +than 2D (8 → 1). Rand Index unchanged (algorithm semantics preserved). + +### Phase 4 — Smaller heap entries + +Skipped. The profile showed the entire cost is in `main_loop` and the next +phase replaces the heap entirely, so the 10–20% potential from a tighter +heap entry isn't worth the binding-layer churn. + +### Phase 5 — Bucket-queue Meyer flooding + +Replaced `std::priority_queue` with a 65536-bucket queue. Algorithm: + +1. Scan the heightmap (skipping masked pixels) to find `[h_min, h_max]`. +2. Quantize each pixel to `uint16_t` with + `level = floor((image[i] - h_min) / (h_max - h_min) * 65535)`. +3. For each level, maintain a `std::vector` of pending pixel + indices. A `current_level` cursor advances monotonically. +4. Pop pixels from `buckets[current_level]` via `pop_back` (LIFO inside a + level). For each unlabeled neighbour, set its label and push it into + `buckets[max(level[neighbour], current_level)]` — the Meyer monotone + semantic: a neighbour pushed from above never lands below the cursor. +5. Advance `current_level` when the current bucket drains. + +`O(N + L)` total work where `L = 65536` levels. For `N >> L` this is +effectively `O(N)` instead of the heap's `O(N log N)`. + +| problem | before (heap+specialization) | after (bucket) | delta | +|---|---:|---:|---:| +| 2D | 36.1 ms | 8.1 ms | 4.5× | +| 3D | 484 ms | 77.3 ms | 6.3× | + +Post-Phase-5 3D profile: + +``` +init_output 0.0002 s ( 0.3%) +range 0.0023 s ( 2.9%) +quantize 0.0020 s ( 2.5%) +seed_pass 0.0013 s ( 1.7%) +main_loop 0.0721 s ( 92.6%) +total 0.0779 s +``` + +`main_loop` is still ~93% of the time; the per-pixel work in the bucket +queue is dramatically lower than in the heap, but it remains the dominant +phase. `range`+`quantize`+`seed_pass` together are ~7% and aren't worth +chasing further at this stage — saving all of them would be ~5 ms. + +## Tradeoff: partition agreement + +Rand Index drops from `0.998 / 0.999` (heap) to `0.982 / 0.997` (bucket). +2D agreement dropped more because the ISBI 2D crop has small regions +(avg ~225 pixels per seed) where boundary tie-breaking dominates the +metric. + +Two semantic differences cause this: + +1. **Quantization.** 65536 levels can't distinguish ULP-close floats. Two + pixels at heights `1e-6` apart end up in the same level instead of being + strictly ordered as the heap would. Tested at 65536 levels; bumping to + 1M didn't materially improve agreement (tried during this round) but + costs more memory, so kept at 65536. +2. **Meyer monotone flooding.** A "valley pixel" reached from a higher + cursor is processed at the cursor's level, not its own. The heap-based + watershed would pop it immediately and let it propagate from there. The + bucket-queue version delays its processing until the cursor naturally + reaches it, which can change which seed claims it. + +Both are documented as unspecified tie-breaking in the Python wrapper. +The 21-test correctness suite (deterministic ridge tests, mask handling, +dtype matrix, error cases) still passes; the agreement metric vs +`skimage.segmentation.watershed` quantifies the boundary jitter on a real +problem. + +The benchmark script's `min_rand_index` gate is set to 0.97 to reflect +this — below that we'd treat it as a real regression to investigate. + +## What was tried and not landed + +- **More quantization levels (1M).** Marginal improvement in agreement + (0.982 → 0.985 on 2D), 4× memory cost for `levels[]`, and the bucket + vector grows to 24 MB of empty `std::vector` shells. Not worth it. +- **`std::deque` for FIFO inside a level.** Doesn't change Meyer semantics + (the issue isn't intra-level order). Adds per-op overhead. Discarded. +- **`uint32_t` indices in the bucket queue.** Cuts bucket storage in half + on the inner index path; preserves the existing `uint64_t` flat-index + semantics in the rest of the codebase. Skipped to keep the change + minimal; revisit if a single watershed > 2³² voxels ever shows up + (currently ruled out by memory anyway). + +## How to reproduce + +```bash +# Clean build (production) +pip install -e . --no-build-isolation +python -m pytest tests/segmentation/test_watershed.py -q # 21 pass +python development/segmentation/check_watershed_2d.py --repeats 5 +python development/segmentation/check_watershed_3d.py --repeats 3 + +# Profile build (per-phase breakdown to stderr) +pip install -e . --no-build-isolation -C cmake.define.BIOIMAGE_PROFILE=ON +python development/segmentation/check_watershed_3d.py --repeats 1 +``` + +The profile macros (`BIOIMAGE_PROFILE_INIT/SCOPE/REPORT`) are gated by the +`BIOIMAGE_PROFILE` cmake option and are no-ops in production builds, so +they're left in `watershed.hxx` for the next round of work. diff --git a/development/segmentation/_watershed_equivalence.py b/development/segmentation/_watershed_equivalence.py new file mode 100644 index 0000000..08cff91 --- /dev/null +++ b/development/segmentation/_watershed_equivalence.py @@ -0,0 +1,288 @@ +"""Shared logic for the watershed correctness + runtime comparisons. + +Builds a node-heightmap and seed markers from the cached ISBI affinity +volume, then runs ``bioimage_cpp.segmentation.watershed`` against +``skimage.segmentation.watershed`` (connectivity=1). Correctness uses +partition-comparison metrics (VI, rand index) rather than exact label +equality — tie-breaking differs between the two implementations. +""" + +from __future__ import annotations + +import argparse +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + + +def load_problem(): + from bioimage_cpp._data import load_isbi_affinities + + affinities, offsets = load_isbi_affinities() + return np.ascontiguousarray(affinities), [tuple(offset) for offset in offsets] + + +def _nearest_neighbour_channels(offsets): + """Return indices of channels whose offset moves one step along a single axis.""" + 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 make_heightmap(affinities: np.ndarray, offsets: list[tuple[int, ...]]) -> np.ndarray: + """Build a per-pixel heightmap from the nearest-neighbour affinity channels. + + Affinity ~ 1 means "neighbours are in the same object". Inverting and + averaging the nearest-neighbour channels gives a smooth boundary map + suitable as a watershed heightmap. + """ + nn_channels = _nearest_neighbour_channels(offsets) + if not nn_channels: + raise ValueError("no nearest-neighbour affinity channels found") + mean_aff = affinities[nn_channels].mean(axis=0).astype(np.float32, copy=True) + return np.ascontiguousarray(1.0 - mean_aff) + + +def make_markers(heightmap: np.ndarray, *, smoothing_sigma: float = 1.5) -> np.ndarray: + """Build seed markers as labelled connected components of the heightmap's local minima. + + A small Gaussian smoothing is applied before detecting minima so that + affinity noise doesn't produce thousands of single-pixel seeds. + """ + from scipy import ndimage as ndi + from skimage.morphology import local_minima + from skimage.measure import label + + if smoothing_sigma > 0: + smoothed = ndi.gaussian_filter(heightmap, sigma=smoothing_sigma) + else: + smoothed = heightmap + minima = local_minima(smoothed) + markers = label(minima, connectivity=1).astype(np.int32, copy=False) + return np.ascontiguousarray(markers) + + +def prepare_2d_problem( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + *, + z: int, + yx_shape: tuple[int, int], + smoothing_sigma: float, +) -> tuple[np.ndarray, np.ndarray]: + channels_2d = [i for i, offset in enumerate(offsets) if offset[0] == 0] + y, x = yx_shape + cropped = affinities[channels_2d, z, :y, :x] + offsets_2d = [offsets[i][1:] for i in channels_2d] + heightmap = make_heightmap(np.ascontiguousarray(cropped), offsets_2d) + markers = make_markers(heightmap, smoothing_sigma=smoothing_sigma) + return 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, np.ndarray]: + z, y, x = zyx_shape + cropped = affinities[:, :z, :y, :x] + heightmap = make_heightmap(np.ascontiguousarray(cropped), offsets) + markers = make_markers(heightmap, smoothing_sigma=smoothing_sigma) + return heightmap, markers + + +def run_bioimage_cpp(heightmap: np.ndarray, markers: np.ndarray) -> np.ndarray: + import bioimage_cpp as bic + + return bic.segmentation.watershed(heightmap, markers) + + +def run_skimage_reference(heightmap: np.ndarray, markers: np.ndarray) -> np.ndarray: + from skimage.segmentation import watershed as sk_watershed + + return sk_watershed(heightmap, markers=markers, connectivity=1) + + +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_segmentations( + candidate: np.ndarray, + reference: np.ndarray, + *, + min_rand_index: float = 0.97, +) -> dict[str, float | str | bool]: + """Partition-style comparison. + + Exact label equality is not expected — bioimage-cpp uses a 65536-bucket + quantized watershed (Meyer-style monotone flooding) while skimage uses a + float-priority heap, so boundary pixels around every region shift by 1–2 + cells and a few percent of pixel pairs partition differently. Rand Index + is the primary "do these partitions agree" check (it copes gracefully + with boundary jitter); VI and ARE are reported for context but not gated. + """ + 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)) + equivalent = ri >= min_rand_index + metrics: dict[str, float | str | bool] = { + "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, + "equivalent": equivalent, + } + if not equivalent: + print( + f"WARNING: rand index {ri:.6g} below threshold {min_rand_index:.6g} — " + "watershed partitions disagree substantially" + ) + return metrics + + +def time_functions_interleaved( + first: Callable[[np.ndarray, np.ndarray], np.ndarray], + second: Callable[[np.ndarray, np.ndarray], np.ndarray], + heightmap: np.ndarray, + markers: np.ndarray, + repeats: int, +) -> tuple[list[float], np.ndarray, list[float], np.ndarray]: + def timed_call(run): + start = perf_counter() + result = run(heightmap, markers) + return perf_counter() - start, result + + # Warm up imports, JIT/Cython compilation, allocator caches. + first(heightmap, markers) + second(heightmap, markers) + + 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, + heightmap: np.ndarray, + markers: np.ndarray, + metrics: dict[str, float | str | bool], + bic_timings: list[float], + ref_timings: list[float], +): + bic_median = median(bic_timings) + ref_median = median(ref_timings) + speedup = ref_median / bic_median if bic_median > 0 else float("inf") + n_markers = int(markers.max()) + + print(f"Watershed {ndim}D comparison") + print(f"heightmap shape: {heightmap.shape}, dtype: {heightmap.dtype}") + print(f"markers: {n_markers} seeds (dtype={markers.dtype})") + 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"exact label equality: {metrics['exact_label_equality']}") + print(f"within thresholds: {metrics['equivalent']}") + print(f"bioimage-cpp median runtime: {bic_median:.6f} s") + print(f"skimage reference median runtime: {ref_median:.6f} s") + print(f"reference / bioimage-cpp runtime ratio: {speedup:.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: + heightmap, markers = prepare_2d_problem( + affinities, + offsets, + z=z, + yx_shape=yx_shape, + smoothing_sigma=smoothing_sigma, + ) + elif ndim == 3: + 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}") + + ref_timings, ref_seg, bic_timings, bic_seg = time_functions_interleaved( + run_skimage_reference, + run_bioimage_cpp, + heightmap, + markers, + repeats, + ) + metrics = compare_segmentations(bic_seg, ref_seg) + print_report( + ndim=ndim, + heightmap=heightmap, + markers=markers, + 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( + "--smoothing-sigma", + type=float, + default=1.5, + help="Gaussian sigma applied to the heightmap before finding local minima.", + ) diff --git a/development/segmentation/_watershed_from_affinities_equivalence.py b/development/segmentation/_watershed_from_affinities_equivalence.py new file mode 100644 index 0000000..666fa9f --- /dev/null +++ b/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.", + ) diff --git a/development/segmentation/check_watershed_2d.py b/development/segmentation/check_watershed_2d.py new file mode 100644 index 0000000..1e42127 --- /dev/null +++ b/development/segmentation/check_watershed_2d.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import argparse + +from _watershed_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and skimage 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() diff --git a/development/segmentation/check_watershed_3d.py b/development/segmentation/check_watershed_3d.py new file mode 100644 index 0000000..1274c2e --- /dev/null +++ b/development/segmentation/check_watershed_3d.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import argparse + +from _watershed_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and skimage 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() diff --git a/development/segmentation/check_watershed_from_affinities_2d.py b/development/segmentation/check_watershed_from_affinities_2d.py new file mode 100644 index 0000000..b9b5a66 --- /dev/null +++ b/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() diff --git a/development/segmentation/check_watershed_from_affinities_3d.py b/development/segmentation/check_watershed_from_affinities_3d.py new file mode 100644 index 0000000..1a60d0d --- /dev/null +++ b/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() diff --git a/examples/segmentation/watershed.py b/examples/segmentation/watershed.py new file mode 100644 index 0000000..58ac0e1 --- /dev/null +++ b/examples/segmentation/watershed.py @@ -0,0 +1,104 @@ +"""Apply both watershed variants to the cached ISBI volume and visualize. + +Runs ``bic.segmentation.watershed`` on a node heightmap derived from the +nearest-neighbour affinity channels (``1 - mean(NN affs)``) and +``bic.segmentation.watershed_from_affinities`` directly on those same +channels. Markers come from labelled local minima of the smoothed +heightmap, so both algorithms see the same seeds and the two segmentations +can be compared side-by-side in napari. +""" + +import napari +import numpy as np +from scipy import ndimage as ndi +from skimage.feature import peak_local_max +from skimage.measure import label + +import bioimage_cpp as bic +from bioimage_cpp._data import load_isbi_affinities, load_isbi_raw + + +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 _filter_2d(affinities, offsets, z=0): + channels_2d = [i for i, offset in enumerate(offsets) if offset[0] == 0] + affinities_2d = np.ascontiguousarray(affinities[channels_2d, z]) + offsets_2d = [tuple(offset[1:]) for offset in (offsets[i] for i in channels_2d)] + return affinities_2d, offsets_2d + + +def build_inputs( + affinities, + offsets, + *, + smoothing_sigma=4.0, + min_distance=8, +): + """Return (nn_affinities, nn_offsets, heightmap, markers). + + Seeds are local maxima of the smoothed inverted-boundary score (i.e. + the "interior-ness" of each pixel). Heavily smoothing the noisy + per-pixel NN-affinity map before peak-finding suppresses the spurious + minima on membranes and within cells that one gets from raw + ``local_minima``, while ``min_distance`` enforces one seed per cell. + """ + nn_channels = _nearest_neighbour_channels(offsets) + nn_affinities = np.ascontiguousarray(affinities[nn_channels]) + nn_offsets = [tuple(offsets[i]) for i in nn_channels] + + # Boundary probability heightmap: ~0 inside cells, ~1 on membranes. + heightmap = (1.0 - nn_affinities.mean(axis=0)).astype(np.float32) + + interior_score = 1.0 - ndi.gaussian_filter(heightmap, sigma=smoothing_sigma) + peak_coords = peak_local_max( + interior_score, + min_distance=min_distance, + exclude_border=False, + ) + seed_mask = np.zeros_like(interior_score, dtype=bool) + if peak_coords.size > 0: + seed_mask[tuple(peak_coords.T)] = True + markers = label(seed_mask, connectivity=1).astype(np.int32, copy=False) + + return nn_affinities, nn_offsets, heightmap, np.ascontiguousarray(markers) + + +def main(): + affinities, offsets = load_isbi_affinities() + raw = load_isbi_raw() + + run_2d = True + if run_2d: + affinities, offsets = _filter_2d(affinities, offsets, z=0) + raw = raw[0] + + nn_affinities, nn_offsets, heightmap, markers = build_inputs(affinities, offsets) + + print(f"Heightmap shape: {heightmap.shape}, markers: {int(markers.max())} seeds") + + print("Running watershed on heightmap...") + seg_heightmap = bic.segmentation.watershed(heightmap, markers) + print("Running watershed_from_affinities...") + seg_affinity = bic.segmentation.watershed_from_affinities( + nn_affinities, nn_offsets, markers, + ) + + viewer = napari.Viewer() + viewer.add_image(raw, name="raw") + viewer.add_image(heightmap, name="heightmap (1 - mean NN affinities)") + viewer.add_image(nn_affinities, name="NN affinities", channel_axis=0) + viewer.add_labels(markers, name="markers (distance-transform peaks)") + viewer.add_labels(seg_heightmap, name="watershed (heightmap)") + viewer.add_labels(seg_affinity, name="watershed_from_affinities") + napari.run() + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/segmentation/watershed.hxx b/include/bioimage_cpp/segmentation/watershed.hxx new file mode 100644 index 0000000..54952fb --- /dev/null +++ b/include/bioimage_cpp/segmentation/watershed.hxx @@ -0,0 +1,980 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/profile.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp { + +namespace detail_ws { + +// Bucket-queue marker watershed. Quantises the heightmap to `n_levels` +// uint16 levels and floods one level at a time with Meyer-style monotone +// semantics: a neighbour pushed from level L always lands at +// max(level[neighbour], L). This is O(N + n_levels) — effectively O(N) for +// typical bioimage volumes — and replaces the std::priority_queue path that +// the first version of this function used. +// +// Tie-breaking on equal levels (including ULP-level differences in the +// quantised range) is intentionally unspecified, matching the documented +// API contract. + +constexpr std::size_t k_n_levels = 65536; + +template +void watershed_2d( + const ConstArrayView &image, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + BIOIMAGE_PROFILE_INIT(profile); + + const bool has_mask = !mask.shape.empty(); + const std::int64_t Y = image.shape[0]; + const std::int64_t X = image.shape[1]; + const std::uint64_t N = + static_cast(Y) * static_cast(X); + + { + BIOIMAGE_PROFILE_SCOPE(profile, "init_output"); + for (std::uint64_t node = 0; node < N; ++node) { + out.data[node] = LabelT{0}; + } + } + + HeightT h_min{}; + HeightT h_max{}; + bool any_valid = false; + { + BIOIMAGE_PROFILE_SCOPE(profile, "range"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + const HeightT v = image.data[node]; + if (!any_valid) { + h_min = v; + h_max = v; + any_valid = true; + } else { + if (v < h_min) { + h_min = v; + } + if (v > h_max) { + h_max = v; + } + } + } + } + if (!any_valid) { + BIOIMAGE_PROFILE_REPORT(profile); + return; + } + + std::vector levels(static_cast(N), 0); + { + BIOIMAGE_PROFILE_SCOPE(profile, "quantize"); + const double range = static_cast(h_max) - static_cast(h_min); + if (range > 0.0) { + const double scale = static_cast(k_n_levels - 1) / range; + const double max_level = static_cast(k_n_levels - 1); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + double q = + (static_cast(image.data[node]) - static_cast(h_min)) * scale; + if (q < 0.0) { + q = 0.0; + } else if (q > max_level) { + q = max_level; + } + levels[node] = static_cast(q); + } + } + } + + std::vector> buckets(k_n_levels); + std::size_t current_level = k_n_levels; + { + BIOIMAGE_PROFILE_SCOPE(profile, "seed_pass"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + const auto marker = markers.data[node]; + if (marker != LabelT{0}) { + out.data[node] = marker; + const std::size_t lvl = levels[node]; + buckets[lvl].push_back(node); + if (lvl < current_level) { + current_level = lvl; + } + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "main_loop"); + while (current_level < k_n_levels) { + auto &bucket = buckets[current_level]; + while (!bucket.empty()) { + const std::uint64_t node = bucket.back(); + bucket.pop_back(); + const auto label = out.data[node]; + const std::int64_t y = static_cast(node) / X; + const std::int64_t x = static_cast(node) - y * X; + + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (x > 0) { + const std::uint64_t n = node - 1; + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + } + ++current_level; + } + } + + BIOIMAGE_PROFILE_REPORT(profile); +} + +template +void watershed_3d( + const ConstArrayView &image, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + BIOIMAGE_PROFILE_INIT(profile); + + const bool has_mask = !mask.shape.empty(); + const std::int64_t Z = image.shape[0]; + const std::int64_t Y = image.shape[1]; + const std::int64_t X = image.shape[2]; + const std::int64_t YX = Y * X; + const std::uint64_t N = + static_cast(Z) * static_cast(YX); + + { + BIOIMAGE_PROFILE_SCOPE(profile, "init_output"); + for (std::uint64_t node = 0; node < N; ++node) { + out.data[node] = LabelT{0}; + } + } + + HeightT h_min{}; + HeightT h_max{}; + bool any_valid = false; + { + BIOIMAGE_PROFILE_SCOPE(profile, "range"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + const HeightT v = image.data[node]; + if (!any_valid) { + h_min = v; + h_max = v; + any_valid = true; + } else { + if (v < h_min) { + h_min = v; + } + if (v > h_max) { + h_max = v; + } + } + } + } + if (!any_valid) { + BIOIMAGE_PROFILE_REPORT(profile); + return; + } + + std::vector levels(static_cast(N), 0); + { + BIOIMAGE_PROFILE_SCOPE(profile, "quantize"); + const double range = static_cast(h_max) - static_cast(h_min); + if (range > 0.0) { + const double scale = static_cast(k_n_levels - 1) / range; + const double max_level = static_cast(k_n_levels - 1); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + double q = + (static_cast(image.data[node]) - static_cast(h_min)) * scale; + if (q < 0.0) { + q = 0.0; + } else if (q > max_level) { + q = max_level; + } + levels[node] = static_cast(q); + } + } + } + + std::vector> buckets(k_n_levels); + std::size_t current_level = k_n_levels; + { + BIOIMAGE_PROFILE_SCOPE(profile, "seed_pass"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) { + continue; + } + const auto marker = markers.data[node]; + if (marker != LabelT{0}) { + out.data[node] = marker; + const std::size_t lvl = levels[node]; + buckets[lvl].push_back(node); + if (lvl < current_level) { + current_level = lvl; + } + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "main_loop"); + while (current_level < k_n_levels) { + auto &bucket = buckets[current_level]; + while (!bucket.empty()) { + const std::uint64_t node = bucket.back(); + bucket.pop_back(); + const auto label = out.data[node]; + const std::int64_t z = static_cast(node) / YX; + const std::int64_t rem = static_cast(node) - z * YX; + const std::int64_t y = rem / X; + const std::int64_t x = rem - y * X; + + if (z > 0) { + const std::uint64_t n = node - static_cast(YX); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (z + 1 < Z) { + const std::uint64_t n = node + static_cast(YX); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (x > 0) { + const std::uint64_t n = node - 1; + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + if ((!has_mask || mask.data[n] != 0) && out.data[n] == LabelT{0}) { + out.data[n] = label; + std::size_t lvl = levels[n]; + if (lvl < current_level) lvl = current_level; + buckets[lvl].push_back(n); + } + } + } + ++current_level; + } + } + + BIOIMAGE_PROFILE_REPORT(profile); +} + +// Bucket entry for affinity-based watershed. Each entry is a pending claim +// for `target` by `source_label`; the actual labelling happens at pop time +// because competing claims with different priorities can target the same +// pixel. +template +struct AffinityWatershedEntry { + std::uint64_t target; + LabelT source_label; +}; + +template +void watershed_2d_from_affinities( + const ConstArrayView &affinities, + const std::array &channel_for_axis, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + BIOIMAGE_PROFILE_INIT(profile); + + const bool has_mask = !mask.shape.empty(); + const std::size_t C = static_cast(affinities.shape[0]); + const std::int64_t Y = affinities.shape[1]; + const std::int64_t X = affinities.shape[2]; + const std::uint64_t N = + static_cast(Y) * static_cast(X); + + const std::size_t cy = channel_for_axis[0]; + const std::size_t cx = channel_for_axis[1]; + + { + BIOIMAGE_PROFILE_SCOPE(profile, "init_output"); + for (std::uint64_t node = 0; node < N; ++node) { + out.data[node] = LabelT{0}; + } + } + + AffT a_min{}; + AffT a_max{}; + bool any_valid = false; + { + BIOIMAGE_PROFILE_SCOPE(profile, "range"); + const std::size_t total = C * static_cast(N); + for (std::size_t i = 0; i < total; ++i) { + if (has_mask) { + const std::uint64_t node = static_cast(i) % N; + if (mask.data[node] == 0) { + continue; + } + } + const AffT v = affinities.data[i]; + if (!any_valid) { + a_min = v; + a_max = v; + any_valid = true; + } else { + if (v < a_min) a_min = v; + if (v > a_max) a_max = v; + } + } + } + if (!any_valid) { + BIOIMAGE_PROFILE_REPORT(profile); + return; + } + + std::vector levels(C * static_cast(N), 0); + { + BIOIMAGE_PROFILE_SCOPE(profile, "quantize"); + const double a_max_d = static_cast(a_max); + const double range = a_max_d - static_cast(a_min); + if (range > 0.0) { + const double scale = static_cast(k_n_levels - 1) / range; + const double max_level = static_cast(k_n_levels - 1); + const std::size_t total = C * static_cast(N); + for (std::size_t i = 0; i < total; ++i) { + double q = + (a_max_d - static_cast(affinities.data[i])) * scale; + if (q < 0.0) { + q = 0.0; + } else if (q > max_level) { + q = max_level; + } + levels[i] = static_cast(q); + } + } + } + + using Entry = AffinityWatershedEntry; + std::vector> buckets(k_n_levels); + std::size_t current_level = k_n_levels; + + auto try_push_seed = [&](std::uint64_t neighbor, LabelT label, std::size_t edge_lvl) { + if (has_mask && mask.data[neighbor] == 0) return; + if (out.data[neighbor] != LabelT{0}) return; + buckets[edge_lvl].push_back(Entry{neighbor, label}); + if (edge_lvl < current_level) current_level = edge_lvl; + }; + auto try_push_main = [&](std::uint64_t neighbor, LabelT label, std::size_t edge_lvl) { + if (has_mask && mask.data[neighbor] == 0) return; + if (out.data[neighbor] != LabelT{0}) return; + if (edge_lvl < current_level) edge_lvl = current_level; + buckets[edge_lvl].push_back(Entry{neighbor, label}); + }; + + { + BIOIMAGE_PROFILE_SCOPE(profile, "seed_pass"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) continue; + const auto marker = markers.data[node]; + if (marker == LabelT{0}) continue; + out.data[node] = marker; + + const std::int64_t y = static_cast(node) / X; + const std::int64_t x = static_cast(node) - y * X; + + if constexpr (Sign < 0) { + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + try_push_seed(n, marker, levels[cy * N + node]); + } + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + try_push_seed(n, marker, levels[cy * N + n]); + } + if (x > 0) { + const std::uint64_t n = node - 1; + try_push_seed(n, marker, levels[cx * N + node]); + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + try_push_seed(n, marker, levels[cx * N + n]); + } + } else { + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + try_push_seed(n, marker, levels[cy * N + node]); + } + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + try_push_seed(n, marker, levels[cy * N + n]); + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + try_push_seed(n, marker, levels[cx * N + node]); + } + if (x > 0) { + const std::uint64_t n = node - 1; + try_push_seed(n, marker, levels[cx * N + n]); + } + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "main_loop"); + while (current_level < k_n_levels) { + auto &bucket = buckets[current_level]; + while (!bucket.empty()) { + const Entry e = bucket.back(); + bucket.pop_back(); + const std::uint64_t u = e.target; + if (out.data[u] != LabelT{0}) continue; + out.data[u] = e.source_label; + const LabelT label = e.source_label; + + const std::int64_t y = static_cast(u) / X; + const std::int64_t x = static_cast(u) - y * X; + + if constexpr (Sign < 0) { + if (y > 0) { + const std::uint64_t n = u - static_cast(X); + try_push_main(n, label, levels[cy * N + u]); + } + if (y + 1 < Y) { + const std::uint64_t n = u + static_cast(X); + try_push_main(n, label, levels[cy * N + n]); + } + if (x > 0) { + const std::uint64_t n = u - 1; + try_push_main(n, label, levels[cx * N + u]); + } + if (x + 1 < X) { + const std::uint64_t n = u + 1; + try_push_main(n, label, levels[cx * N + n]); + } + } else { + if (y + 1 < Y) { + const std::uint64_t n = u + static_cast(X); + try_push_main(n, label, levels[cy * N + u]); + } + if (y > 0) { + const std::uint64_t n = u - static_cast(X); + try_push_main(n, label, levels[cy * N + n]); + } + if (x + 1 < X) { + const std::uint64_t n = u + 1; + try_push_main(n, label, levels[cx * N + u]); + } + if (x > 0) { + const std::uint64_t n = u - 1; + try_push_main(n, label, levels[cx * N + n]); + } + } + } + ++current_level; + } + } + + BIOIMAGE_PROFILE_REPORT(profile); +} + +template +void watershed_3d_from_affinities( + const ConstArrayView &affinities, + const std::array &channel_for_axis, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + BIOIMAGE_PROFILE_INIT(profile); + + const bool has_mask = !mask.shape.empty(); + const std::size_t C = static_cast(affinities.shape[0]); + const std::int64_t Z = affinities.shape[1]; + const std::int64_t Y = affinities.shape[2]; + const std::int64_t X = affinities.shape[3]; + const std::int64_t YX = Y * X; + const std::uint64_t N = + static_cast(Z) * static_cast(YX); + + const std::size_t cz = channel_for_axis[0]; + const std::size_t cy = channel_for_axis[1]; + const std::size_t cx = channel_for_axis[2]; + + { + BIOIMAGE_PROFILE_SCOPE(profile, "init_output"); + for (std::uint64_t node = 0; node < N; ++node) { + out.data[node] = LabelT{0}; + } + } + + AffT a_min{}; + AffT a_max{}; + bool any_valid = false; + { + BIOIMAGE_PROFILE_SCOPE(profile, "range"); + const std::size_t total = C * static_cast(N); + for (std::size_t i = 0; i < total; ++i) { + if (has_mask) { + const std::uint64_t node = static_cast(i) % N; + if (mask.data[node] == 0) { + continue; + } + } + const AffT v = affinities.data[i]; + if (!any_valid) { + a_min = v; + a_max = v; + any_valid = true; + } else { + if (v < a_min) a_min = v; + if (v > a_max) a_max = v; + } + } + } + if (!any_valid) { + BIOIMAGE_PROFILE_REPORT(profile); + return; + } + + std::vector levels(C * static_cast(N), 0); + { + BIOIMAGE_PROFILE_SCOPE(profile, "quantize"); + const double a_max_d = static_cast(a_max); + const double range = a_max_d - static_cast(a_min); + if (range > 0.0) { + const double scale = static_cast(k_n_levels - 1) / range; + const double max_level = static_cast(k_n_levels - 1); + const std::size_t total = C * static_cast(N); + for (std::size_t i = 0; i < total; ++i) { + double q = + (a_max_d - static_cast(affinities.data[i])) * scale; + if (q < 0.0) { + q = 0.0; + } else if (q > max_level) { + q = max_level; + } + levels[i] = static_cast(q); + } + } + } + + using Entry = AffinityWatershedEntry; + std::vector> buckets(k_n_levels); + std::size_t current_level = k_n_levels; + + auto try_push_seed = [&](std::uint64_t neighbor, LabelT label, std::size_t edge_lvl) { + if (has_mask && mask.data[neighbor] == 0) return; + if (out.data[neighbor] != LabelT{0}) return; + buckets[edge_lvl].push_back(Entry{neighbor, label}); + if (edge_lvl < current_level) current_level = edge_lvl; + }; + auto try_push_main = [&](std::uint64_t neighbor, LabelT label, std::size_t edge_lvl) { + if (has_mask && mask.data[neighbor] == 0) return; + if (out.data[neighbor] != LabelT{0}) return; + if (edge_lvl < current_level) edge_lvl = current_level; + buckets[edge_lvl].push_back(Entry{neighbor, label}); + }; + + { + BIOIMAGE_PROFILE_SCOPE(profile, "seed_pass"); + for (std::uint64_t node = 0; node < N; ++node) { + if (has_mask && mask.data[node] == 0) continue; + const auto marker = markers.data[node]; + if (marker == LabelT{0}) continue; + out.data[node] = marker; + + const std::int64_t z = static_cast(node) / YX; + const std::int64_t rem = static_cast(node) - z * YX; + const std::int64_t y = rem / X; + const std::int64_t x = rem - y * X; + + if constexpr (Sign < 0) { + if (z > 0) { + const std::uint64_t n = node - static_cast(YX); + try_push_seed(n, marker, levels[cz * N + node]); + } + if (z + 1 < Z) { + const std::uint64_t n = node + static_cast(YX); + try_push_seed(n, marker, levels[cz * N + n]); + } + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + try_push_seed(n, marker, levels[cy * N + node]); + } + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + try_push_seed(n, marker, levels[cy * N + n]); + } + if (x > 0) { + const std::uint64_t n = node - 1; + try_push_seed(n, marker, levels[cx * N + node]); + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + try_push_seed(n, marker, levels[cx * N + n]); + } + } else { + if (z + 1 < Z) { + const std::uint64_t n = node + static_cast(YX); + try_push_seed(n, marker, levels[cz * N + node]); + } + if (z > 0) { + const std::uint64_t n = node - static_cast(YX); + try_push_seed(n, marker, levels[cz * N + n]); + } + if (y + 1 < Y) { + const std::uint64_t n = node + static_cast(X); + try_push_seed(n, marker, levels[cy * N + node]); + } + if (y > 0) { + const std::uint64_t n = node - static_cast(X); + try_push_seed(n, marker, levels[cy * N + n]); + } + if (x + 1 < X) { + const std::uint64_t n = node + 1; + try_push_seed(n, marker, levels[cx * N + node]); + } + if (x > 0) { + const std::uint64_t n = node - 1; + try_push_seed(n, marker, levels[cx * N + n]); + } + } + } + } + + { + BIOIMAGE_PROFILE_SCOPE(profile, "main_loop"); + while (current_level < k_n_levels) { + auto &bucket = buckets[current_level]; + while (!bucket.empty()) { + const Entry e = bucket.back(); + bucket.pop_back(); + const std::uint64_t u = e.target; + if (out.data[u] != LabelT{0}) continue; + out.data[u] = e.source_label; + const LabelT label = e.source_label; + + const std::int64_t z = static_cast(u) / YX; + const std::int64_t rem = static_cast(u) - z * YX; + const std::int64_t y = rem / X; + const std::int64_t x = rem - y * X; + + if constexpr (Sign < 0) { + if (z > 0) { + const std::uint64_t n = u - static_cast(YX); + try_push_main(n, label, levels[cz * N + u]); + } + if (z + 1 < Z) { + const std::uint64_t n = u + static_cast(YX); + try_push_main(n, label, levels[cz * N + n]); + } + if (y > 0) { + const std::uint64_t n = u - static_cast(X); + try_push_main(n, label, levels[cy * N + u]); + } + if (y + 1 < Y) { + const std::uint64_t n = u + static_cast(X); + try_push_main(n, label, levels[cy * N + n]); + } + if (x > 0) { + const std::uint64_t n = u - 1; + try_push_main(n, label, levels[cx * N + u]); + } + if (x + 1 < X) { + const std::uint64_t n = u + 1; + try_push_main(n, label, levels[cx * N + n]); + } + } else { + if (z + 1 < Z) { + const std::uint64_t n = u + static_cast(YX); + try_push_main(n, label, levels[cz * N + u]); + } + if (z > 0) { + const std::uint64_t n = u - static_cast(YX); + try_push_main(n, label, levels[cz * N + n]); + } + if (y + 1 < Y) { + const std::uint64_t n = u + static_cast(X); + try_push_main(n, label, levels[cy * N + u]); + } + if (y > 0) { + const std::uint64_t n = u - static_cast(X); + try_push_main(n, label, levels[cy * N + n]); + } + if (x + 1 < X) { + const std::uint64_t n = u + 1; + try_push_main(n, label, levels[cx * N + u]); + } + if (x > 0) { + const std::uint64_t n = u - 1; + try_push_main(n, label, levels[cx * N + n]); + } + } + } + ++current_level; + } + } + + BIOIMAGE_PROFILE_REPORT(profile); +} + +} // namespace detail_ws + +// Marker-controlled flooding watershed on a 2D or 3D image. +// +// `image` is the heightmap. `markers` carries non-zero seed labels that get +// propagated to neighbouring pixels in order of increasing height. If `mask` +// is non-empty, only pixels with a non-zero mask value participate; the +// remaining pixels stay 0 in the output. Connectivity is 1 (axis-aligned +// 4-neighbours in 2D, 6-neighbours in 3D). +// +// Internally uses a 65536-bucket queue with Meyer-style monotone flooding. +// Tie-breaking on equal heights (including 1-ULP differences within a +// quantisation bucket) is unspecified, as documented in the Python wrapper. +template +void watershed( + const ConstArrayView &image, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + if (image.ndim() != 2 && image.ndim() != 3) { + throw std::invalid_argument( + "image must have ndim 2 or 3, got ndim=" + std::to_string(image.ndim()) + ); + } + if (markers.shape != image.shape) { + throw std::invalid_argument("markers shape must match image shape"); + } + if (out.shape != image.shape) { + throw std::invalid_argument("out shape must match image shape"); + } + const bool has_mask = !mask.shape.empty(); + if (has_mask && mask.shape != image.shape) { + throw std::invalid_argument("mask shape must match image shape"); + } + + if (image.ndim() == 2) { + detail_ws::watershed_2d(image, markers, mask, out); + } else { + detail_ws::watershed_3d(image, markers, mask, out); + } +} + +// Marker-controlled watershed driven by nearest-neighbour edge affinities. +// +// `affinities` has shape `(C, *spatial)` where `C == spatial_ndim` and each +// channel encodes the affinity of one axis-aligned edge per pixel. `offsets` +// must contain exactly `C` nearest-neighbour offsets (each with one entry +// ±1 and the rest zero), all with the same sign, and together covering every +// spatial axis exactly once. Higher affinity = stronger bond = processed +// first. +template +void watershed_from_affinities( + const ConstArrayView &affinities, + const std::vector> &offsets, + const ConstArrayView &markers, + const ConstArrayView &mask, + const ArrayView &out +) { + if (affinities.ndim() != 3 && affinities.ndim() != 4) { + throw std::invalid_argument( + "affinities must have ndim 3 or 4, got ndim=" + std::to_string(affinities.ndim()) + ); + } + const auto spatial_ndim = static_cast(affinities.ndim() - 1); + if (static_cast(affinities.shape[0]) != spatial_ndim) { + throw std::invalid_argument( + "affinities channel count must equal spatial ndim, got channels=" + + std::to_string(affinities.shape[0]) + ", spatial ndim=" + + std::to_string(spatial_ndim) + ); + } + + std::vector spatial_shape( + affinities.shape.begin() + 1, + affinities.shape.end() + ); + if (markers.shape != spatial_shape) { + throw std::invalid_argument("markers shape must match affinities spatial shape"); + } + if (out.shape != spatial_shape) { + throw std::invalid_argument("out shape must match affinities spatial shape"); + } + const bool has_mask = !mask.shape.empty(); + if (has_mask && mask.shape != spatial_shape) { + throw std::invalid_argument("mask shape must match affinities spatial shape"); + } + + if (offsets.size() != spatial_ndim) { + throw std::invalid_argument( + "offsets count must equal affinities channel count, got offsets=" + + std::to_string(offsets.size()) + ", channels=" + + std::to_string(spatial_ndim) + ); + } + + int sign = 0; + std::array channel_for_axis{}; + std::array axis_seen{false, false, false}; + for (std::size_t c = 0; c < spatial_ndim; ++c) { + if (offsets[c].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::size_t nonzero_axis = 0; + int nonzero_count = 0; + int channel_sign = 0; + for (std::size_t a = 0; a < spatial_ndim; ++a) { + const auto v = offsets[c][a]; + if (v == 0) continue; + ++nonzero_count; + nonzero_axis = a; + if (v == 1) channel_sign = 1; + else if (v == -1) channel_sign = -1; + else nonzero_count = -1; + } + if (nonzero_count != 1) { + throw std::invalid_argument( + "each offset must be a nearest-neighbour offset (one entry of value +1 or -1, rest 0)" + ); + } + if (sign == 0) { + sign = channel_sign; + } else if (sign != channel_sign) { + throw std::invalid_argument( + "all offsets must have the same sign (positive or negative direction), " + "mixing positive and negative is not supported" + ); + } + if (axis_seen[nonzero_axis]) { + throw std::invalid_argument( + "each spatial axis must be covered by exactly one offset; " + "axis " + std::to_string(nonzero_axis) + " is covered more than once" + ); + } + axis_seen[nonzero_axis] = true; + channel_for_axis[nonzero_axis] = c; + } + for (std::size_t a = 0; a < spatial_ndim; ++a) { + if (!axis_seen[a]) { + throw std::invalid_argument( + "each spatial axis must be covered by exactly one offset; " + "axis " + std::to_string(a) + " is not covered" + ); + } + } + + if (spatial_ndim == 2) { + const std::array cfa{channel_for_axis[0], channel_for_axis[1]}; + if (sign < 0) { + detail_ws::watershed_2d_from_affinities<-1, AffT, LabelT>( + affinities, cfa, markers, mask, out + ); + } else { + detail_ws::watershed_2d_from_affinities<+1, AffT, LabelT>( + affinities, cfa, markers, mask, out + ); + } + } else { + const std::array cfa{ + channel_for_axis[0], channel_for_axis[1], channel_for_axis[2] + }; + if (sign < 0) { + detail_ws::watershed_3d_from_affinities<-1, AffT, LabelT>( + affinities, cfa, markers, mask, out + ); + } else { + detail_ws::watershed_3d_from_affinities<+1, AffT, LabelT>( + affinities, cfa, markers, mask, out + ); + } + } +} + +} // namespace bioimage_cpp diff --git a/src/bindings/segmentation.cxx b/src/bindings/segmentation.cxx index 10154d3..59fccf4 100644 --- a/src/bindings/segmentation.cxx +++ b/src/bindings/segmentation.cxx @@ -3,8 +3,10 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/segmentation/mutex_watershed.hxx" #include "bioimage_cpp/segmentation/semantic_mutex_watershed.hxx" +#include "bioimage_cpp/segmentation/watershed.hxx" #include +#include #include #include @@ -12,6 +14,7 @@ #include #include #include +#include #include #include @@ -174,6 +177,149 @@ std::pair semantic_mutex_watershed_grid_t( ); } +template +nb::ndarray watershed_t( + nb::ndarray image, + nb::ndarray markers, + std::optional> mask +) { + if (markers.ndim() != image.ndim()) { + throw std::invalid_argument("markers shape must match image shape"); + } + std::vector image_shape(image.ndim()); + std::vector out_shape(image.ndim()); + for (std::size_t axis = 0; axis < image.ndim(); ++axis) { + image_shape[axis] = static_cast(image.shape(axis)); + out_shape[axis] = image.shape(axis); + if (markers.shape(axis) != image.shape(axis)) { + throw std::invalid_argument("markers shape must match image shape"); + } + } + + std::vector mask_shape; + const std::uint8_t *mask_data = nullptr; + if (mask.has_value()) { + if (mask->ndim() != image.ndim()) { + throw std::invalid_argument("mask shape must match image shape"); + } + mask_shape.resize(mask->ndim()); + for (std::size_t axis = 0; axis < mask->ndim(); ++axis) { + mask_shape[axis] = static_cast(mask->shape(axis)); + if (mask->shape(axis) != image.shape(axis)) { + throw std::invalid_argument("mask shape must match image shape"); + } + } + mask_data = mask->data(); + } + + const auto number_of_nodes = std::accumulate( + image_shape.begin(), + image_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + ); + auto *data = new LabelT[static_cast(number_of_nodes)](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + + ConstArrayView image_view{image.data(), image_shape, {}}; + ConstArrayView markers_view{markers.data(), image_shape, {}}; + ConstArrayView mask_view{mask_data, mask_shape, {}}; + ArrayView out_view{data, image_shape, {}}; + + { + nb::gil_scoped_release release; + watershed(image_view, markers_view, mask_view, out_view); + } + + return nb::ndarray( + data, + out_shape.size(), + out_shape.data(), + owner + ); +} + +template +nb::ndarray watershed_from_affinities_t( + nb::ndarray affinities, + const std::vector> &offsets, + nb::ndarray markers, + std::optional> mask +) { + if (affinities.ndim() != 3 && affinities.ndim() != 4) { + throw std::invalid_argument( + "affinities must have ndim 3 or 4, got ndim=" + + std::to_string(affinities.ndim()) + ); + } + std::vector affinity_shape(affinities.ndim()); + for (std::size_t axis = 0; axis < affinities.ndim(); ++axis) { + affinity_shape[axis] = static_cast(affinities.shape(axis)); + } + + const std::size_t spatial_ndim = affinities.ndim() - 1; + std::vector spatial_shape(spatial_ndim); + std::vector out_shape(spatial_ndim); + for (std::size_t axis = 0; axis < spatial_ndim; ++axis) { + spatial_shape[axis] = + static_cast(affinities.shape(axis + 1)); + out_shape[axis] = affinities.shape(axis + 1); + } + + if (markers.ndim() != spatial_ndim) { + throw std::invalid_argument("markers shape must match affinities spatial shape"); + } + for (std::size_t axis = 0; axis < spatial_ndim; ++axis) { + if (markers.shape(axis) != affinities.shape(axis + 1)) { + throw std::invalid_argument("markers shape must match affinities spatial shape"); + } + } + + std::vector mask_shape; + const std::uint8_t *mask_data = nullptr; + if (mask.has_value()) { + if (mask->ndim() != spatial_ndim) { + throw std::invalid_argument("mask shape must match affinities spatial shape"); + } + mask_shape.resize(mask->ndim()); + for (std::size_t axis = 0; axis < spatial_ndim; ++axis) { + mask_shape[axis] = static_cast(mask->shape(axis)); + if (mask->shape(axis) != affinities.shape(axis + 1)) { + throw std::invalid_argument("mask shape must match affinities spatial shape"); + } + } + mask_data = mask->data(); + } + + const auto number_of_nodes = 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; } + ); + auto *data = new LabelT[static_cast(number_of_nodes)](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + + ConstArrayView affinities_view{affinities.data(), affinity_shape, {}}; + ConstArrayView markers_view{markers.data(), spatial_shape, {}}; + ConstArrayView mask_view{mask_data, mask_shape, {}}; + ArrayView out_view{data, spatial_shape, {}}; + + { + nb::gil_scoped_release release; + watershed_from_affinities( + affinities_view, offsets, markers_view, mask_view, out_view + ); + } + + return nb::ndarray( + data, + out_shape.size(), + out_shape.data(), + owner + ); +} + } // namespace void bind_segmentation(nb::module_ &m) { @@ -215,6 +361,144 @@ void bind_segmentation(nb::module_ &m) { nb::arg("number_of_offsets"), "Run semantic mutex watershed on a 2D or 3D image-derived grid graph with float64 affinities." ); + + m.def( + "_watershed_float32_uint32", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float32 image with uint32 markers." + ); + m.def( + "_watershed_float32_uint64", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float32 image with uint64 markers." + ); + m.def( + "_watershed_float32_int32", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float32 image with int32 markers." + ); + m.def( + "_watershed_float32_int64", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float32 image with int64 markers." + ); + m.def( + "_watershed_float64_uint32", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float64 image with uint32 markers." + ); + m.def( + "_watershed_float64_uint64", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float64 image with uint64 markers." + ); + m.def( + "_watershed_float64_int32", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float64 image with int32 markers." + ); + m.def( + "_watershed_float64_int64", + &watershed_t, + nb::arg("image"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Marker-controlled watershed on a 2D or 3D float64 image with int64 markers." + ); + + m.def( + "_watershed_from_affinities_float32_uint32", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float32 affinities with uint32 markers." + ); + m.def( + "_watershed_from_affinities_float32_uint64", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float32 affinities with uint64 markers." + ); + m.def( + "_watershed_from_affinities_float32_int32", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float32 affinities with int32 markers." + ); + m.def( + "_watershed_from_affinities_float32_int64", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float32 affinities with int64 markers." + ); + m.def( + "_watershed_from_affinities_float64_uint32", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float64 affinities with uint32 markers." + ); + m.def( + "_watershed_from_affinities_float64_uint64", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float64 affinities with uint64 markers." + ); + m.def( + "_watershed_from_affinities_float64_int32", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float64 affinities with int32 markers." + ); + m.def( + "_watershed_from_affinities_float64_int64", + &watershed_from_affinities_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("markers"), + nb::arg("mask") = nb::none(), + "Affinity-based watershed on 2D or 3D nearest-neighbour float64 affinities with int64 markers." + ); } } // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/segmentation/__init__.py b/src/bioimage_cpp/segmentation/__init__.py index 7324596..825e7f4 100644 --- a/src/bioimage_cpp/segmentation/__init__.py +++ b/src/bioimage_cpp/segmentation/__init__.py @@ -1,5 +1,11 @@ """Segmentation algorithms.""" from .mutex_watershed import mutex_watershed, semantic_mutex_watershed +from .watershed import watershed, watershed_from_affinities -__all__ = ["mutex_watershed", "semantic_mutex_watershed"] +__all__ = [ + "mutex_watershed", + "semantic_mutex_watershed", + "watershed", + "watershed_from_affinities", +] diff --git a/src/bioimage_cpp/segmentation/watershed.py b/src/bioimage_cpp/segmentation/watershed.py new file mode 100644 index 0000000..08d3a6b --- /dev/null +++ b/src/bioimage_cpp/segmentation/watershed.py @@ -0,0 +1,261 @@ +"""Marker-controlled watershed.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from .. import _core + +_WATERSHED_BY_DTYPE: dict[np.dtype, dict[np.dtype, object]] = { + np.dtype("float32"): { + np.dtype("uint32"): _core._watershed_float32_uint32, + np.dtype("uint64"): _core._watershed_float32_uint64, + np.dtype("int32"): _core._watershed_float32_int32, + np.dtype("int64"): _core._watershed_float32_int64, + }, + np.dtype("float64"): { + np.dtype("uint32"): _core._watershed_float64_uint32, + np.dtype("uint64"): _core._watershed_float64_uint64, + np.dtype("int32"): _core._watershed_float64_int32, + np.dtype("int64"): _core._watershed_float64_int64, + }, +} + +_WATERSHED_FROM_AFFINITIES_BY_DTYPE: dict[np.dtype, dict[np.dtype, object]] = { + np.dtype("float32"): { + np.dtype("uint32"): _core._watershed_from_affinities_float32_uint32, + np.dtype("uint64"): _core._watershed_from_affinities_float32_uint64, + np.dtype("int32"): _core._watershed_from_affinities_float32_int32, + np.dtype("int64"): _core._watershed_from_affinities_float32_int64, + }, + np.dtype("float64"): { + np.dtype("uint32"): _core._watershed_from_affinities_float64_uint32, + np.dtype("uint64"): _core._watershed_from_affinities_float64_uint64, + np.dtype("int32"): _core._watershed_from_affinities_float64_int32, + np.dtype("int64"): _core._watershed_from_affinities_float64_int64, + }, +} + + +def watershed( + image: np.ndarray, + markers: np.ndarray, + *, + mask: np.ndarray | None = None, +) -> np.ndarray: + """Run a marker-controlled watershed on a 2D or 3D heightmap. + + Pixels with non-zero ``markers`` values are treated as seeds. Their + label is propagated to neighbouring pixels in order of increasing + ``image`` value, using axis-aligned connectivity (4-neighbours in 2D, + 6-neighbours in 3D). Tie-breaking on equal heights is unspecified. + + Parameters + ---------- + image: + Heightmap with shape ``(y, x)`` or ``(z, y, x)``. Supported dtypes + are ``float32`` and ``float64``. Non-contiguous arrays are copied. + markers: + Seed array with the same shape as ``image``. Non-zero entries are + treated as seeds. Supported dtypes are ``uint32``, ``uint64``, + ``int32``, ``int64``. The output dtype matches ``markers``. + mask: + Optional boolean foreground mask with the same shape as ``image``. + ``False`` pixels are excluded from the flooding and stay ``0`` in + the output. A seed under a ``False`` pixel is ignored. + + Returns + ------- + np.ndarray + Segmentation labels with the same shape as ``image`` and the same + dtype as ``markers``. Pixels that the flooding never reaches stay + at ``0``. + """ + image_array = np.asarray(image) + markers_array = np.asarray(markers) + + if image_array.ndim not in (2, 3): + raise ValueError( + "image must have ndim 2 or 3, got ndim=" + str(image_array.ndim) + ) + if markers_array.shape != image_array.shape: + raise ValueError( + "markers shape must match image shape, got " + f"markers shape={markers_array.shape}, image shape={image_array.shape}" + ) + + try: + by_marker = _WATERSHED_BY_DTYPE[image_array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _WATERSHED_BY_DTYPE) + raise TypeError( + f"image must have one of dtypes ({supported}), got dtype={image_array.dtype}" + ) from error + + try: + run = by_marker[markers_array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in by_marker) + raise TypeError( + f"markers must have one of dtypes ({supported}), got dtype={markers_array.dtype}" + ) from error + + mask_arg = None + if mask is not None: + mask_array = np.asarray(mask) + if mask_array.shape != image_array.shape: + raise ValueError( + "mask shape must match image shape, got " + f"mask shape={mask_array.shape}, image shape={image_array.shape}" + ) + if mask_array.dtype != np.dtype("bool"): + raise TypeError(f"mask must have dtype bool, got dtype={mask_array.dtype}") + mask_arg = np.ascontiguousarray(mask_array.view(np.uint8)) + + image_c = np.ascontiguousarray(image_array) + markers_c = np.ascontiguousarray(markers_array) + return run(image_c, markers_c, mask_arg) + + +def watershed_from_affinities( + affinities: np.ndarray, + offsets: Sequence[Sequence[int]], + markers: np.ndarray, + *, + mask: np.ndarray | None = None, +) -> np.ndarray: + """Run a marker-controlled watershed on nearest-neighbour edge affinities. + + Higher-affinity edges are processed first. Each pixel is claimed by the + seed that reaches it via the strongest affinity path; ties on equal + affinities are unspecified. + + Parameters + ---------- + affinities: + Affinity tensor of shape ``(C, y, x)`` for 2D data or + ``(C, z, y, x)`` for 3D, where ``C`` equals the spatial ``ndim``. + Channel ``c`` holds the affinity of one nearest-neighbour edge per + pixel, in the direction given by ``offsets[c]``. Supported dtypes + are ``float32`` and ``float64``. Non-contiguous arrays are copied. + offsets: + One nearest-neighbour offset per channel. Each offset is a tuple of + length ``ndim`` with exactly one entry equal to ``+1`` or ``-1`` and + the rest zero. All offsets must have the same sign — mixing + positive and negative directions is not supported. Each spatial + axis must be covered by exactly one offset. + markers: + Seed array with shape ``affinities.shape[1:]``. Non-zero entries + are seeds. Supported dtypes are ``uint32``, ``uint64``, ``int32``, + ``int64``; the output dtype matches. + mask: + Optional boolean foreground mask with shape + ``affinities.shape[1:]``. ``False`` pixels stay ``0`` in the + output; a seed under a ``False`` pixel is ignored. + + Returns + ------- + np.ndarray + Segmentation labels with shape ``affinities.shape[1:]`` and the + same dtype as ``markers``. + """ + affinities_array = np.asarray(affinities) + markers_array = np.asarray(markers) + + if affinities_array.ndim not in (3, 4): + raise ValueError( + "affinities must have ndim 3 or 4, got ndim=" + str(affinities_array.ndim) + ) + spatial_ndim = affinities_array.ndim - 1 + n_channels = int(affinities_array.shape[0]) + if n_channels != spatial_ndim: + raise ValueError( + "affinities channel count must equal spatial ndim, got " + f"channels={n_channels}, spatial ndim={spatial_ndim}" + ) + if markers_array.shape != affinities_array.shape[1:]: + raise ValueError( + "markers shape must match affinities spatial shape, got " + f"markers shape={markers_array.shape}, " + f"affinities spatial shape={affinities_array.shape[1:]}" + ) + + try: + by_marker = _WATERSHED_FROM_AFFINITIES_BY_DTYPE[affinities_array.dtype] + except KeyError as error: + supported = ", ".join( + str(dtype) for dtype in _WATERSHED_FROM_AFFINITIES_BY_DTYPE + ) + raise TypeError( + "affinities must have one of dtypes " + f"({supported}), got dtype={affinities_array.dtype}" + ) from error + + try: + run = by_marker[markers_array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in by_marker) + raise TypeError( + f"markers must have one of dtypes ({supported}), got dtype={markers_array.dtype}" + ) from error + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + if len(normalized_offsets) != n_channels: + raise ValueError( + "offsets count must equal affinities channel count, got " + f"offsets={len(normalized_offsets)}, channels={n_channels}" + ) + + sign: int | None = None + axes_seen: set[int] = set() + for offset in normalized_offsets: + if len(offset) != spatial_ndim: + raise ValueError( + "each offset must have length matching the spatial ndim, got " + f"spatial ndim={spatial_ndim}" + ) + nonzero = [(a, v) for a, v in enumerate(offset) if v != 0] + if len(nonzero) != 1 or nonzero[0][1] not in (-1, 1): + raise ValueError( + "each offset must be a nearest-neighbour offset (one entry of " + f"value +1 or -1, rest 0), got {offset}" + ) + axis, value = nonzero[0] + if sign is None: + sign = value + elif sign != value: + raise ValueError( + "all offsets must have the same sign (positive or negative " + "direction); mixing positive and negative is not supported" + ) + if axis in axes_seen: + raise ValueError( + "each spatial axis must be covered by exactly one offset; " + f"axis {axis} is covered more than once" + ) + axes_seen.add(axis) + missing = set(range(spatial_ndim)) - axes_seen + if missing: + raise ValueError( + "each spatial axis must be covered by exactly one offset; " + f"missing axis {sorted(missing)[0]}" + ) + + mask_arg = None + if mask is not None: + mask_array = np.asarray(mask) + if mask_array.shape != affinities_array.shape[1:]: + raise ValueError( + "mask shape must match affinities spatial shape, got " + f"mask shape={mask_array.shape}, " + f"affinities spatial shape={affinities_array.shape[1:]}" + ) + if mask_array.dtype != np.dtype("bool"): + raise TypeError(f"mask must have dtype bool, got dtype={mask_array.dtype}") + mask_arg = np.ascontiguousarray(mask_array.view(np.uint8)) + + affinities_c = np.ascontiguousarray(affinities_array) + markers_c = np.ascontiguousarray(markers_array) + return run(affinities_c, normalized_offsets, markers_c, mask_arg) diff --git a/tests/segmentation/test_watershed.py b/tests/segmentation/test_watershed.py new file mode 100644 index 0000000..02ecb9b --- /dev/null +++ b/tests/segmentation/test_watershed.py @@ -0,0 +1,157 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_watershed_2d_ridge_assigns_each_side_to_its_marker(): + # Strictly distinct heights on each side of the ridge make the result + # tie-break-free. + image = np.array([[0.0, 2.0, 9.0, 1.0, 0.0]], dtype=np.float32) + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint32) + + labels = bic.segmentation.watershed(image, markers) + + np.testing.assert_array_equal(labels, np.array([[1, 1, 2, 2, 2]], dtype=np.uint32)) + assert labels.dtype == np.uint32 + assert labels.shape == image.shape + + +def test_watershed_3d_ridge_assigns_each_side_to_its_marker(): + # 3D version of the 2D ridge test — exercises the 3D neighbour code path + # while keeping the answer tie-break-free. + image = np.array([[[0.0, 2.0, 9.0, 1.0, 0.0]]], dtype=np.float64) + markers = np.array([[[1, 0, 0, 0, 2]]], dtype=np.int64) + + labels = bic.segmentation.watershed(image, markers) + + np.testing.assert_array_equal( + labels, np.array([[[1, 1, 2, 2, 2]]], dtype=np.int64) + ) + assert labels.dtype == np.int64 + assert labels.shape == (1, 1, 5) + + +def test_watershed_3d_exercises_all_axis_neighbours(): + # A single marker at the origin of a (2,2,2) volume should reach every + # cell when the heightmap is uniform — every axis-aligned neighbour is + # visited. + image = np.zeros((2, 2, 2), dtype=np.float32) + markers = np.zeros((2, 2, 2), dtype=np.uint32) + markers[0, 0, 0] = 7 + + labels = bic.segmentation.watershed(image, markers) + + np.testing.assert_array_equal(labels, np.full((2, 2, 2), 7, dtype=np.uint32)) + + +def test_watershed_returns_zero_where_no_marker_reaches(): + image = np.zeros((1, 3), dtype=np.float32) + markers = np.zeros((1, 3), dtype=np.uint32) + + labels = bic.segmentation.watershed(image, markers) + + np.testing.assert_array_equal(labels, np.zeros((1, 3), dtype=np.uint32)) + + +def test_watershed_mask_excludes_pixels_from_flooding(): + image = np.zeros((1, 5), dtype=np.float32) + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint64) + mask = np.array([[True, True, False, True, True]]) + + labels = bic.segmentation.watershed(image, markers, mask=mask) + + # The middle column is masked out and stays 0 — it also separates the two + # flood fronts so each marker stays on its own side. + np.testing.assert_array_equal( + labels, np.array([[1, 1, 0, 2, 2]], dtype=np.uint64) + ) + + +def test_watershed_marker_under_masked_pixel_is_ignored(): + image = np.zeros((1, 3), dtype=np.float32) + markers = np.array([[0, 7, 0]], dtype=np.int32) + mask = np.array([[True, False, True]]) + + labels = bic.segmentation.watershed(image, markers, mask=mask) + + # The only marker sits under a False mask pixel, so nothing floods. + np.testing.assert_array_equal(labels, np.zeros((1, 3), dtype=np.int32)) + + +@pytest.mark.parametrize("image_dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("marker_dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_watershed_dtype_matrix(image_dtype, marker_dtype): + image = np.array([[0.0, 2.0, 9.0, 1.0, 0.0]], dtype=image_dtype) + markers = np.array([[1, 0, 0, 0, 2]], dtype=marker_dtype) + + labels = bic.segmentation.watershed(image, markers) + + assert labels.dtype == np.dtype(marker_dtype) + assert labels.shape == image.shape + np.testing.assert_array_equal( + labels, np.array([[1, 1, 2, 2, 2]], dtype=marker_dtype) + ) + + +def test_watershed_is_deterministic(): + rng = np.random.default_rng(0) + image = rng.random((10, 12), dtype=np.float32) + markers = np.zeros((10, 12), dtype=np.uint32) + markers[0, 0] = 1 + markers[9, 11] = 2 + + first = bic.segmentation.watershed(image, markers) + second = bic.segmentation.watershed(image, markers) + + np.testing.assert_array_equal(first, second) + + +def test_watershed_rejects_1d_image(): + image = np.zeros(5, dtype=np.float32) + markers = np.zeros(5, dtype=np.uint32) + + with pytest.raises(ValueError, match="ndim"): + bic.segmentation.watershed(image, markers) + + +def test_watershed_rejects_mismatched_markers_shape(): + image = np.zeros((3, 3), dtype=np.float32) + markers = np.zeros((3, 4), dtype=np.uint32) + + with pytest.raises(ValueError, match="markers shape"): + bic.segmentation.watershed(image, markers) + + +def test_watershed_rejects_mismatched_mask_shape(): + image = np.zeros((3, 3), dtype=np.float32) + markers = np.zeros((3, 3), dtype=np.uint32) + mask = np.ones((3, 4), dtype=bool) + + with pytest.raises(ValueError, match="mask shape"): + bic.segmentation.watershed(image, markers, mask=mask) + + +def test_watershed_rejects_non_bool_mask(): + image = np.zeros((3, 3), dtype=np.float32) + markers = np.zeros((3, 3), dtype=np.uint32) + mask = np.ones((3, 3), dtype=np.uint8) + + with pytest.raises(TypeError, match="mask must have dtype bool"): + bic.segmentation.watershed(image, markers, mask=mask) + + +def test_watershed_rejects_unsupported_image_dtype(): + image = np.zeros((3, 3), dtype=np.int16) + markers = np.zeros((3, 3), dtype=np.uint32) + + with pytest.raises(TypeError, match="image must have one of dtypes"): + bic.segmentation.watershed(image, markers) + + +def test_watershed_rejects_unsupported_markers_dtype(): + image = np.zeros((3, 3), dtype=np.float32) + markers = np.zeros((3, 3), dtype=np.uint8) + + with pytest.raises(TypeError, match="markers must have one of dtypes"): + bic.segmentation.watershed(image, markers) diff --git a/tests/segmentation/test_watershed_from_affinities.py b/tests/segmentation/test_watershed_from_affinities.py new file mode 100644 index 0000000..2b5d152 --- /dev/null +++ b/tests/segmentation/test_watershed_from_affinities.py @@ -0,0 +1,299 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _make_2d_ridge_neg(): + # Shape (2, 1, 5); negative-direction offsets [(-1, 0), (0, -1)]. + # Channel 1 (x-axis): aff[1, 0, x] = edge between (0, x-1) and (0, x). + aff = np.zeros((2, 1, 5), dtype=np.float32) + aff[1, 0, 1] = 0.9 # edge 0-1: strong + aff[1, 0, 2] = 0.5 # edge 1-2 + aff[1, 0, 3] = 0.1 # edge 2-3: weak ridge + aff[1, 0, 4] = 0.9 # edge 3-4: strong + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint32) + return aff, markers + + +def _make_2d_ridge_pos(): + # Shape (2, 1, 5); positive-direction offsets [(1, 0), (0, 1)]. + # Channel 1 (x-axis): aff[1, 0, x] = edge between (0, x) and (0, x+1). + aff = np.zeros((2, 1, 5), dtype=np.float32) + aff[1, 0, 0] = 0.9 # edge 0-1 + aff[1, 0, 1] = 0.5 # edge 1-2 + aff[1, 0, 2] = 0.1 # edge 2-3 + aff[1, 0, 3] = 0.9 # edge 3-4 + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint32) + return aff, markers + + +def test_watershed_from_affinities_2d_negative_direction(): + aff, markers = _make_2d_ridge_neg() + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1, 2, 2]], dtype=np.uint32) + ) + assert labels.dtype == np.uint32 + assert labels.shape == (1, 5) + + +def test_watershed_from_affinities_2d_positive_direction(): + aff, markers = _make_2d_ridge_pos() + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(1, 0), (0, 1)], markers=markers, + ) + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1, 2, 2]], dtype=np.uint32) + ) + + +def test_watershed_from_affinities_2d_permuted_offsets(): + # Same problem as the negative-direction ridge, but pass offsets in + # reversed axis order. The function must rebuild the axis-channel + # mapping; the result is unchanged. + aff_yx = np.zeros((2, 1, 5), dtype=np.float32) + aff_yx[1, 0, 1] = 0.9 + aff_yx[1, 0, 2] = 0.5 + aff_yx[1, 0, 3] = 0.1 + aff_yx[1, 0, 4] = 0.9 + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint32) + + # Swap channels: channel 0 is now the x-axis edges, channel 1 is the + # (unused) y-axis edges. Offsets describe the swap explicitly. + aff_xy = aff_yx[[1, 0]].copy() + labels = bic.segmentation.watershed_from_affinities( + aff_xy, offsets=[(0, -1), (-1, 0)], markers=markers, + ) + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1, 2, 2]], dtype=np.uint32) + ) + + +def test_watershed_from_affinities_3d_negative_direction(): + aff = np.zeros((3, 1, 1, 5), dtype=np.float32) + aff[2, 0, 0, 1] = 0.9 + aff[2, 0, 0, 2] = 0.5 + aff[2, 0, 0, 3] = 0.1 + aff[2, 0, 0, 4] = 0.9 + markers = np.array([[[1, 0, 0, 0, 2]]], dtype=np.int64) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0, 0), (0, -1, 0), (0, 0, -1)], markers=markers, + ) + np.testing.assert_array_equal( + labels, np.array([[[1, 1, 1, 2, 2]]], dtype=np.int64) + ) + assert labels.dtype == np.int64 + + +def test_watershed_from_affinities_3d_positive_direction(): + aff = np.zeros((3, 1, 1, 5), dtype=np.float32) + aff[2, 0, 0, 0] = 0.9 + aff[2, 0, 0, 1] = 0.5 + aff[2, 0, 0, 2] = 0.1 + aff[2, 0, 0, 3] = 0.9 + markers = np.array([[[1, 0, 0, 0, 2]]], dtype=np.uint64) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(1, 0, 0), (0, 1, 0), (0, 0, 1)], markers=markers, + ) + np.testing.assert_array_equal( + labels, np.array([[[1, 1, 1, 2, 2]]], dtype=np.uint64) + ) + + +def test_watershed_from_affinities_floods_uniform_region(): + # Single marker in a uniform-affinity volume should claim every pixel. + aff = np.ones((2, 3, 3), dtype=np.float32) + markers = np.zeros((3, 3), dtype=np.uint32) + markers[0, 0] = 5 + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + np.testing.assert_array_equal(labels, np.full((3, 3), 5, dtype=np.uint32)) + + +def test_watershed_from_affinities_returns_zero_where_no_marker(): + aff = np.ones((2, 1, 3), dtype=np.float32) + markers = np.zeros((1, 3), dtype=np.uint32) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + np.testing.assert_array_equal(labels, np.zeros((1, 3), dtype=np.uint32)) + + +def test_watershed_from_affinities_mask_excludes_pixels(): + aff = np.ones((2, 1, 5), dtype=np.float32) + markers = np.array([[1, 0, 0, 0, 2]], dtype=np.uint64) + mask = np.array([[True, True, False, True, True]]) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, mask=mask, + ) + np.testing.assert_array_equal( + labels, np.array([[1, 1, 0, 2, 2]], dtype=np.uint64) + ) + + +def test_watershed_from_affinities_marker_under_masked_pixel_is_ignored(): + aff = np.ones((2, 1, 3), dtype=np.float32) + markers = np.array([[0, 7, 0]], dtype=np.int32) + mask = np.array([[True, False, True]]) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, mask=mask, + ) + np.testing.assert_array_equal(labels, np.zeros((1, 3), dtype=np.int32)) + + +@pytest.mark.parametrize("aff_dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("marker_dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_watershed_from_affinities_dtype_matrix(aff_dtype, marker_dtype): + aff = np.zeros((2, 1, 5), dtype=aff_dtype) + aff[1, 0, 1] = 0.9 + aff[1, 0, 2] = 0.5 + aff[1, 0, 3] = 0.1 + aff[1, 0, 4] = 0.9 + markers = np.array([[1, 0, 0, 0, 2]], dtype=marker_dtype) + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + assert labels.dtype == np.dtype(marker_dtype) + assert labels.shape == (1, 5) + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1, 2, 2]], dtype=marker_dtype) + ) + + +def test_watershed_from_affinities_is_deterministic(): + rng = np.random.default_rng(0) + aff = rng.random((2, 8, 10), dtype=np.float32) + markers = np.zeros((8, 10), dtype=np.uint32) + markers[0, 0] = 1 + markers[7, 9] = 2 + + first = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + second = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + np.testing.assert_array_equal(first, second) + + +def test_watershed_from_affinities_rejects_wrong_ndim(): + aff = np.ones((5,), dtype=np.float32) + markers = np.zeros((5,), dtype=np.uint32) + with pytest.raises(ValueError, match="ndim"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1,)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_channel_mismatch(): + aff = np.ones((3, 4, 4), dtype=np.float32) # 3 channels for 2D spatial + markers = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="channel count"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1), (-1, 0)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_non_nearest_neighbour_offset(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="nearest-neighbour"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-2, 0), (0, -1)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_mixed_signs(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="same sign"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, 1)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_duplicate_axis(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="more than once"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (-1, 0)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_markers_shape_mismatch(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 5), dtype=np.uint32) + with pytest.raises(ValueError, match="markers shape"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_mask_shape_mismatch(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint32) + mask = np.ones((4, 5), dtype=bool) + with pytest.raises(ValueError, match="mask shape"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, mask=mask, + ) + + +def test_watershed_from_affinities_rejects_non_bool_mask(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint32) + mask = np.ones((4, 4), dtype=np.uint8) + with pytest.raises(TypeError, match="mask must have dtype bool"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, mask=mask, + ) + + +def test_watershed_from_affinities_rejects_unsupported_aff_dtype(): + aff = np.ones((2, 4, 4), dtype=np.int16) + markers = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(TypeError, match="affinities must have one of dtypes"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + + +def test_watershed_from_affinities_rejects_unsupported_marker_dtype(): + aff = np.ones((2, 4, 4), dtype=np.float32) + markers = np.zeros((4, 4), dtype=np.uint8) + with pytest.raises(TypeError, match="markers must have one of dtypes"): + bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + + +def test_watershed_from_affinities_full_coverage_on_random_input(): + # Sanity invariant on random input: every pixel ends up with one of the + # marker labels (no zeros), and the set of output labels is exactly the + # marker label set. (Exact agreement with the node-based watershed is + # not expected — the algorithms have different priority semantics.) + rng = np.random.default_rng(42) + aff = rng.random((2, 32, 32), dtype=np.float32) + markers = np.zeros((32, 32), dtype=np.uint32) + markers[0, 0] = 1 + markers[31, 31] = 2 + markers[0, 31] = 3 + markers[31, 0] = 4 + + labels = bic.segmentation.watershed_from_affinities( + aff, offsets=[(-1, 0), (0, -1)], markers=markers, + ) + assert set(np.unique(labels).tolist()) == {1, 2, 3, 4} + assert (labels > 0).all()