MWS implementation WIP

Author: constantinpape

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
 )
 

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`.

include/bioimage_cpp/segmentation/mutex_watershed.hxx

@@ -0,0 +1,228 @@
+#pragma once
+
+#include "bioimage_cpp/array_view.hxx"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+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<std::uint64_t> parents_;
+    std::vector<std::uint64_t> ranks_;
+};
+
+using MutexStorage = std::vector<std::unordered_set<std::uint64_t>>;
+
+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<std::ptrdiff_t> c_order_strides(const std::vector<std::ptrdiff_t> &shape) {
+    std::vector<std::ptrdiff_t> strides(shape.size(), 1);
+    for (std::ptrdiff_t axis = static_cast<std::ptrdiff_t>(shape.size()) - 2; axis >= 0; --axis) {
+        strides[static_cast<std::size_t>(axis)] =
+            strides[static_cast<std::size_t>(axis + 1)] * shape[static_cast<std::size_t>(axis + 1)];
+    }
+    return strides;
+}
+
+inline bool is_valid_grid_edge(
+    const std::uint64_t node,
+    const std::vector<std::ptrdiff_t> &offset,
+    const std::vector<std::ptrdiff_t> &shape,
+    const std::vector<std::ptrdiff_t> &strides
+) {
+    for (std::size_t axis = 0; axis < shape.size(); ++axis) {
+        const auto coord = static_cast<std::ptrdiff_t>(node / static_cast<std::uint64_t>(strides[axis])) % shape[axis];
+        const auto neighbor = coord + offset[axis];
+        if (neighbor < 0 || neighbor >= shape[axis]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+template <class T>
+void mutex_watershed_grid(
+    const ConstArrayView<T> &affinities,
+    const std::vector<std::vector<std::ptrdiff_t>> &offsets,
+    const std::size_t number_of_attractive_channels,
+    const ArrayView<std::uint64_t> &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<std::size_t>(affinities.shape[0]);
+    const auto spatial_ndim = static_cast<std::size_t>(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<std::ptrdiff_t> 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::uint64_t>(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<std::ptrdiff_t> 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<std::uint64_t> edge_order(static_cast<std::size_t>(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<std::size_t>(number_of_nodes));
+    MutexStorage mutexes(static_cast<std::size_t>(number_of_nodes));
+
+    for (const auto edge_id : edge_order) {
+        const auto channel = static_cast<std::size_t>(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<std::int64_t>(u) + static_cast<std::int64_t>(offset_strides[channel]);
+        const auto v = static_cast<std::uint64_t>(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<std::uint64_t, std::uint64_t> 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<std::uint64_t>(label_map.size() + 1);
+            found = label_map.emplace(root, next_label).first;
+        }
+        out.data[node] = found->second;
+    }
+}
+
+} // namespace bioimage_cpp

src/bindings/module.cxx

@@ -1,3 +1,4 @@
+#include "segmentation.hxx"
 #include "utils.hxx"
 
 #include <nanobind/nanobind.h>
@@ -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);
 }

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 <nanobind/ndarray.h>
+#include <nanobind/stl/vector.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <numeric>
+#include <vector>
+
+namespace nb = nanobind;
+
+namespace bioimage_cpp::bindings {
+namespace {
+
+template <class T>
+using AffinityArray = nb::ndarray<nb::numpy, const T, nb::c_contig>;
+
+using LabelArray = nb::ndarray<nb::numpy, std::uint64_t, nb::c_contig>;
+
+template <class T>
+LabelArray mutex_watershed_grid_t(
+    AffinityArray<T> affinities,
+    const std::vector<std::vector<std::ptrdiff_t>> &offsets,
+    const std::size_t number_of_attractive_channels
+) {
+    std::vector<std::ptrdiff_t> affinity_shape(affinities.ndim());
+    for (std::size_t axis = 0; axis < affinities.ndim(); ++axis) {
+        affinity_shape[axis] = static_cast<std::ptrdiff_t>(affinities.shape(axis));
+    }
+
+    std::vector<std::size_t> label_shape;
+    label_shape.reserve(affinities.ndim() > 0 ? affinities.ndim() - 1 : 0);
+    std::vector<std::ptrdiff_t> 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<std::ptrdiff_t>(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<std::size_t>(number_of_nodes)]();
+    nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast<std::uint64_t *>(p); });
+
+    ConstArrayView<T> affinities_view{
+        affinities.data(),
+        affinity_shape,
+        {},
+    };
+    ArrayView<std::uint64_t> out_view{
+        data,
+        label_view_shape,
+        {},
+    };
+
+    {
+        nb::gil_scoped_release release;
+        mutex_watershed_grid<T>(
+            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<float>,
+        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<double>,
+        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

src/bindings/segmentation.hxx

@@ -0,0 +1,9 @@
+#pragma once
+
+#include <nanobind/nanobind.h>
+
+namespace bioimage_cpp::bindings {
+
+void bind_segmentation(nanobind::module_ &m);
+
+} // namespace bioimage_cpp::bindings

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"]

src/bioimage_cpp/segmentation/__init__.py

@@ -0,0 +1,5 @@
+"""Segmentation algorithms."""
+
+from .mutex_watershed import mutex_watershed
+
+__all__ = ["mutex_watershed"]