Skip to content

Commit 2af75e8

Browse files
constantinpapeconstantinpape
committed
Implement mutex watershed clustering
1 parent 4a6b523 commit 2af75e8

6 files changed

Lines changed: 491 additions & 46 deletions

File tree

‎include/bioimage_cpp/detail/mutex_storage.hxx‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#pragma once
2+
3+
#include <cstdint>
4+
#include <unordered_set>
5+
#include <vector>
6+
7+
namespace bioimage_cpp {
8+
9+
using MutexStorage = std::vector<std::unordered_set<std::uint64_t>>;
10+
11+
inline bool check_mutex(
12+
const std::uint64_t first,
13+
const std::uint64_t second,
14+
const MutexStorage &mutexes
15+
) {
16+
const auto &first_mutexes = mutexes[first];
17+
const auto &second_mutexes = mutexes[second];
18+
if (first_mutexes.size() < second_mutexes.size()) {
19+
return first_mutexes.find(second) != first_mutexes.end();
20+
}
21+
return second_mutexes.find(first) != second_mutexes.end();
22+
}
23+
24+
inline void insert_mutex(
25+
const std::uint64_t first,
26+
const std::uint64_t second,
27+
MutexStorage &mutexes
28+
) {
29+
mutexes[first].insert(second);
30+
mutexes[second].insert(first);
31+
}
32+
33+
inline void merge_mutexes(
34+
const std::uint64_t root_from,
35+
const std::uint64_t root_to,
36+
MutexStorage &mutexes
37+
) {
38+
auto &mutexes_from = mutexes[root_from];
39+
auto &mutexes_to = mutexes[root_to];
40+
41+
for (const auto other_root : mutexes_from) {
42+
auto &other_mutexes = mutexes[other_root];
43+
other_mutexes.erase(root_from);
44+
if (other_root != root_to) {
45+
other_mutexes.insert(root_to);
46+
mutexes_to.insert(other_root);
47+
}
48+
}
49+
mutexes_to.erase(root_from);
50+
mutexes_to.erase(root_to);
51+
mutexes_from.clear();
52+
}
53+
54+
} // namespace bioimage_cpp

