From 930148cf9d6682c069383fccd4493e63204f9773 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Tue, 12 May 2026 13:50:36 +0200 Subject: [PATCH 1/4] Implement structure and first functionality --- .github/workflows/tests.yml | 20 ++ .github/workflows/wheels.yml | 22 ++ .gitignore | 26 ++ AGENTS.md | 454 ++++++++++++++++++++++++++++ CMakeLists.txt | 22 ++ README.md | 35 ++- include/bioimage_cpp/array_view.hxx | 30 ++ include/bioimage_cpp/take_dict.hxx | 41 +++ pyproject.toml | 34 +++ src/bindings/module.cxx | 10 + src/bindings/utils.cxx | 113 +++++++ src/bindings/utils.hxx | 9 + src/bioimage_cpp/__init__.py | 6 + src/bioimage_cpp/_version.py | 1 + src/bioimage_cpp/utils.py | 69 +++++ src/cpp/take_dict.cxx | 32 ++ tests/test_utils_take_dict.py | 75 +++++ 17 files changed, 998 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/tests.yml create mode 100644 .github/workflows/wheels.yml create mode 100644 .gitignore create mode 100644 AGENTS.md create mode 100644 CMakeLists.txt create mode 100644 include/bioimage_cpp/array_view.hxx create mode 100644 include/bioimage_cpp/take_dict.hxx create mode 100644 pyproject.toml create mode 100644 src/bindings/module.cxx create mode 100644 src/bindings/utils.cxx create mode 100644 src/bindings/utils.hxx create mode 100644 src/bioimage_cpp/__init__.py create mode 100644 src/bioimage_cpp/_version.py create mode 100644 src/bioimage_cpp/utils.py create mode 100644 src/cpp/take_dict.cxx create mode 100644 tests/test_utils_take_dict.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..5404549 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,20 @@ +name: tests + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Install package + run: python -m pip install -v ".[test]" + - name: Test import + run: python -c "import bioimage_cpp" + - name: Run tests + run: pytest -q diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml new file mode 100644 index 0000000..94e2103 --- /dev/null +++ b/.github/workflows/wheels.yml @@ -0,0 +1,22 @@ +name: wheels + +on: + workflow_dispatch: + release: + types: [published] + +jobs: + wheels: + name: ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-13, macos-14, windows-latest] + steps: + - uses: actions/checkout@v4 + - uses: pypa/cibuildwheel@v2.18 + env: + CIBW_BUILD: "cp310-* cp311-* cp312-*" + CIBW_TEST_REQUIRES: pytest + CIBW_TEST_COMMAND: "python -c \"import bioimage_cpp\" && pytest {project}/tests" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..983b472 --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +# Python caches +__pycache__/ +*.py[cod] +*$py.class + +# Python packaging/build outputs +*.egg-info/ +build/ +dist/ +wheelhouse/ + +# scikit-build-core / CMake outputs +_skbuild/ +CMakeFiles/ +CMakeCache.txt +cmake_install.cmake +compile_commands.json + +# Test and coverage artifacts +.pytest_cache/ +.coverage +htmlcov/ + +# Local virtual environments +.venv/ +venv/ diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..5279916 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,454 @@ +# AGENTS.md — bioimage-cpp + +This repository implements `bioimage-cpp`: a small C++/Python package for bioimage-analysis algorithms that are currently hard to install from existing C++-heavy libraries. The primary goal is to provide reliable PyPI wheels with a minimal dependency footprint. + +The package should be designed as a focused implementation of selected bioimage-analysis functionality, not as a compatibility layer for older libraries and not as a full reimplementation of large libraries such as `nifty` or `vigra`. + +## Core principles + +- Keep the package easy to build and easy to ship on PyPI. +- Prefer simple, dependency-light C++20 over large template libraries or external binary dependencies. +- Avoid dependencies that complicate wheels unless they are absolutely necessary. +- Do not introduce a C++ image container or multidimensional-array framework as a public dependency. +- Treat NumPy arrays as the Python-side data model and lightweight C++ array views as the C++-side data model. +- Keep I/O outside the C++ core. File formats such as TIFF, HDF5, zarr, N5, OME-NGFF, etc. should be handled by existing Python libraries, not by this package. +- Prefer explicit, narrow APIs over broad compatibility with old libraries. +- Keep the C++ algorithmic core independent of Python, NumPy, and nanobind. Nanobind should only be used in the binding layer. + +## Dependencies + +### Preferred dependencies + +- C++20 standard library. +- `nanobind` for Python bindings. +- `numpy` for array inputs and outputs. +- `scikit-build-core` as the Python build backend. +- `cmake` for configuring the C++ build. +- `pytest` for tests. +- `cibuildwheel` for wheel builds in CI. + +### Avoid in the core package + +Do not add these dependencies to the core package unless there is a strong reason and the wheel-building consequences have been considered carefully: + +- `xtensor` +- `Boost.Python` +- `pybind11` +- `vigra` +- `nifty` +- HDF5 +- z5 / N5 libraries +- libtiff, libpng, libjpeg, OpenEXR +- FFTW +- solver libraries with nontrivial binary dependencies +- conda-only dependencies +- OpenMP, unless the packaging implications are handled explicitly + +If optional functionality really requires a heavy dependency, isolate it behind an optional extension or a separate package. The default PyPI wheel should remain small and robust. + +## Suggested repository layout + +Use a `src` layout for the Python package and keep C++ headers, C++ implementation files, and Python bindings clearly separated. + +```text +bioimage-cpp/ +├── AGENTS.md +├── CMakeLists.txt +├── pyproject.toml +├── README.md +├── include/ +│ └── bioimage_cpp/ +│ ├── array_view.hxx +│ ├── dtype_dispatch.hxx +│ ├── errors.hxx +│ └── ... +├── src/ +│ ├── bioimage_cpp/ +│ │ ├── __init__.py +│ │ ├── _version.py +│ │ └── ... +│ ├── cpp/ +│ │ └── ... +│ └── bindings/ +│ ├── module.cxx +│ ├── array_conversions.cxx +│ └── ... +├── tests/ +│ ├── test_*.py +│ └── reference/ +│ └── ... +└── .github/ + └── workflows/ + ├── tests.yml + └── wheels.yml +``` + +The exact file names may change, but the separation should remain: + +- `include/bioimage_cpp/`: public or semi-public C++ headers. +- `src/cpp/`: C++ implementations. +- `src/bindings/`: Python binding code only. +- `src/bioimage_cpp/`: Python package code. +- `tests/`: Python-level tests and, where useful, C++ behavior tests exposed through Python. + +## Build system + +Use `scikit-build-core` with CMake. Avoid `setup.py`-based builds. + +A minimal `pyproject.toml` should look conceptually like this: + +```toml +[build-system] +requires = ["scikit-build-core", "nanobind"] +build-backend = "scikit_build_core.build" + +[project] +name = "bioimage-cpp" +requires-python = ">=3.10" +dependencies = ["numpy"] + +[tool.scikit-build] +cmake.version = ">=3.21" +wheel.packages = ["src/bioimage_cpp"] +``` + +The package name on PyPI should be `bioimage-cpp`. The import name should be `bioimage_cpp`. + +If support for Python 3.10 and 3.11 is dropped later, the project may switch to Python >=3.12 and nanobind stable-ABI wheels. Do not enable nanobind `STABLE_ABI` while the package still supports Python <3.12. + +## CMake guidelines + +- Build one main Python extension module, for example `bioimage_cpp._core`. +- Use C++20 consistently. +- Use nanobind, not pybind11. +- Keep CMake simple and portable. +- Avoid platform-specific compiler flags unless they are guarded carefully. +- Do not assume conda is available. +- Do not assume system libraries such as HDF5, TIFF, or Boost are available. +- Do not enable native architecture flags such as `-march=native` in wheel builds. +- Use hidden symbol visibility where appropriate. +- Keep the extension self-contained. + +A typical CMake structure should be: + +```cmake +cmake_minimum_required(VERSION 3.21) +project(bioimage_cpp LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) +find_package(nanobind CONFIG REQUIRED) + +nanobind_add_module(_core + NB_STATIC + src/bindings/module.cxx + src/bindings/array_conversions.cxx + src/cpp/some_algorithm.cxx +) + +target_include_directories(_core PRIVATE include) +target_compile_features(_core PRIVATE cxx_std_20) +``` + +If the minimum Python version is raised to 3.12 or later, stable-ABI builds can be considered: + +```cmake +nanobind_add_module(_core + STABLE_ABI + NB_STATIC + src/bindings/module.cxx + src/bindings/array_conversions.cxx + src/cpp/some_algorithm.cxx +) +``` + +Adjust this as the project grows, but avoid turning the build into a large dependency-discovery system. + +## Wheel building + +Use `cibuildwheel` in GitHub Actions for binary wheels. + +The initial target platforms should be: + +- Linux x86_64, using manylinux wheels. +- macOS x86_64. +- macOS arm64. +- Windows x86_64. + +Later targets such as Linux aarch64 can be added once the core wheels are stable. + +Important wheel-building rules: + +- Wheels must not depend on conda. +- Wheels must not require users to have a local compiler at install time. +- Avoid external shared libraries unless they are bundled correctly and legally. +- Test the installed wheel, not only the source tree. +- Test importability with `python -c "import bioimage_cpp"`. +- Test core functions on small arrays after wheel installation. +- Do not use `-march=native`, unguarded OpenMP, or platform-specific compiler assumptions in release wheels. + +## Python API design + +The public Python API should be explicit and simple. Prefer functions such as: + +```python +import bioimage_cpp as bic + +result = bic.some_function(input_array, labels, ...) +``` + +Avoid exposing complex C++ classes unless they are genuinely needed. If a C++ object is exposed, its lifetime and ownership semantics must be obvious from Python. + +Do not implement a `vigra` or `nifty` compatibility namespace. The API should be designed for `bioimage-cpp` directly. + +## Multidimensional-array handling + +### General rule + +Use NumPy arrays at the Python boundary and lightweight C++ views internally. + +Do not use `xtensor`, `vigra::MultiArray`, or another multidimensional-array framework as a core dependency. + +### C++ array view + +Use a simple internal view type that stores: + +- data pointer +- shape +- strides +- number of dimensions +- dtype handled by template dispatch + +For example: + +```cpp +#include +#include + +template +struct ArrayView { + T* data = nullptr; + std::vector shape; + std::vector strides; + + std::ptrdiff_t ndim() const { + return static_cast(shape.size()); + } +}; +``` + +For read-only inputs, use `const T*` or a separate `ConstArrayView`. + +The C++ algorithm implementation must not depend directly on Python, NumPy, or nanobind. + +### Binding-layer responsibilities + +The Python binding layer is responsible for: + +- accepting `nanobind::ndarray` inputs +- checking dimensionality +- checking or converting dtype +- checking writability for output or in-place arguments +- checking shape compatibility between inputs +- checking C-contiguity if the C++ implementation requires it +- converting nanobind array metadata into `ArrayView` +- allocating output arrays +- dispatching to the correct C++ template instantiation + +Use nanobind only at this boundary. Do not pass `nanobind::ndarray` or other nanobind types into the algorithmic core. + +A typical binding file should use a local namespace alias: + +```cpp +#include +#include + +namespace nb = nanobind; +``` + +Example pattern: + +```cpp +nb::ndarray, nb::c_contig> +some_function( + nb::ndarray, nb::c_contig> labels +) { + // Validate Python-facing arguments here. + // Convert labels to ConstArrayView. + // Release the GIL. + // Run the C++ implementation. + // Return a newly allocated output array. +} +``` + +Prefer readable, explicit binding code over overly generic binding templates. + +### Contiguity policy + +Prefer a two-level policy: + +1. Fast C++ kernels assume C-contiguous arrays. +2. Python wrappers convert inputs with `np.ascontiguousarray` when copying is acceptable. + +For large arrays, avoid hidden expensive copies where possible. Public Python functions should make the copy policy clear, for example with an argument such as: + +```python +def some_function(array, *, copy=True): + ... +``` + +or by documenting that inputs are converted to contiguous arrays. + +If an algorithm supports arbitrary strides, this should be intentional and tested. Do not accidentally support some strided cases while failing on others. + +### Shape and axis conventions + +Use NumPy conventions throughout: + +- arrays are indexed in row-major order +- dimensions are ordered as they appear in the NumPy array +- spatial axes should not be silently reordered +- coordinate outputs should be in NumPy axis order + +For an array with shape `(z, y, x)`, coordinates should be reported as `(z, y, x)` unless a function explicitly documents otherwise. + +### Dtype policy + +Do not silently accept every dtype. Each function should define supported dtypes explicitly. + +Common policies: + +- label arrays: `uint32`, `uint64`, `int32`, `int64` +- image arrays: `uint8`, `uint16`, `uint32`, `float32`, `float64` +- output statistics: usually `float64` for numeric stability, unless there is a reason to preserve dtype +- index arrays: prefer `int64` or `uint64` at the Python boundary + +Use explicit dtype dispatch in the binding layer. Avoid large, uncontrolled template instantiation matrices. + +### Memory ownership and lifetime + +- Never store raw pointers to array data beyond the duration of the function call unless the owning Python object is kept alive explicitly. +- Do not return C++ views into temporary arrays. +- Do not return references to internal buffers unless ownership is clear and tested. +- Prefer returning newly allocated arrays. +- For in-place operations, require writable arrays and document mutation clearly. +- If a nanobind array wraps memory owned by C++, attach an explicit owner/capsule so the lifetime is correct. + +### GIL handling + +For long-running C++ kernels, release the Python GIL in the binding layer: + +```cpp +{ + nb::gil_scoped_release release; + run_algorithm(...); +} +``` + +Only release the GIL after all Python-object validation and array metadata extraction have completed. Do not touch Python objects while the GIL is released. + +### Error handling + +C++ code should throw standard exceptions such as `std::invalid_argument` or `std::runtime_error`. The binding layer should allow nanobind to translate these into Python exceptions. + +Error messages should include: + +- the offending argument name +- expected dimensionality or dtype +- actual dimensionality or dtype +- expected shape relationship, if applicable + +Example: + +```text +labels must have ndim >= 2, got ndim=1 +``` + +## Testing guidelines + +Tests should run against the installed Python package. + +For each public function, test: + +- small 2D arrays +- small 3D arrays where applicable +- nontrivial labels or regions +- background label behavior +- dtype variants +- non-contiguous inputs, if accepted +- incorrect shapes +- incorrect dtypes +- empty or degenerate inputs +- deterministic behavior + +Reference tests against older libraries may be useful during development, but they are not part of the public API contract. Do not require `vigra` or `nifty` for the default test suite, wheel tests, or source builds. + +## Coding style + +### C++ + +- Use clear, boring C++20. +- Prefer small functions over large template-heavy abstractions. +- Keep algorithm code independent of nanobind. +- Use `std::vector`, `std::array`, `std::span` where available and appropriate. +- Avoid clever metaprogramming unless it clearly reduces maintenance burden. +- Use explicit integer types where overflow matters. +- Be careful with signed/unsigned conversions for shapes, strides, and labels. + +### Python + +- Keep Python wrappers thin. +- Validate user-facing arguments in Python when this improves error messages. +- Use NumPy for lightweight preprocessing only. +- Avoid adding dependencies for small convenience features. + +## Performance guidelines + +- Start with correct, portable implementations. +- Benchmark before introducing complexity. +- Prefer algorithmic improvements over build-level optimizations. +- Avoid `-march=native` in distributed wheels. +- Consider releasing the GIL for expensive kernels. +- Consider optional threading only after the single-threaded implementation is stable. +- If threading is added, make it portable and controllable by the user. + +## Documentation expectations + +The README should explain: + +- what `bioimage-cpp` is +- what it intentionally is not +- how to install it from PyPI +- how to build it from source +- minimal examples for the public API +- the dependency-light design philosophy + +Public functions should have concise docstrings documenting: + +- input shapes +- supported dtypes +- output shapes and dtypes +- copy behavior +- background-label behavior, if relevant +- axis and coordinate conventions + +## What an agent should do when modifying this repository + +When adding or changing code: + +1. Keep the build dependency-light. +2. Do not add external C++ dependencies without strong justification. +3. Keep algorithm code separate from binding code. +4. Validate arrays at the Python/binding boundary. +5. Use nanobind only in the binding layer. +6. Use internal C++ array views rather than external multidimensional-array libraries. +7. Add tests for new behavior. +8. Preserve PyPI wheel portability. +9. Prefer clear error messages over permissive behavior. +10. Do not add I/O format support to the C++ core. +11. Do not add compatibility namespaces for older libraries. + +When unsure, choose the simpler design that is easier to build, test, and ship as wheels. diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..388e439 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,22 @@ +cmake_minimum_required(VERSION 3.21) +project(bioimage_cpp LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) +find_package(nanobind CONFIG REQUIRED) + +nanobind_add_module(_core + NB_STATIC + src/bindings/module.cxx + src/bindings/utils.cxx + src/cpp/take_dict.cxx +) + +target_include_directories(_core PRIVATE include) +target_compile_features(_core PRIVATE cxx_std_20) + +install(TARGETS _core LIBRARY DESTINATION bioimage_cpp) diff --git a/README.md b/README.md index e86eea6..23c7e9b 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,35 @@ # bioimage-cpp -Stand-alone implementation of image analysis and segmentaton functionality in C++ with minimal python bindings + +Stand-alone implementation of image analysis and segmentaton functionality in C++ with minimal python bindings. + +`bioimage-cpp` is a small C++/Python package for dependency-light bioimage-analysis algorithms that are difficult to install from larger C++ libraries. It exposes NumPy-based Python APIs backed by a focused C++20 core. + +## Install + +```bash +python -m pip install bioimage-cpp +``` + +## Build from source + +```bash +python -m pip install -v ".[test]" +pytest -q +``` + +## Example + +```python +import numpy as np +import bioimage_cpp as bic + +labels = np.array([1, 3, 2, 1], dtype=np.uint64) +relabeling = {1: 10, 2: 20, 3: 30} + +out = bic.utils.take_dict(relabeling, labels) +# array([10, 30, 20, 10], dtype=uint64) +``` + +## Scope + +The project is not a compatibility layer for `nifty`, `vigra`, or other large libraries. It keeps I/O and heavy dependencies out of the C++ core; callers should use existing Python packages for file formats and pass NumPy arrays into `bioimage-cpp`. diff --git a/include/bioimage_cpp/array_view.hxx b/include/bioimage_cpp/array_view.hxx new file mode 100644 index 0000000..6a17c78 --- /dev/null +++ b/include/bioimage_cpp/array_view.hxx @@ -0,0 +1,30 @@ +#pragma once + +#include +#include + +namespace bioimage_cpp { + +template +struct ArrayView { + T *data = nullptr; + std::vector shape; + std::vector strides; + + [[nodiscard]] std::ptrdiff_t ndim() const { + return static_cast(shape.size()); + } +}; + +template +struct ConstArrayView { + const T *data = nullptr; + std::vector shape; + std::vector strides; + + [[nodiscard]] std::ptrdiff_t ndim() const { + return static_cast(shape.size()); + } +}; + +} // namespace bioimage_cpp diff --git a/include/bioimage_cpp/take_dict.hxx b/include/bioimage_cpp/take_dict.hxx new file mode 100644 index 0000000..ac1fca0 --- /dev/null +++ b/include/bioimage_cpp/take_dict.hxx @@ -0,0 +1,41 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" + +#include +#include +#include +#include +#include + +namespace bioimage_cpp { + +template +void take_dict( + const std::unordered_map &relabeling, + const ConstArrayView &to_relabel, + const ArrayView &out +) { + if (to_relabel.shape != out.shape) { + throw std::invalid_argument( + "out shape must match to_relabel shape" + ); + } + + const auto n = std::accumulate( + to_relabel.shape.begin(), + to_relabel.shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + ); + for (std::ptrdiff_t i = 0; i < n; ++i) { + const T value = to_relabel.data[i]; + const auto found = relabeling.find(value); + if (found == relabeling.end()) { + throw std::out_of_range("relabeling is missing key " + std::to_string(value)); + } + out.data[i] = found->second; + } +} + +} // namespace bioimage_cpp diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f3522ba --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,34 @@ +[build-system] +requires = ["scikit-build-core>=0.8", "nanobind>=2.0"] +build-backend = "scikit_build_core.build" + +[project] +name = "bioimage-cpp" +version = "0.1.0" +description = "Dependency-light C++ bioimage analysis algorithms with Python bindings" +readme = "README.md" +requires-python = ">=3.10" +dependencies = ["numpy"] +license = { text = "MIT" } +authors = [ + { name = "bioimage-cpp contributors" }, +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Programming Language :: C++", + "Programming Language :: C++ :: 20", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Topic :: Scientific/Engineering :: Image Processing", +] + +[project.optional-dependencies] +test = ["pytest"] + +[tool.scikit-build] +cmake.version = ">=3.21" +wheel.packages = ["src/bioimage_cpp"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx new file mode 100644 index 0000000..bc3a0be --- /dev/null +++ b/src/bindings/module.cxx @@ -0,0 +1,10 @@ +#include "utils.hxx" + +#include + +namespace nb = nanobind; + +NB_MODULE(_core, m) { + m.doc() = "C++ extension module for bioimage_cpp."; + bioimage_cpp::bindings::bind_utils(m); +} diff --git a/src/bindings/utils.cxx b/src/bindings/utils.cxx new file mode 100644 index 0000000..b51de32 --- /dev/null +++ b/src/bindings/utils.cxx @@ -0,0 +1,113 @@ +#include "utils.hxx" + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/take_dict.hxx" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +template +using InputArray = nb::ndarray; + +template +using OutputArray = nb::ndarray; + +template +std::unordered_map dict_to_map(const nb::dict &relabeling) { + std::unordered_map result; + result.reserve(static_cast(nb::len(relabeling))); + + for (auto item : relabeling) { + const T key = nb::cast(item.first); + const T value = nb::cast(item.second); + result.emplace(key, value); + } + return result; +} + +template +OutputArray take_dict_t(const nb::dict &relabeling, InputArray to_relabel) { + std::vector ndarray_shape(to_relabel.ndim()); + std::vector view_shape(to_relabel.ndim()); + for (std::size_t axis = 0; axis < to_relabel.ndim(); ++axis) { + ndarray_shape[axis] = to_relabel.shape(axis); + view_shape[axis] = static_cast(to_relabel.shape(axis)); + } + + const auto n = std::accumulate( + view_shape.begin(), + view_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + ); + auto *data = new T[static_cast(n)](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + + ConstArrayView input{ + to_relabel.data(), + view_shape, + {}, + }; + ArrayView output{ + data, + view_shape, + {}, + }; + const auto map = dict_to_map(relabeling); + + { + nb::gil_scoped_release release; + take_dict(map, input, output); + } + + return OutputArray(data, ndarray_shape.size(), ndarray_shape.data(), owner); +} + +} // namespace + +void bind_utils(nb::module_ &m) { + m.def( + "_take_dict_uint32", + &take_dict_t, + nb::arg("relabeling"), + nb::arg("to_relabel"), + "Map a contiguous uint32 array through an integer dictionary." + ); + m.def( + "_take_dict_uint64", + &take_dict_t, + nb::arg("relabeling"), + nb::arg("to_relabel"), + "Map a contiguous uint64 array through an integer dictionary." + ); + m.def( + "_take_dict_int32", + &take_dict_t, + nb::arg("relabeling"), + nb::arg("to_relabel"), + "Map a contiguous int32 array through an integer dictionary." + ); + m.def( + "_take_dict_int64", + &take_dict_t, + nb::arg("relabeling"), + nb::arg("to_relabel"), + "Map a contiguous int64 array through an integer dictionary." + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/utils.hxx b/src/bindings/utils.hxx new file mode 100644 index 0000000..82eb1ec --- /dev/null +++ b/src/bindings/utils.hxx @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace bioimage_cpp::bindings { + +void bind_utils(nanobind::module_ &m); + +} // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/__init__.py b/src/bioimage_cpp/__init__.py new file mode 100644 index 0000000..0b5ebdd --- /dev/null +++ b/src/bioimage_cpp/__init__.py @@ -0,0 +1,6 @@ +"""Dependency-light bioimage analysis algorithms backed by C++.""" + +from ._version import __version__ +from . import utils + +__all__ = ["__version__", "utils"] diff --git a/src/bioimage_cpp/_version.py b/src/bioimage_cpp/_version.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/src/bioimage_cpp/_version.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/src/bioimage_cpp/utils.py b/src/bioimage_cpp/utils.py new file mode 100644 index 0000000..798b680 --- /dev/null +++ b/src/bioimage_cpp/utils.py @@ -0,0 +1,69 @@ +"""Utility algorithms for array relabeling and related operations.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import numpy as np + +from . import _core + +_TAKE_DICT_BY_DTYPE = { + np.dtype("uint32"): _core._take_dict_uint32, + np.dtype("uint64"): _core._take_dict_uint64, + np.dtype("int32"): _core._take_dict_int32, + np.dtype("int64"): _core._take_dict_int64, +} + + +def take_dict(relabeling: Mapping[Any, Any], to_relabel: np.ndarray) -> np.ndarray: + """Map an integer array through a dictionary. + + Parameters + ---------- + relabeling: + Mapping from input labels to output labels. All keys present in + ``to_relabel`` must exist in the mapping. + to_relabel: + NumPy array with dtype ``uint32``, ``uint64``, ``int32``, or + ``int64``. Non-contiguous inputs are copied to a contiguous array + before calling the C++ kernel. + + Returns + ------- + np.ndarray + Array with the same shape and dtype as ``to_relabel``. + + Raises + ------ + TypeError + If ``relabeling`` is not a mapping or ``to_relabel`` has an + unsupported dtype. + IndexError + If a value in ``to_relabel`` is missing from ``relabeling``. + """ + if not isinstance(relabeling, Mapping): + raise TypeError("relabeling must be a mapping") + + array = np.asarray(to_relabel) + dtype = array.dtype + try: + take = _TAKE_DICT_BY_DTYPE[dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _TAKE_DICT_BY_DTYPE) + raise TypeError( + f"to_relabel must have one of dtypes ({supported}), got dtype={dtype}" + ) from error + + contiguous = np.ascontiguousarray(array) + try: + return take(dict(relabeling), contiguous) + except IndexError: + raise + except RuntimeError as error: + # nanobind translates std::out_of_range to IndexError in recent + # versions, but older builds may surface it as RuntimeError. + if "missing key" in str(error): + raise IndexError(str(error)) from error + raise diff --git a/src/cpp/take_dict.cxx b/src/cpp/take_dict.cxx new file mode 100644 index 0000000..f9dd407 --- /dev/null +++ b/src/cpp/take_dict.cxx @@ -0,0 +1,32 @@ +#include "bioimage_cpp/take_dict.hxx" + +#include +#include + +namespace bioimage_cpp { + +template void take_dict( + const std::unordered_map &, + const ConstArrayView &, + const ArrayView & +); + +template void take_dict( + const std::unordered_map &, + const ConstArrayView &, + const ArrayView & +); + +template void take_dict( + const std::unordered_map &, + const ConstArrayView &, + const ArrayView & +); + +template void take_dict( + const std::unordered_map &, + const ConstArrayView &, + const ArrayView & +); + +} // namespace bioimage_cpp diff --git a/tests/test_utils_take_dict.py b/tests/test_utils_take_dict.py new file mode 100644 index 0000000..45a7f32 --- /dev/null +++ b/tests/test_utils_take_dict.py @@ -0,0 +1,75 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_take_dict(dtype): + to_relabel = np.array([1, 3, 2, 1], dtype=dtype) + relabeling = { + dtype(1).item(): dtype(10).item(), + dtype(2).item(): dtype(20).item(), + dtype(3).item(): dtype(30).item(), + } + + out = bic.utils.take_dict(relabeling, to_relabel) + + assert out.dtype == to_relabel.dtype + np.testing.assert_array_equal(out, np.array([10, 30, 20, 10], dtype=dtype)) + + +def test_take_dict_accepts_non_contiguous_input(): + to_relabel = np.array([0, 9, 1, 9, 2], dtype=np.uint64)[::2] + + out = bic.utils.take_dict({0: 5, 1: 6, 2: 7}, to_relabel) + + assert out.flags.c_contiguous + np.testing.assert_array_equal(out, np.array([5, 6, 7], dtype=np.uint64)) + + +def test_take_dict_preserves_2d_shape(): + to_relabel = np.array([[1, 2, 3], [3, 2, 1]], dtype=np.uint32) + + out = bic.utils.take_dict({1: 10, 2: 20, 3: 30}, to_relabel) + + assert out.shape == to_relabel.shape + assert out.dtype == to_relabel.dtype + np.testing.assert_array_equal( + out, np.array([[10, 20, 30], [30, 20, 10]], dtype=np.uint32) + ) + + +def test_take_dict_preserves_3d_shape(): + to_relabel = np.array([[[1, 2], [3, 1]], [[2, 3], [1, 2]]], dtype=np.int64) + + out = bic.utils.take_dict({1: -1, 2: -2, 3: -3}, to_relabel) + + assert out.shape == to_relabel.shape + np.testing.assert_array_equal( + out, np.array([[[-1, -2], [-3, -1]], [[-2, -3], [-1, -2]]], dtype=np.int64) + ) + + +def test_take_dict_accepts_non_contiguous_2d_input(): + to_relabel = np.array([[1, 9, 2], [3, 9, 1]], dtype=np.uint64)[:, ::2] + + out = bic.utils.take_dict({1: 10, 2: 20, 3: 30}, to_relabel) + + assert out.shape == (2, 2) + assert out.flags.c_contiguous + np.testing.assert_array_equal(out, np.array([[10, 20], [30, 10]], dtype=np.uint64)) + + +def test_take_dict_rejects_missing_key(): + with pytest.raises(IndexError, match="missing key 2"): + bic.utils.take_dict({1: 10}, np.array([1, 2], dtype=np.uint32)) + + +def test_take_dict_rejects_unsupported_dtype(): + with pytest.raises(TypeError, match="to_relabel must have one of dtypes"): + bic.utils.take_dict({1: 10}, np.array([1, 2], dtype=np.uint16)) + +def test_take_dict_rejects_non_mapping(): + with pytest.raises(TypeError, match="relabeling must be a mapping"): + bic.utils.take_dict([(1, 10)], np.array([1], dtype=np.uint64)) From 60883214a79ee5b2ba972855c89d0452276ff8a7 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Tue, 12 May 2026 14:11:20 +0200 Subject: [PATCH 2/4] MWS implementation WIP --- CMakeLists.txt | 2 + README.md | 11 + .../segmentation/mutex_watershed.hxx | 228 ++++++++++++++++++ src/bindings/module.cxx | 2 + src/bindings/segmentation.cxx | 99 ++++++++ src/bindings/segmentation.hxx | 9 + src/bioimage_cpp/__init__.py | 3 +- src/bioimage_cpp/segmentation/__init__.py | 5 + .../segmentation/mutex_watershed.py | 77 ++++++ src/cpp/segmentation/mutex_watershed.cxx | 21 ++ tests/test_segmentation_mutex_watershed.py | 104 ++++++++ 11 files changed, 560 insertions(+), 1 deletion(-) create mode 100644 include/bioimage_cpp/segmentation/mutex_watershed.hxx create mode 100644 src/bindings/segmentation.cxx create mode 100644 src/bindings/segmentation.hxx create mode 100644 src/bioimage_cpp/segmentation/__init__.py create mode 100644 src/bioimage_cpp/segmentation/mutex_watershed.py create mode 100644 src/cpp/segmentation/mutex_watershed.cxx create mode 100644 tests/test_segmentation_mutex_watershed.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 388e439..b3adef2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,7 +12,9 @@ find_package(nanobind CONFIG REQUIRED) nanobind_add_module(_core NB_STATIC src/bindings/module.cxx + src/bindings/segmentation.cxx src/bindings/utils.cxx + src/cpp/segmentation/mutex_watershed.cxx src/cpp/take_dict.cxx ) diff --git a/README.md b/README.md index 23c7e9b..af49860 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,17 @@ out = bic.utils.take_dict(relabeling, labels) # array([10, 30, 20, 10], dtype=uint64) ``` +```python +affinities = np.ones((2, 4, 4), dtype=np.float32) +offsets = [[0, 1], [1, 0]] + +segmentation = bic.segmentation.mutex_watershed( + affinities, + offsets, + number_of_attractive_channels=2, +) +``` + ## Scope The project is not a compatibility layer for `nifty`, `vigra`, or other large libraries. It keeps I/O and heavy dependencies out of the C++ core; callers should use existing Python packages for file formats and pass NumPy arrays into `bioimage-cpp`. diff --git a/include/bioimage_cpp/segmentation/mutex_watershed.hxx b/include/bioimage_cpp/segmentation/mutex_watershed.hxx new file mode 100644 index 0000000..5d933be --- /dev/null +++ b/include/bioimage_cpp/segmentation/mutex_watershed.hxx @@ -0,0 +1,228 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp { + +class DisjointSets { +public: + explicit DisjointSets(const std::size_t size) : parents_(size), ranks_(size, 0) { + std::iota(parents_.begin(), parents_.end(), std::uint64_t{0}); + } + + std::uint64_t find(const std::uint64_t node) { + if (parents_[node] != node) { + parents_[node] = find(parents_[node]); + } + return parents_[node]; + } + + std::uint64_t unite_roots(std::uint64_t first, std::uint64_t second) { + if (ranks_[first] < ranks_[second]) { + std::swap(first, second); + } + parents_[second] = first; + if (ranks_[first] == ranks_[second]) { + ++ranks_[first]; + } + return first; + } + +private: + std::vector parents_; + std::vector ranks_; +}; + +using MutexStorage = std::vector>; + +inline bool check_mutex( + const std::uint64_t first, + const std::uint64_t second, + const MutexStorage &mutexes +) { + const auto &first_mutexes = mutexes[first]; + const auto &second_mutexes = mutexes[second]; + if (first_mutexes.size() < second_mutexes.size()) { + return first_mutexes.find(second) != first_mutexes.end(); + } + return second_mutexes.find(first) != second_mutexes.end(); +} + +inline void insert_mutex( + const std::uint64_t first, + const std::uint64_t second, + MutexStorage &mutexes +) { + mutexes[first].insert(second); + mutexes[second].insert(first); +} + +inline void merge_mutexes( + const std::uint64_t root_from, + const std::uint64_t root_to, + MutexStorage &mutexes +) { + auto &mutexes_from = mutexes[root_from]; + auto &mutexes_to = mutexes[root_to]; + + for (const auto other_root : mutexes_from) { + auto &other_mutexes = mutexes[other_root]; + other_mutexes.erase(root_from); + if (other_root != root_to) { + other_mutexes.insert(root_to); + mutexes_to.insert(other_root); + } + } + mutexes_to.erase(root_from); + mutexes_to.erase(root_to); + mutexes_from.clear(); +} + +inline std::vector c_order_strides(const std::vector &shape) { + std::vector strides(shape.size(), 1); + for (std::ptrdiff_t axis = static_cast(shape.size()) - 2; axis >= 0; --axis) { + strides[static_cast(axis)] = + strides[static_cast(axis + 1)] * shape[static_cast(axis + 1)]; + } + return strides; +} + +inline bool is_valid_grid_edge( + const std::uint64_t node, + const std::vector &offset, + const std::vector &shape, + const std::vector &strides +) { + for (std::size_t axis = 0; axis < shape.size(); ++axis) { + const auto coord = static_cast(node / static_cast(strides[axis])) % shape[axis]; + const auto neighbor = coord + offset[axis]; + if (neighbor < 0 || neighbor >= shape[axis]) { + return false; + } + } + return true; +} + +template +void mutex_watershed_grid( + const ConstArrayView &affinities, + const std::vector> &offsets, + const std::size_t number_of_attractive_channels, + const ArrayView &out +) { + if (affinities.ndim() != 3 && affinities.ndim() != 4) { + throw std::invalid_argument( + "affinities must have shape (channels, y, x) or (channels, z, y, x), got ndim=" + + std::to_string(affinities.ndim()) + ); + } + if (offsets.empty()) { + throw std::invalid_argument("offsets must not be empty"); + } + + const auto number_of_channels = static_cast(affinities.shape[0]); + const auto spatial_ndim = static_cast(affinities.ndim() - 1); + if (offsets.size() != number_of_channels) { + throw std::invalid_argument( + "offsets length must match affinities channel count, got offsets length=" + + std::to_string(offsets.size()) + ", channels=" + std::to_string(number_of_channels) + ); + } + if (number_of_attractive_channels > number_of_channels) { + throw std::invalid_argument("number_of_attractive_channels must be <= number of channels"); + } + for (const auto &offset : offsets) { + if (offset.size() != spatial_ndim) { + throw std::invalid_argument( + "each offset must have length matching the spatial ndim, got spatial ndim=" + + std::to_string(spatial_ndim) + ); + } + } + + std::vector spatial_shape( + affinities.shape.begin() + 1, + affinities.shape.end() + ); + if (out.shape != spatial_shape) { + throw std::invalid_argument("out shape must match affinities spatial shape"); + } + + const auto number_of_nodes = static_cast(std::accumulate( + spatial_shape.begin(), + spatial_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + )); + const auto spatial_strides = c_order_strides(spatial_shape); + + std::vector offset_strides(number_of_channels, 0); + for (std::size_t channel = 0; channel < number_of_channels; ++channel) { + for (std::size_t axis = 0; axis < spatial_ndim; ++axis) { + offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis]; + } + } + + const auto number_of_edges = number_of_nodes * number_of_channels; + std::vector edge_order(static_cast(number_of_edges)); + std::iota(edge_order.begin(), edge_order.end(), std::uint64_t{0}); + std::sort(edge_order.begin(), edge_order.end(), [&](const std::uint64_t first, const std::uint64_t second) { + const T first_weight = affinities.data[first]; + const T second_weight = affinities.data[second]; + if (first_weight == second_weight) { + return first < second; + } + return first_weight > second_weight; + }); + + DisjointSets sets(static_cast(number_of_nodes)); + MutexStorage mutexes(static_cast(number_of_nodes)); + + for (const auto edge_id : edge_order) { + const auto channel = static_cast(edge_id / number_of_nodes); + const auto u = edge_id % number_of_nodes; + if (!is_valid_grid_edge(u, offsets[channel], spatial_shape, spatial_strides)) { + continue; + } + + const auto v_signed = static_cast(u) + static_cast(offset_strides[channel]); + const auto v = static_cast(v_signed); + std::uint64_t root_u = sets.find(u); + std::uint64_t root_v = sets.find(v); + if (root_u == root_v || check_mutex(root_u, root_v, mutexes)) { + continue; + } + + const bool is_mutex_edge = channel >= number_of_attractive_channels; + if (is_mutex_edge) { + insert_mutex(root_u, root_v, mutexes); + } else { + const auto new_root = sets.unite_roots(root_u, root_v); + const auto old_root = (new_root == root_u) ? root_v : root_u; + merge_mutexes(old_root, new_root, mutexes); + } + } + + std::unordered_map label_map; + for (std::uint64_t node = 0; node < number_of_nodes; ++node) { + const auto root = sets.find(node); + auto found = label_map.find(root); + if (found == label_map.end()) { + const auto next_label = static_cast(label_map.size() + 1); + found = label_map.emplace(root, next_label).first; + } + out.data[node] = found->second; + } +} + +} // namespace bioimage_cpp diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx index bc3a0be..ac1230e 100644 --- a/src/bindings/module.cxx +++ b/src/bindings/module.cxx @@ -1,3 +1,4 @@ +#include "segmentation.hxx" #include "utils.hxx" #include @@ -6,5 +7,6 @@ namespace nb = nanobind; NB_MODULE(_core, m) { m.doc() = "C++ extension module for bioimage_cpp."; + bioimage_cpp::bindings::bind_segmentation(m); bioimage_cpp::bindings::bind_utils(m); } diff --git a/src/bindings/segmentation.cxx b/src/bindings/segmentation.cxx new file mode 100644 index 0000000..7416ab4 --- /dev/null +++ b/src/bindings/segmentation.cxx @@ -0,0 +1,99 @@ +#include "segmentation.hxx" + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/segmentation/mutex_watershed.hxx" + +#include +#include + +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +template +using AffinityArray = nb::ndarray; + +using LabelArray = nb::ndarray; + +template +LabelArray mutex_watershed_grid_t( + AffinityArray affinities, + const std::vector> &offsets, + const std::size_t number_of_attractive_channels +) { + std::vector affinity_shape(affinities.ndim()); + for (std::size_t axis = 0; axis < affinities.ndim(); ++axis) { + affinity_shape[axis] = static_cast(affinities.shape(axis)); + } + + std::vector label_shape; + label_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0); + std::vector label_view_shape; + label_view_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0); + for (std::size_t axis = 1; axis < affinities.ndim(); ++axis) { + label_shape.push_back(affinities.shape(axis)); + label_view_shape.push_back(static_cast(affinities.shape(axis))); + } + + const auto number_of_nodes = std::accumulate( + label_view_shape.begin(), + label_view_shape.end(), + std::ptrdiff_t{1}, + [](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; } + ); + auto *data = new std::uint64_t[static_cast(number_of_nodes)](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + + ConstArrayView affinities_view{ + affinities.data(), + affinity_shape, + {}, + }; + ArrayView out_view{ + data, + label_view_shape, + {}, + }; + + { + nb::gil_scoped_release release; + mutex_watershed_grid( + affinities_view, + offsets, + number_of_attractive_channels, + out_view + ); + } + + return LabelArray(data, label_shape.size(), label_shape.data(), owner); +} + +} // namespace + +void bind_segmentation(nb::module_ &m) { + m.def( + "_mutex_watershed_grid_float32", + &mutex_watershed_grid_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("number_of_attractive_channels"), + "Run mutex watershed on a 2D or 3D image-derived grid graph with float32 affinities." + ); + m.def( + "_mutex_watershed_grid_float64", + &mutex_watershed_grid_t, + nb::arg("affinities"), + nb::arg("offsets"), + nb::arg("number_of_attractive_channels"), + "Run mutex watershed on a 2D or 3D image-derived grid graph with float64 affinities." + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/segmentation.hxx b/src/bindings/segmentation.hxx new file mode 100644 index 0000000..8502e84 --- /dev/null +++ b/src/bindings/segmentation.hxx @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace bioimage_cpp::bindings { + +void bind_segmentation(nanobind::module_ &m); + +} // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/__init__.py b/src/bioimage_cpp/__init__.py index 0b5ebdd..ba81546 100644 --- a/src/bioimage_cpp/__init__.py +++ b/src/bioimage_cpp/__init__.py @@ -1,6 +1,7 @@ """Dependency-light bioimage analysis algorithms backed by C++.""" from ._version import __version__ +from . import segmentation from . import utils -__all__ = ["__version__", "utils"] +__all__ = ["__version__", "segmentation", "utils"] diff --git a/src/bioimage_cpp/segmentation/__init__.py b/src/bioimage_cpp/segmentation/__init__.py new file mode 100644 index 0000000..cf4b303 --- /dev/null +++ b/src/bioimage_cpp/segmentation/__init__.py @@ -0,0 +1,5 @@ +"""Segmentation algorithms.""" + +from .mutex_watershed import mutex_watershed + +__all__ = ["mutex_watershed"] diff --git a/src/bioimage_cpp/segmentation/mutex_watershed.py b/src/bioimage_cpp/segmentation/mutex_watershed.py new file mode 100644 index 0000000..26f9814 --- /dev/null +++ b/src/bioimage_cpp/segmentation/mutex_watershed.py @@ -0,0 +1,77 @@ +"""Segmentation algorithms.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from .. import _core + +_MUTEX_WATERSHED_BY_DTYPE = { + np.dtype("float32"): _core._mutex_watershed_grid_float32, + np.dtype("float64"): _core._mutex_watershed_grid_float64, +} + + +def mutex_watershed( + affinities: np.ndarray, + offsets: Sequence[Sequence[int]], + number_of_attractive_channels: int, +) -> np.ndarray: + """Run mutex watershed on a 2D or 3D image-derived grid graph. + + Parameters + ---------- + affinities: + Array with shape ``(channels, y, x)`` for 2D data or + ``(channels, z, y, x)`` for 3D data. Supported dtypes are + ``float32`` and ``float64``. Non-contiguous arrays are copied. + offsets: + One offset per channel, in NumPy axis order. Each offset has length 2 + for 2D affinities or length 3 for 3D affinities. + number_of_attractive_channels: + The first this many affinity channels are attractive merge edges. The + remaining channels are mutex edges. + + Returns + ------- + np.ndarray + Consecutive 1-based ``uint64`` segmentation labels with shape + ``affinities.shape[1:]``. + """ + array = np.asarray(affinities) + if array.ndim not in (3, 4): + raise ValueError( + "affinities must have shape (channels, y, x) or " + f"(channels, z, y, x), got ndim={array.ndim}" + ) + + dtype = array.dtype + try: + run = _MUTEX_WATERSHED_BY_DTYPE[dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in _MUTEX_WATERSHED_BY_DTYPE) + raise TypeError( + f"affinities must have one of dtypes ({supported}), got dtype={dtype}" + ) from error + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + spatial_ndim = array.ndim - 1 + if len(normalized_offsets) != array.shape[0]: + raise ValueError( + "offsets length must match affinities channel count, got " + f"offsets length={len(normalized_offsets)}, channels={array.shape[0]}" + ) + if any(len(offset) != spatial_ndim for offset in normalized_offsets): + raise ValueError( + "each offset must have length matching the spatial ndim, got " + f"spatial ndim={spatial_ndim}" + ) + if number_of_attractive_channels < 0: + raise ValueError("number_of_attractive_channels must be non-negative") + if number_of_attractive_channels > array.shape[0]: + raise ValueError("number_of_attractive_channels must be <= number of channels") + + contiguous = np.ascontiguousarray(array) + return run(contiguous, normalized_offsets, number_of_attractive_channels) diff --git a/src/cpp/segmentation/mutex_watershed.cxx b/src/cpp/segmentation/mutex_watershed.cxx new file mode 100644 index 0000000..135a556 --- /dev/null +++ b/src/cpp/segmentation/mutex_watershed.cxx @@ -0,0 +1,21 @@ +#include "bioimage_cpp/segmentation/mutex_watershed.hxx" + +#include + +namespace bioimage_cpp { + +template void mutex_watershed_grid( + const ConstArrayView &, + const std::vector> &, + std::size_t, + const ArrayView & +); + +template void mutex_watershed_grid( + const ConstArrayView &, + const std::vector> &, + std::size_t, + const ArrayView & +); + +} // namespace bioimage_cpp diff --git a/tests/test_segmentation_mutex_watershed.py b/tests/test_segmentation_mutex_watershed.py new file mode 100644 index 0000000..c29bb0f --- /dev/null +++ b/tests/test_segmentation_mutex_watershed.py @@ -0,0 +1,104 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_mutex_watershed_2d_attractive_edges_merge_all_pixels(): + affinities = np.ones((2, 3, 4), dtype=np.float32) + offsets = [[0, 1], [1, 0]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=2 + ) + + assert labels.dtype == np.uint64 + assert labels.shape == (3, 4) + np.testing.assert_array_equal(labels, np.ones((3, 4), dtype=np.uint64)) + + +def test_mutex_watershed_2d_mutex_edge_blocks_lower_attractive_edge(): + affinities = np.array([[[0.5, 0.0]], [[0.9, 0.0]]], dtype=np.float64) + offsets = [[0, 1], [0, 1]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + np.testing.assert_array_equal(labels, np.array([[1, 2]], dtype=np.uint64)) + + +def test_mutex_watershed_2d_higher_attractive_edge_wins_before_mutex_edge(): + affinities = np.array([[[0.9, 0.0]], [[0.5, 0.0]]], dtype=np.float32) + offsets = [[0, 1], [0, 1]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + np.testing.assert_array_equal(labels, np.array([[1, 1]], dtype=np.uint64)) + + +def test_mutex_watershed_2d_respects_grid_boundaries(): + affinities = np.ones((1, 2, 3), dtype=np.float32) + offsets = [[0, 1]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1], [2, 2, 2]], dtype=np.uint64) + ) + + +def test_mutex_watershed_2d_accepts_negative_offsets(): + affinities = np.ones((1, 2, 3), dtype=np.float32) + offsets = [[0, -1]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + + np.testing.assert_array_equal( + labels, np.array([[1, 1, 1], [2, 2, 2]], dtype=np.uint64) + ) + + +def test_mutex_watershed_3d_attractive_edges_merge_all_pixels(): + affinities = np.ones((3, 2, 2, 2), dtype=np.float32) + offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=3 + ) + + assert labels.shape == (2, 2, 2) + np.testing.assert_array_equal(labels, np.ones((2, 2, 2), dtype=np.uint64)) + + +def test_mutex_watershed_rejects_wrong_affinity_dtype(): + with pytest.raises(TypeError, match="affinities must have one of dtypes"): + bic.segmentation.mutex_watershed( + np.ones((1, 2, 2), dtype=np.uint8), + [[0, 1]], + number_of_attractive_channels=1, + ) + + +def test_mutex_watershed_rejects_wrong_offset_length(): + with pytest.raises(ValueError, match="each offset must have length"): + bic.segmentation.mutex_watershed( + np.ones((1, 2, 2), dtype=np.float32), + [[0, 1, 0]], + number_of_attractive_channels=1, + ) + + +def test_mutex_watershed_rejects_wrong_number_of_offsets(): + with pytest.raises(ValueError, match="offsets length must match"): + bic.segmentation.mutex_watershed( + np.ones((2, 2, 2), dtype=np.float32), + [[0, 1]], + number_of_attractive_channels=1, + ) From c564987544fe860450f99a466d0ae7c408c4efc0 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 14 May 2026 12:01:00 +0200 Subject: [PATCH 3/4] Add mutex watershed tests and development scripts --- .gitignore | 3 + .../_mutex_watershed_equivalence.py | 227 ++++++++++++++++++ .../segmentation/check_mutex_watershed_2d.py | 40 +++ .../segmentation/check_mutex_watershed_3d.py | 34 +++ .../test_mutex_watershed.py} | 0 5 files changed, 304 insertions(+) create mode 100644 development/segmentation/_mutex_watershed_equivalence.py create mode 100644 development/segmentation/check_mutex_watershed_2d.py create mode 100644 development/segmentation/check_mutex_watershed_3d.py rename tests/{test_segmentation_mutex_watershed.py => segmentation/test_mutex_watershed.py} (100%) diff --git a/.gitignore b/.gitignore index 983b472..1ce03dc 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,6 @@ htmlcov/ # Local virtual environments .venv/ venv/ + +# Data +*.h5 diff --git a/development/segmentation/_mutex_watershed_equivalence.py b/development/segmentation/_mutex_watershed_equivalence.py new file mode 100644 index 0000000..ea51b3f --- /dev/null +++ b/development/segmentation/_mutex_watershed_equivalence.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_DATA_PREFIX = PROJECT_ROOT / "examples" / "segmentation" / "isbi-data-" + + +def load_problem(data_prefix: Path | str = DEFAULT_DATA_PREFIX): + from elf.segmentation.utils import load_mutex_watershed_problem + + affinities, offsets = load_mutex_watershed_problem(prefix=str(data_prefix)) + return np.ascontiguousarray(affinities), [tuple(offset) for offset in offsets] + + +def prepare_2d_problem( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + z: int, + yx_shape: tuple[int, int], +): + 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] + return np.ascontiguousarray(cropped), offsets_2d, 2 + + +def prepare_3d_problem( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + zyx_shape: tuple[int, int, int], +): + z, y, x = zyx_shape + cropped = affinities[:, :z, :y, :x] + return np.ascontiguousarray(cropped), offsets, 3 + + +def run_bioimage_cpp( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, +) -> np.ndarray: + import bioimage_cpp as bic + + affs = affinities.copy() + affs[:number_of_attractive_channels] *= -1 + affs[:number_of_attractive_channels] += 1 + return bic.segmentation.mutex_watershed( + affs, + offsets, + number_of_attractive_channels=number_of_attractive_channels, + ) + + +def run_affogato_reference( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, +) -> np.ndarray: + from elf.segmentation.mutex_watershed import mutex_watershed + + spatial_ndim = len(offsets[0]) + if number_of_attractive_channels != spatial_ndim: + raise ValueError( + "the elf mutex_watershed wrapper assumes one attractive channel per " + f"spatial axis, got ndim={spatial_ndim}, attractive channels=" + f"{number_of_attractive_channels}" + ) + return mutex_watershed(affinities.copy(), offsets, strides=[1] * spatial_ndim) + + +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, + *, + max_vi: float = 1.0e-10, + max_are: float = 1.0e-10, + min_rand_index: float = 1.0 - 1.0e-10, +) -> 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)) + equivalent = ( + vi_split <= max_vi + and vi_merge <= max_vi + and adapted_rand_error <= max_are + and 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: + raise AssertionError( + "mutex watershed results differ: " + f"VI split={vi_split:.6g}, VI merge={vi_merge:.6g}, " + f"adapted rand error={adapted_rand_error:.6g}, " + f"rand index={ri:.12g}, exact labels={exact_equal}" + ) + return metrics + + +def time_function( + run: Callable[[np.ndarray, list[tuple[int, ...]], int], np.ndarray], + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + number_of_attractive_channels: int, + repeats: int, +) -> tuple[list[float], np.ndarray]: + timings = [] + result = None + for _ in range(repeats): + start = perf_counter() + result = run(affinities, offsets, number_of_attractive_channels) + timings.append(perf_counter() - start) + assert result is not None + return timings, result + + +def print_report( + *, + ndim: int, + affinities: 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") + + print(f"Mutex watershed {ndim}D equivalence check") + print(f"affinities shape: {affinities.shape}, dtype: {affinities.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']:.12g}" + ) + print(f"exact label equality: {metrics['exact_label_equality']}") + print(f"bioimage-cpp median runtime: {bic_median:.6f} s") + print(f"affogato reference median runtime: {ref_median:.6f} s") + print(f"reference / bioimage-cpp runtime ratio: {speedup:.3f}x") + + +def run_check( + *, + ndim: int, + repeats: int, + data_prefix: Path | str, + z: int, + yx_shape: tuple[int, int], + zyx_shape: tuple[int, int, int], +): + affinities, offsets = load_problem(data_prefix) + if ndim == 2: + affs, used_offsets, attractive_channels = prepare_2d_problem( + affinities, offsets, z=z, yx_shape=yx_shape + ) + elif ndim == 3: + affs, used_offsets, attractive_channels = prepare_3d_problem( + affinities, offsets, zyx_shape=zyx_shape + ) + else: + raise ValueError(f"ndim must be 2 or 3, got {ndim}") + + ref_timings, ref_seg = time_function( + run_affogato_reference, affs, used_offsets, attractive_channels, repeats + ) + bic_timings, bic_seg = time_function( + run_bioimage_cpp, affs, used_offsets, attractive_channels, repeats + ) + metrics = compare_segmentations(bic_seg, ref_seg) + print_report( + ndim=ndim, + affinities=affs, + metrics=metrics, + bic_timings=bic_timings, + ref_timings=ref_timings, + ) + + +def add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--data-prefix", + type=Path, + default=DEFAULT_DATA_PREFIX, + help=( + "Path prefix for the ISBI mutex watershed data. The loader expects " + "'test.h5' and 'train.h5' suffixes." + ), + ) + parser.add_argument( + "--repeats", + type=int, + default=3, + help="Number of timed runs for each implementation.", + ) diff --git a/development/segmentation/check_mutex_watershed_2d.py b/development/segmentation/check_mutex_watershed_2d.py new file mode 100644 index 0000000..8960639 --- /dev/null +++ b/development/segmentation/check_mutex_watershed_2d.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import argparse + +from _mutex_watershed_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and affogato mutex watershed in 2D." + ) + add_common_arguments(parser) + parser.add_argument( + "--z", + type=int, + default=0, + help="Z slice from the example affinities used for the 2D check.", + ) + parser.add_argument( + "--shape", + type=int, + nargs=2, + metavar=("Y", "X"), + default=(128, 128), + help="Spatial crop shape for the 2D check.", + ) + args = parser.parse_args() + + run_check( + ndim=2, + repeats=args.repeats, + data_prefix=args.data_prefix, + z=args.z, + yx_shape=tuple(args.shape), + zyx_shape=(0, 0, 0), + ) + + +if __name__ == "__main__": + main() diff --git a/development/segmentation/check_mutex_watershed_3d.py b/development/segmentation/check_mutex_watershed_3d.py new file mode 100644 index 0000000..45ae843 --- /dev/null +++ b/development/segmentation/check_mutex_watershed_3d.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import argparse + +from _mutex_watershed_equivalence import add_common_arguments, run_check + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare bioimage-cpp and affogato mutex watershed in 3D." + ) + add_common_arguments(parser) + parser.add_argument( + "--shape", + type=int, + nargs=3, + metavar=("Z", "Y", "X"), + default=(6, 96, 96), + help="Spatial crop shape for the 3D check.", + ) + args = parser.parse_args() + + run_check( + ndim=3, + repeats=args.repeats, + data_prefix=args.data_prefix, + z=0, + yx_shape=(0, 0), + zyx_shape=tuple(args.shape), + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_segmentation_mutex_watershed.py b/tests/segmentation/test_mutex_watershed.py similarity index 100% rename from tests/test_segmentation_mutex_watershed.py rename to tests/segmentation/test_mutex_watershed.py From b87a82e13136425f1b93665c1ea13a7a9c8098bc Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 14 May 2026 12:14:30 +0200 Subject: [PATCH 4/4] Add stride and mask support to MWS --- examples/segmentation/mutex_watershed.py | 63 ++++++++ .../segmentation/mutex_watershed.hxx | 8 + src/bindings/segmentation.cxx | 14 ++ .../segmentation/mutex_watershed.py | 152 +++++++++++++++++- src/cpp/segmentation/mutex_watershed.cxx | 2 + tests/segmentation/test_mutex_watershed.py | 75 +++++++++ 6 files changed, 313 insertions(+), 1 deletion(-) create mode 100644 examples/segmentation/mutex_watershed.py diff --git a/examples/segmentation/mutex_watershed.py b/examples/segmentation/mutex_watershed.py new file mode 100644 index 0000000..7477e48 --- /dev/null +++ b/examples/segmentation/mutex_watershed.py @@ -0,0 +1,63 @@ +import napari + +from elf.segmentation.utils import load_mutex_watershed_problem +from elf.io import open_file + + +def mws_bic(affinities, offsets): + import bioimage_cpp as bic + print("Start MWS ...") + affs = affinities.copy() + ndim = len(offsets[0]) + affs[:ndim] *= -1 + affs[:ndim] += 1 + segmentation = bic.segmentation.mutex_watershed( + affs, offsets, number_of_attractive_channels=ndim + ) + print("done ...") + return segmentation + + +def mws_affogato(affinities, offsets): + from elf.segmentation.mutex_watershed import mutex_watershed as mws + print("Start MWS ...") + affs = affinities.copy() + segmentation = mws(affs, offsets, strides=[1, 1, 1]) + print("done ...") + return segmentation + + +def _filter_2d(affinities, offsets): + chans_2d = [i for i, off in enumerate(offsets) if off[0] == 0] + affinities = affinities[chans_2d][:, 0] + offsets = [off[1:] for i, off in enumerate(offsets) if i in chans_2d] + return affinities, offsets + + +def main(): + prefix = "isbi-data-" + data_path = f"{prefix}test.h5" + affinities, offsets = load_mutex_watershed_problem(prefix=prefix) + + check_2d = True + if check_2d: + affinities, offsets = _filter_2d(affinities, offsets) + + use_reference_impl = False + if use_reference_impl: + segmentation = mws_affogato(affinities, offsets) + else: + segmentation = mws_bic(affinities, offsets) + + with open_file(data_path, "r") as f: + raw = f["raw"][0] if check_2d else f["raw"][:] + + viewer = napari.Viewer() + viewer.add_image(raw, name="raw") + viewer.add_image(affinities, name="affinities") + viewer.add_labels(segmentation, name="mws-segmentation") + napari.run() + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/segmentation/mutex_watershed.hxx b/include/bioimage_cpp/segmentation/mutex_watershed.hxx index 5d933be..b96eb71 100644 --- a/include/bioimage_cpp/segmentation/mutex_watershed.hxx +++ b/include/bioimage_cpp/segmentation/mutex_watershed.hxx @@ -116,6 +116,7 @@ inline bool is_valid_grid_edge( template void mutex_watershed_grid( const ConstArrayView &affinities, + const ConstArrayView &valid_edges, const std::vector> &offsets, const std::size_t number_of_attractive_channels, const ArrayView &out @@ -129,6 +130,9 @@ void mutex_watershed_grid( if (offsets.empty()) { throw std::invalid_argument("offsets must not be empty"); } + if (valid_edges.shape != affinities.shape) { + throw std::invalid_argument("valid_edges shape must match affinities shape"); + } const auto number_of_channels = static_cast(affinities.shape[0]); const auto spatial_ndim = static_cast(affinities.ndim() - 1); @@ -189,6 +193,10 @@ void mutex_watershed_grid( MutexStorage mutexes(static_cast(number_of_nodes)); for (const auto edge_id : edge_order) { + if (valid_edges.data[edge_id] == 0) { + continue; + } + const auto channel = static_cast(edge_id / number_of_nodes); const auto u = edge_id % number_of_nodes; if (!is_valid_grid_edge(u, offsets[channel], spatial_shape, spatial_strides)) { diff --git a/src/bindings/segmentation.cxx b/src/bindings/segmentation.cxx index 7416ab4..d37bdc4 100644 --- a/src/bindings/segmentation.cxx +++ b/src/bindings/segmentation.cxx @@ -20,11 +20,13 @@ namespace { template using AffinityArray = nb::ndarray; +using ValidEdgeArray = nb::ndarray; using LabelArray = nb::ndarray; template LabelArray mutex_watershed_grid_t( AffinityArray affinities, + ValidEdgeArray valid_edges, const std::vector> &offsets, const std::size_t number_of_attractive_channels ) { @@ -32,6 +34,10 @@ LabelArray mutex_watershed_grid_t( for (std::size_t axis = 0; axis < affinities.ndim(); ++axis) { affinity_shape[axis] = static_cast(affinities.shape(axis)); } + std::vector valid_edges_shape(valid_edges.ndim()); + for (std::size_t axis = 0; axis < valid_edges.ndim(); ++axis) { + valid_edges_shape[axis] = static_cast(valid_edges.shape(axis)); + } std::vector label_shape; label_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0); @@ -56,6 +62,11 @@ LabelArray mutex_watershed_grid_t( affinity_shape, {}, }; + ConstArrayView valid_edges_view{ + valid_edges.data(), + valid_edges_shape, + {}, + }; ArrayView out_view{ data, label_view_shape, @@ -66,6 +77,7 @@ LabelArray mutex_watershed_grid_t( nb::gil_scoped_release release; mutex_watershed_grid( affinities_view, + valid_edges_view, offsets, number_of_attractive_channels, out_view @@ -82,6 +94,7 @@ void bind_segmentation(nb::module_ &m) { "_mutex_watershed_grid_float32", &mutex_watershed_grid_t, nb::arg("affinities"), + nb::arg("valid_edges"), nb::arg("offsets"), nb::arg("number_of_attractive_channels"), "Run mutex watershed on a 2D or 3D image-derived grid graph with float32 affinities." @@ -90,6 +103,7 @@ void bind_segmentation(nb::module_ &m) { "_mutex_watershed_grid_float64", &mutex_watershed_grid_t, nb::arg("affinities"), + nb::arg("valid_edges"), nb::arg("offsets"), nb::arg("number_of_attractive_channels"), "Run mutex watershed on a 2D or 3D image-derived grid graph with float64 affinities." diff --git a/src/bioimage_cpp/segmentation/mutex_watershed.py b/src/bioimage_cpp/segmentation/mutex_watershed.py index 26f9814..d500dd0 100644 --- a/src/bioimage_cpp/segmentation/mutex_watershed.py +++ b/src/bioimage_cpp/segmentation/mutex_watershed.py @@ -18,6 +18,10 @@ def mutex_watershed( affinities: np.ndarray, offsets: Sequence[Sequence[int]], number_of_attractive_channels: int, + *, + strides: Sequence[int] | None = None, + randomized_strides: bool = False, + mask: np.ndarray | None = None, ) -> np.ndarray: """Run mutex watershed on a 2D or 3D image-derived grid graph. @@ -33,6 +37,17 @@ def mutex_watershed( number_of_attractive_channels: The first this many affinity channels are attractive merge edges. The remaining channels are mutex edges. + strides: + Optional spatial sub-sampling strides for mutex edges. Attractive + channels are always kept. If given, it must have one positive integer + per spatial dimension. + randomized_strides: + If ``True``, sub-sample mutex edges randomly with probability + ``1 / np.prod(strides)`` instead of on a regular grid. + mask: + Optional boolean foreground mask with shape ``affinities.shape[1:]``. + Edges touching ``False`` pixels are ignored and ``False`` pixels are + labeled as background ``0`` in the output. Returns ------- @@ -73,5 +88,140 @@ def mutex_watershed( if number_of_attractive_channels > array.shape[0]: raise ValueError("number_of_attractive_channels must be <= number of channels") + normalized_strides = _normalize_strides(strides, spatial_ndim, randomized_strides) + valid_edges = _compute_valid_edges( + array.shape, + normalized_offsets, + number_of_attractive_channels, + normalized_strides, + randomized_strides, + mask, + ) + contiguous = np.ascontiguousarray(array) - return run(contiguous, normalized_offsets, number_of_attractive_channels) + labels = run( + contiguous, + valid_edges, + normalized_offsets, + number_of_attractive_channels, + ) + if mask is not None: + labels[np.logical_not(np.asarray(mask))] = 0 + return labels + + +def _normalize_strides( + strides: Sequence[int] | None, + spatial_ndim: int, + randomized_strides: bool, +) -> tuple[int, ...] | None: + if strides is None: + if randomized_strides: + raise ValueError("randomized_strides requires strides") + return None + + normalized = tuple(int(stride) for stride in strides) + if len(normalized) != spatial_ndim: + raise ValueError( + "strides length must match the spatial ndim, got " + f"strides length={len(normalized)}, spatial ndim={spatial_ndim}" + ) + if any(stride <= 0 for stride in normalized): + raise ValueError("strides must contain only positive integers") + return normalized + + +def _valid_source_slices( + image_shape: tuple[int, ...], + offset: tuple[int, ...], +) -> tuple[slice, ...] | None: + slices = [] + for axis_size, step in zip(image_shape, offset, strict=True): + if step > 0: + if step >= axis_size: + return None + slices.append(slice(0, axis_size - step)) + elif step < 0: + if -step >= axis_size: + return None + slices.append(slice(-step, axis_size)) + else: + slices.append(slice(None)) + return tuple(slices) + + +def _neighbor_slices( + image_shape: tuple[int, ...], + offset: tuple[int, ...], +) -> tuple[slice, ...] | None: + slices = [] + for axis_size, step in zip(image_shape, offset, strict=True): + if step > 0: + if step >= axis_size: + return None + slices.append(slice(step, axis_size)) + elif step < 0: + if -step >= axis_size: + return None + slices.append(slice(0, axis_size + step)) + else: + slices.append(slice(None)) + return tuple(slices) + + +def _compute_valid_edges( + affinity_shape: tuple[int, ...], + offsets: Sequence[tuple[int, ...]], + number_of_attractive_channels: int, + strides: tuple[int, ...] | None, + randomized_strides: bool, + mask: np.ndarray | None, +) -> np.ndarray: + image_shape = tuple(int(size) for size in affinity_shape[1:]) + valid_edges = np.zeros(affinity_shape, dtype=bool) + + for channel, offset in enumerate(offsets): + source_slices = _valid_source_slices(image_shape, offset) + if source_slices is not None: + valid_edges[(channel,) + source_slices] = True + + if strides is not None: + stride_edges = np.zeros_like(valid_edges, dtype=bool) + stride_edges[:number_of_attractive_channels] = True + if randomized_strides: + stride_factor = 1.0 / np.prod(strides) + stride_edges[number_of_attractive_channels:] = ( + np.random.random( + valid_edges[number_of_attractive_channels:].shape + ) + < stride_factor + ) + else: + valid_slice = (slice(number_of_attractive_channels, None),) + tuple( + slice(None, None, stride) for stride in strides + ) + stride_edges[valid_slice] = True + valid_edges &= stride_edges + + if mask is not None: + mask_array = np.asarray(mask) + if mask_array.shape != image_shape: + raise ValueError( + "mask shape must match affinities spatial shape, got " + f"mask shape={mask_array.shape}, spatial shape={image_shape}" + ) + if mask_array.dtype != np.dtype("bool"): + raise TypeError(f"mask must have dtype bool, got dtype={mask_array.dtype}") + + invalid = np.logical_not(mask_array) + for channel, offset in enumerate(offsets): + source_slices = _valid_source_slices(image_shape, offset) + neighbor_slices = _neighbor_slices(image_shape, offset) + if source_slices is None or neighbor_slices is None: + continue + touches_invalid = invalid[source_slices] | invalid[neighbor_slices] + channel_valid = valid_edges[(channel,) + source_slices] + channel_valid[touches_invalid] = False + valid_edges[(channel,) + source_slices] = channel_valid + + return np.ascontiguousarray(valid_edges, dtype=np.uint8) diff --git a/src/cpp/segmentation/mutex_watershed.cxx b/src/cpp/segmentation/mutex_watershed.cxx index 135a556..cfac854 100644 --- a/src/cpp/segmentation/mutex_watershed.cxx +++ b/src/cpp/segmentation/mutex_watershed.cxx @@ -6,6 +6,7 @@ namespace bioimage_cpp { template void mutex_watershed_grid( const ConstArrayView &, + const ConstArrayView &, const std::vector> &, std::size_t, const ArrayView & @@ -13,6 +14,7 @@ template void mutex_watershed_grid( template void mutex_watershed_grid( const ConstArrayView &, + const ConstArrayView &, const std::vector> &, std::size_t, const ArrayView & diff --git a/tests/segmentation/test_mutex_watershed.py b/tests/segmentation/test_mutex_watershed.py index c29bb0f..92640c8 100644 --- a/tests/segmentation/test_mutex_watershed.py +++ b/tests/segmentation/test_mutex_watershed.py @@ -77,6 +77,61 @@ def test_mutex_watershed_3d_attractive_edges_merge_all_pixels(): np.testing.assert_array_equal(labels, np.ones((2, 2, 2), dtype=np.uint64)) +def test_mutex_watershed_strides_subsample_mutex_edges(): + affinities = np.zeros((2, 1, 5), dtype=np.float32) + affinities[0, 0, :] = 0.8 + affinities[1, 0, :] = 0.9 + offsets = [[0, 1], [0, 1]] + + dense = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1 + ) + strided = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1, strides=[1, 2] + ) + + np.testing.assert_array_equal(dense, np.array([[1, 2, 3, 4, 5]], dtype=np.uint64)) + np.testing.assert_array_equal(strided, np.array([[1, 2, 2, 3, 3]], dtype=np.uint64)) + + +def test_mutex_watershed_randomized_strides_use_numpy_random_state(): + affinities = np.zeros((2, 1, 5), dtype=np.float32) + affinities[0, 0, :] = 0.8 + affinities[1, 0, :] = 0.9 + offsets = [[0, 1], [0, 1]] + + np.random.seed(17) + first = bic.segmentation.mutex_watershed( + affinities, + offsets, + number_of_attractive_channels=1, + strides=[1, 2], + randomized_strides=True, + ) + np.random.seed(17) + second = bic.segmentation.mutex_watershed( + affinities, + offsets, + number_of_attractive_channels=1, + strides=[1, 2], + randomized_strides=True, + ) + + np.testing.assert_array_equal(first, second) + + +def test_mutex_watershed_mask_sets_background_zero_and_blocks_edges(): + affinities = np.ones((1, 1, 5), dtype=np.float32) + offsets = [[0, 1]] + mask = np.array([[True, True, False, True, True]], dtype=bool) + + labels = bic.segmentation.mutex_watershed( + affinities, offsets, number_of_attractive_channels=1, mask=mask + ) + + np.testing.assert_array_equal(labels, np.array([[1, 1, 0, 3, 3]], dtype=np.uint64)) + + def test_mutex_watershed_rejects_wrong_affinity_dtype(): with pytest.raises(TypeError, match="affinities must have one of dtypes"): bic.segmentation.mutex_watershed( @@ -102,3 +157,23 @@ def test_mutex_watershed_rejects_wrong_number_of_offsets(): [[0, 1]], number_of_attractive_channels=1, ) + + +def test_mutex_watershed_rejects_invalid_strides(): + with pytest.raises(ValueError, match="strides length must match"): + bic.segmentation.mutex_watershed( + np.ones((1, 2, 2), dtype=np.float32), + [[0, 1]], + number_of_attractive_channels=1, + strides=[1, 1, 1], + ) + + +def test_mutex_watershed_rejects_invalid_mask(): + with pytest.raises(TypeError, match="mask must have dtype bool"): + bic.segmentation.mutex_watershed( + np.ones((1, 2, 2), dtype=np.float32), + [[0, 1]], + number_of_attractive_channels=1, + mask=np.ones((2, 2), dtype=np.uint8), + )