‎include/bioimage_cpp/graph/mutex_watershed.hxx‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
#pragma once
2+
3+
#include "bioimage_cpp/detail/mutex_storage.hxx"
4+
#include "bioimage_cpp/detail/relabel.hxx"
5+
#include "bioimage_cpp/detail/union_find.hxx"
6+
#include "bioimage_cpp/graph/undirected_graph.hxx"
7+
8+
#include <algorithm>
9+
#include <array>
10+
#include <cstddef>
11+
#include <cstdint>
12+
#include <stdexcept>
13+
#include <string>
14+
#include <vector>
15+
16+
namespace bioimage_cpp::graph {
17+
18+
// Mutex watershed clustering on an arbitrary undirected graph.
19+
//
20+
// `graph` defines the attractive edges with one cost per edge in `edge_costs`.
21+
// `mutex_uvs` defines the long-range repulsive (mutex) edges as (u, v) pairs
22+
// with one cost per edge in `mutex_costs`. Higher costs win — they are
23+
// processed first when sorting jointly by descending weight.
24+
//
25+
// Returns dense node labels in [0, k) following first-occurrence order.
26+
//
27+
// This is a port of `compute_mws_clustering` from the affogato library,
28+
// adapted to bioimage-cpp's UndirectedGraph and detail/ primitives.
29+
inline std::vector<std::uint64_t> mutex_watershed_clustering(
30+
const UndirectedGraph &graph,
31+
const std::vector<double> &edge_costs,
32+
const std::vector<std::array<std::uint64_t, 2>> &mutex_uvs,
33+
const std::vector<double> &mutex_costs
34+
) {
35+
const auto number_of_edges = static_cast<std::size_t>(graph.number_of_edges());
36+
if (edge_costs.size() != number_of_edges) {
37+
throw std::invalid_argument(
38+
"edge_costs size must match graph.number_of_edges(), got edge_costs size=" +
39+
std::to_string(edge_costs.size()) +
40+
", number_of_edges=" + std::to_string(number_of_edges)
41+
);
42+
}
43+
if (mutex_costs.size() != mutex_uvs.size()) {
44+
throw std::invalid_argument(
45+
"mutex_costs size must match mutex_uvs size, got mutex_costs size=" +
46+
std::to_string(mutex_costs.size()) +
47+
", mutex_uvs size=" + std::to_string(mutex_uvs.size())
48+
);
49+
}
50+
51+
const auto number_of_nodes = static_cast<std::size_t>(graph.number_of_nodes());
52+
const auto number_of_mutex = mutex_uvs.size();
53+
54+
for (std::size_t index = 0; index < number_of_mutex; ++index) {
55+
const auto u = mutex_uvs[index][0];
56+
const auto v = mutex_uvs[index][1];
57+
if (u >= number_of_nodes || v >= number_of_nodes) {
58+
throw std::invalid_argument(
59+
"mutex_uvs endpoints must be < number_of_nodes, got u=" +
60+
std::to_string(u) + ", v=" + std::to_string(v) +
61+
", number_of_nodes=" + std::to_string(number_of_nodes)
62+
);
63+
}
64+
}
65+
66+
struct WeightedEdge {
67+
double weight;
68+
std::uint64_t index;
69+
bool is_mutex;
70+
};
71+
72+
std::vector<WeightedEdge> edge_order;
73+
edge_order.reserve(number_of_edges + number_of_mutex);
74+
for (std::size_t index = 0; index < number_of_edges; ++index) {
75+
edge_order.push_back(
76+
WeightedEdge{edge_costs[index], static_cast<std::uint64_t>(index), false}
77+
);
78+
}
79+
for (std::size_t index = 0; index < number_of_mutex; ++index) {
80+
edge_order.push_back(
81+
WeightedEdge{mutex_costs[index], static_cast<std::uint64_t>(index), true}
82+
);
83+
}
84+
85+
std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) {
86+
if (first.weight != second.weight) {
87+
return first.weight > second.weight;
88+
}
89+
if (first.is_mutex != second.is_mutex) {
90+
return !first.is_mutex;
91+
}
92+
return first.index < second.index;
93+
});
94+
95+
bioimage_cpp::detail::UnionFind sets(number_of_nodes);
96+
MutexStorage mutexes(number_of_nodes);
97+
98+
for (const auto &edge : edge_order) {
99+
std::uint64_t u;
100+
std::uint64_t v;
101+
if (edge.is_mutex) {
102+
const auto &pair = mutex_uvs[static_cast<std::size_t>(edge.index)];
103+
u = pair[0];
104+
v = pair[1];
105+
} else {
106+
const auto uv = graph.uv(edge.index);
107+
u = uv.first;
108+
v = uv.second;
109+
}
110+
if (u == v) {
111+
continue;
112+
}
113+
114+
const auto root_u = sets.find(u);
115+
const auto root_v = sets.find(v);
116+
if (root_u == root_v) {
117+
continue;
118+
}
119+
120+
if (edge.is_mutex) {
121+
insert_mutex(root_u, root_v, mutexes);
122+
} else {
123+
if (check_mutex(root_u, root_v, mutexes)) {
124+
continue;
125+
}
126+
const auto new_root = sets.unite_roots(root_u, root_v);
127+
const auto old_root = (new_root == root_u) ? root_v : root_u;
128+
merge_mutexes(old_root, new_root, mutexes);
129+
}
130+
}
131+
132+
std::vector<std::uint64_t> roots(number_of_nodes);
133+
for (std::size_t node = 0; node < number_of_nodes; ++node) {
134+
roots[node] = sets.find(static_cast<std::uint64_t>(node));
135+
}
136+
return bioimage_cpp::detail::dense_relabel(roots);
137+
}
138+
139+
} // namespace bioimage_cpp::graph

‎include/bioimage_cpp/segmentation/mutex_watershed.hxx‎

Lines changed: 1 addition & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
#include "bioimage_cpp/array_view.hxx"
44
#include "bioimage_cpp/detail/grid.hxx"
5+
#include "bioimage_cpp/detail/mutex_storage.hxx"
56
#include "bioimage_cpp/detail/union_find.hxx"
67

78
#include <algorithm>
@@ -10,62 +11,16 @@
1011
#include <numeric>
1112
#include <stdexcept>
1213
#include <string>
13-
#include <unordered_set>
1414
#include <vector>
1515

1616
namespace bioimage_cpp {
1717

18-
using MutexStorage = std::vector<std::unordered_set<std::uint64_t>>;
19-
2018
template <class T>
2119
struct WeightedGridEdge {
2220
T weight;
2321
std::uint64_t id;
2422
};
2523

26-
inline bool check_mutex(
27-
const std::uint64_t first,
28-
const std::uint64_t second,
29-
const MutexStorage &mutexes
30-
) {
31-
const auto &first_mutexes = mutexes[first];
32-
const auto &second_mutexes = mutexes[second];
33-
if (first_mutexes.size() < second_mutexes.size()) {
34-
return first_mutexes.find(second) != first_mutexes.end();
35-
}
36-
return second_mutexes.find(first) != second_mutexes.end();
37-
}
38-
39-
inline void insert_mutex(
40-
const std::uint64_t first,
41-
const std::uint64_t second,
42-
MutexStorage &mutexes
43-
) {
44-
mutexes[first].insert(second);
45-
mutexes[second].insert(first);
46-
}
47-
48-
inline void merge_mutexes(
49-
const std::uint64_t root_from,
50-
const std::uint64_t root_to,
51-
MutexStorage &mutexes
52-
) {
53-
auto &mutexes_from = mutexes[root_from];
54-
auto &mutexes_to = mutexes[root_to];
55-
56-
for (const auto other_root : mutexes_from) {
57-
auto &other_mutexes = mutexes[other_root];
58-
other_mutexes.erase(root_from);
59-
if (other_root != root_to) {
60-
other_mutexes.insert(root_to);
61-
mutexes_to.insert(other_root);
62-
}
63-
}
64-
mutexes_to.erase(root_from);
65-
mutexes_to.erase(root_to);
66-
mutexes_from.clear();
67-
}
68-
6924
template <class T>
7025
void mutex_watershed_grid(
7126
const ConstArrayView<T> &affinities,

‎src/bindings/graph.cxx‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include "bioimage_cpp/graph/lifted_from_affinities.hxx"
1111
#include "bioimage_cpp/graph/lifted_multicut.hxx"
1212
#include "bioimage_cpp/graph/multicut.hxx"
13+
#include "bioimage_cpp/graph/mutex_watershed.hxx"
1314
#include "bioimage_cpp/graph/multicut/fusion_move.hxx"
1415
#include "bioimage_cpp/graph/multicut/greedy_additive.hxx"
1516
#include "bioimage_cpp/graph/multicut/greedy_fixation.hxx"
@@ -790,6 +791,36 @@ UInt64Array lifted_multicut_kernighan_lin(
790791
return vector_to_uint64_array(label_vector);
791792
}
792793

794+
UInt64Array mutex_watershed_clustering(
795+
const Graph &graph,
796+
ConstDoubleArray edge_costs,
797+
ConstUInt64Array mutex_uvs,
798+
ConstDoubleArray mutex_costs
799+
) {
800+
const auto edge_cost_vector =
801+
double_array_to_vector(edge_costs, "edge_costs", graph.number_of_edges());
802+
require_uv_array(mutex_uvs, "mutex_uvs");
803+
const auto n_mutex = mutex_uvs.shape(0);
804+
const auto mutex_cost_vector =
805+
double_array_to_vector(mutex_costs, "mutex_costs", static_cast<std::uint64_t>(n_mutex));
806+
807+
std::vector<std::array<std::uint64_t, 2>> mutex_uv_vector(n_mutex);
808+
const auto *uv_data = mutex_uvs.data();
809+
for (std::size_t index = 0; index < n_mutex; ++index) {
810+
mutex_uv_vector[index][0] = uv_data[2 * index];
811+
mutex_uv_vector[index][1] = uv_data[2 * index + 1];
812+
}
813+
814+
std::vector<std::uint64_t> labels;
815+
{
816+
nb::gil_scoped_release release;
817+
labels = graph::mutex_watershed_clustering(
818+
graph, edge_cost_vector, mutex_uv_vector, mutex_cost_vector
819+
);
820+
}
821+
return vector_to_uint64_array(labels);
822+
}
823+
793824
UInt64Array multicut_fusion_move(
794825
const Graph &graph,
795826
ConstDoubleArray costs,
@@ -1429,6 +1460,15 @@ void bind_graph(nb::module_ &m) {
14291460
nb::arg("epsilon")
14301461
);
14311462

1463+
m.def(
1464+
"_mutex_watershed_clustering",
1465+
&mutex_watershed_clustering,
1466+
nb::arg("graph"),
1467+
nb::arg("edge_costs"),
1468+
nb::arg("mutex_uvs"),
1469+
nb::arg("mutex_costs")
1470+
);
1471+
14321472
// Lifted multicut sub-solver hierarchy. Same shape as the multicut sub-
14331473
// solver bindings — opaque to Python, used by future fusion-move drivers.
14341474
nb::class_<graph::lifted_multicut::SolverBase>(m, "_LiftedMulticutSolverBase");

‎src/bioimage_cpp/graph/__init__.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,57 @@ def edge_weighted_watershed(
599599
return run(graph, weight_array, seed_array)
600600

601601

602+
def mutex_watershed_clustering(
603+
graph: UndirectedGraph | RegionAdjacencyGraph,
604+
edge_costs,
605+
mutex_uvs,
606+
mutex_costs,
607+
) -> np.ndarray:
608+
"""Mutex watershed clustering on an undirected graph.
609+
610+
Attractive edges come from ``graph`` (one cost per edge in
611+
``edge_costs``); repulsive long-range edges are supplied separately as
612+
``mutex_uvs`` with weights ``mutex_costs``. All edges are jointly sorted
613+
by descending weight and processed in a single pass: an attractive edge
614+
merges its two components unless a mutex constraint already separates
615+
them; a mutex edge installs a constraint between the two current
616+
components.
617+
618+
The input format matches :class:`LiftedMulticutObjective` — the same
619+
``(graph, edge_costs, lifted_uvs, lifted_costs)`` arrays can be passed
620+
here as ``(graph, edge_costs, mutex_uvs, mutex_costs)`` — though the
621+
algorithms differ (mutex constraints are hard; lifted costs are soft).
622+
623+
Parameters
624+
----------
625+
graph:
626+
:class:`UndirectedGraph` or :class:`RegionAdjacencyGraph` defining
627+
the attractive edges.
628+
edge_costs:
629+
1D float64 array of length ``graph.number_of_edges``. Higher
630+
values are more attractive.
631+
mutex_uvs:
632+
``(n_mutex, 2)`` uint64 array of (u, v) pairs for the mutex edges.
633+
mutex_costs:
634+
1D float64 array of length ``n_mutex``. Higher values are stronger
635+
repulsions.
636+
637+
Returns
638+
-------
639+
np.ndarray
640+
``(graph.number_of_nodes,)`` uint64 array. Dense node labels in
641+
``[0, k)`` in first-occurrence order.
642+
"""
643+
edge_cost_array = _as_edge_costs(edge_costs, graph)
644+
mutex_uv_array = _as_uv_array(mutex_uvs, "mutex_uvs")
645+
mutex_cost_array = _as_1d_array(
646+
mutex_costs, np.float64, "mutex_costs", int(mutex_uv_array.shape[0])
647+
)
648+
return _core._mutex_watershed_clustering(
649+
graph, edge_cost_array, mutex_uv_array, mutex_cost_array
650+
)
651+
652+
602653
class MulticutObjective:
603654
"""Multicut objective for an undirected graph and edge costs."""
604655

@@ -1818,6 +1869,7 @@ def _normalize_number_of_threads(number_of_threads: int) -> int:
18181869
"load_multicut_problem",
18191870
"load_multicut_problem_data",
18201871
"multicut_problem_path",
1872+
"mutex_watershed_clustering",
18211873
"project_node_labels_to_pixels",
18221874
"region_adjacency_graph",
18231875
"undirected_graph",

0 commit comments

Comments
 (0)