Skip to content

Commit dbcfe1d

Browse files
constantinpapeconstantinpape
committed
Implement semantic mutex watershed
1 parent 7d5c937 commit dbcfe1d

14 files changed

Lines changed: 1549 additions & 144 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ nanobind_add_module(_core
1919
src/bindings/segmentation.cxx
2020
src/bindings/utils.cxx
2121
src/cpp/segmentation/mutex_watershed.cxx
22+
src/cpp/segmentation/semantic_mutex_watershed.cxx
2223
src/cpp/take_dict.cxx
2324
)
2425

‎include/bioimage_cpp/detail/semantic_labels.hxx‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
#pragma once
2+
3+
#include <cstdint>
4+
#include <vector>
5+
6+
namespace bioimage_cpp {
7+
8+
using SemanticLabeling = std::vector<std::int64_t>;
9+
10+
// Returns true if the two roots already carry different non-negative semantic
11+
// labels. Merging two such roots would conflict the semantic assignments, so
12+
// the caller must skip that merge.
13+
inline bool check_semantic_constraint(
14+
const std::uint64_t first,
15+
const std::uint64_t second,
16+
const SemanticLabeling &semantic_labels
17+
) {
18+
const auto label_first = semantic_labels[first];
19+
const auto label_second = semantic_labels[second];
20+
if (label_first >= 0 && label_second >= 0) {
21+
return label_first != label_second;
22+
}
23+
return false;
24+
}
25+
26+
// Assigns `class_id` to `root` only if `root` has no semantic label yet
27+
// (-1 == unassigned). Already-assigned roots keep their first assignment;
28+
// because semantic edges are processed in descending-weight order this means
29+
// each root keeps the strongest class assignment seen for it.
30+
inline void assign_semantic_label(
31+
const std::uint64_t root,
32+
const std::int64_t class_id,
33+
SemanticLabeling &semantic_labels
34+
) {
35+
if (semantic_labels[root] < 0) {
36+
semantic_labels[root] = class_id;
37+
}
38+
}
39+
40+
// After two roots merge, propagate a non-negative semantic label from one to
41+
// the other if exactly one was assigned. If both are assigned the caller is
42+
// expected to have rejected the merge via `check_semantic_constraint`.
43+
inline void merge_semantic_labels(
44+
const std::uint64_t first,
45+
const std::uint64_t second,
46+
SemanticLabeling &semantic_labels
47+
) {
48+
const auto label_first = semantic_labels[first];
49+
const auto label_second = semantic_labels[second];
50+
if (label_first >= 0 && label_second < 0) {
51+
semantic_labels[second] = label_first;
52+
} else if (label_first < 0 && label_second >= 0) {
53+
semantic_labels[first] = label_second;
54+
}
55+
}
56+
57+
} // namespace bioimage_cpp

‎include/bioimage_cpp/graph/mutex_watershed.hxx‎

Lines changed: 2 additions & 142 deletions
Original file line numberDiff line numberDiff line change
@@ -1,144 +1,4 @@
11
#pragma once
22

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-
//
30-
// Templated on the weight type. Concrete instantiations for `float` and
31-
// `double` are provided by the binding layer; other floating types are
32-
// supported but must be instantiated explicitly.
33-
template <class WeightT>
34-
std::vector<std::uint64_t> mutex_watershed_clustering(
35-
const UndirectedGraph &graph,
36-
const std::vector<WeightT> &edge_costs,
37-
const std::vector<std::array<std::uint64_t, 2>> &mutex_uvs,
38-
const std::vector<WeightT> &mutex_costs
39-
) {
40-
const auto number_of_edges = static_cast<std::size_t>(graph.number_of_edges());
41-
if (edge_costs.size() != number_of_edges) {
42-
throw std::invalid_argument(
43-
"edge_costs size must match graph.number_of_edges(), got edge_costs size=" +
44-
std::to_string(edge_costs.size()) +
45-
", number_of_edges=" + std::to_string(number_of_edges)
46-
);
47-
}
48-
if (mutex_costs.size() != mutex_uvs.size()) {
49-
throw std::invalid_argument(
50-
"mutex_costs size must match mutex_uvs size, got mutex_costs size=" +
51-
std::to_string(mutex_costs.size()) +
52-
", mutex_uvs size=" + std::to_string(mutex_uvs.size())
53-
);
54-
}
55-
56-
const auto number_of_nodes = static_cast<std::size_t>(graph.number_of_nodes());
57-
const auto number_of_mutex = mutex_uvs.size();
58-
59-
for (std::size_t index = 0; index < number_of_mutex; ++index) {
60-
const auto u = mutex_uvs[index][0];
61-
const auto v = mutex_uvs[index][1];
62-
if (u >= number_of_nodes || v >= number_of_nodes) {
63-
throw std::invalid_argument(
64-
"mutex_uvs endpoints must be < number_of_nodes, got u=" +
65-
std::to_string(u) + ", v=" + std::to_string(v) +
66-
", number_of_nodes=" + std::to_string(number_of_nodes)
67-
);
68-
}
69-
}
70-
71-
struct WeightedEdge {
72-
WeightT weight;
73-
std::uint64_t index;
74-
bool is_mutex;
75-
};
76-
77-
std::vector<WeightedEdge> edge_order;
78-
edge_order.reserve(number_of_edges + number_of_mutex);
79-
for (std::size_t index = 0; index < number_of_edges; ++index) {
80-
edge_order.push_back(
81-
WeightedEdge{edge_costs[index], static_cast<std::uint64_t>(index), false}
82-
);
83-
}
84-
for (std::size_t index = 0; index < number_of_mutex; ++index) {
85-
edge_order.push_back(
86-
WeightedEdge{mutex_costs[index], static_cast<std::uint64_t>(index), true}
87-
);
88-
}
89-
90-
std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) {
91-
if (first.weight != second.weight) {
92-
return first.weight > second.weight;
93-
}
94-
if (first.is_mutex != second.is_mutex) {
95-
return !first.is_mutex;
96-
}
97-
return first.index < second.index;
98-
});
99-
100-
bioimage_cpp::detail::UnionFind sets(number_of_nodes);
101-
MutexStorage mutexes(number_of_nodes);
102-
103-
for (const auto &edge : edge_order) {
104-
std::uint64_t u;
105-
std::uint64_t v;
106-
if (edge.is_mutex) {
107-
const auto &pair = mutex_uvs[static_cast<std::size_t>(edge.index)];
108-
u = pair[0];
109-
v = pair[1];
110-
} else {
111-
const auto uv = graph.uv(edge.index);
112-
u = uv.first;
113-
v = uv.second;
114-
}
115-
if (u == v) {
116-
continue;
117-
}
118-
119-
const auto root_u = sets.find(u);
120-
const auto root_v = sets.find(v);
121-
if (root_u == root_v) {
122-
continue;
123-
}
124-
125-
if (edge.is_mutex) {
126-
insert_mutex(root_u, root_v, mutexes);
127-
} else {
128-
if (check_mutex(root_u, root_v, mutexes)) {
129-
continue;
130-
}
131-
const auto new_root = sets.unite_roots(root_u, root_v);
132-
const auto old_root = (new_root == root_u) ? root_v : root_u;
133-
merge_mutexes(old_root, new_root, mutexes);
134-
}
135-
}
136-
137-
std::vector<std::uint64_t> roots(number_of_nodes);
138-
for (std::size_t node = 0; node < number_of_nodes; ++node) {
139-
roots[node] = sets.find(static_cast<std::uint64_t>(node));
140-
}
141-
return bioimage_cpp::detail::dense_relabel(roots);
142-
}
143-
144-
} // namespace bioimage_cpp::graph
3+
#include "bioimage_cpp/graph/mutex_watershed/clustering.hxx"
4+
#include "bioimage_cpp/graph/mutex_watershed/semantic.hxx"

‎include/bioimage_cpp/graph/mutex_watershed/clustering.hxx‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
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+
//
30+
// Templated on the weight type. Concrete instantiations for `float` and
31+
// `double` are provided by the binding layer; other floating types are
32+
// supported but must be instantiated explicitly.
33+
template <class WeightT>
34+
std::vector<std::uint64_t> mutex_watershed_clustering(
35+
const UndirectedGraph &graph,
36+
const std::vector<WeightT> &edge_costs,
37+
const std::vector<std::array<std::uint64_t, 2>> &mutex_uvs,
38+
const std::vector<WeightT> &mutex_costs
39+
) {
40+
const auto number_of_edges = static_cast<std::size_t>(graph.number_of_edges());
41+
if (edge_costs.size() != number_of_edges) {
42+
throw std::invalid_argument(
43+
"edge_costs size must match graph.number_of_edges(), got edge_costs size=" +
44+
std::to_string(edge_costs.size()) +
45+
", number_of_edges=" + std::to_string(number_of_edges)
46+
);
47+
}
48+
if (mutex_costs.size() != mutex_uvs.size()) {
49+
throw std::invalid_argument(
50+
"mutex_costs size must match mutex_uvs size, got mutex_costs size=" +
51+
std::to_string(mutex_costs.size()) +
52+
", mutex_uvs size=" + std::to_string(mutex_uvs.size())
53+
);
54+
}
55+
56+
const auto number_of_nodes = static_cast<std::size_t>(graph.number_of_nodes());
57+
const auto number_of_mutex = mutex_uvs.size();
58+
59+
for (std::size_t index = 0; index < number_of_mutex; ++index) {
60+
const auto u = mutex_uvs[index][0];
61+
const auto v = mutex_uvs[index][1];
62+
if (u >= number_of_nodes || v >= number_of_nodes) {
63+
throw std::invalid_argument(
64+
"mutex_uvs endpoints must be < number_of_nodes, got u=" +
65+
std::to_string(u) + ", v=" + std::to_string(v) +
66+
", number_of_nodes=" + std::to_string(number_of_nodes)
67+
);
68+
}
69+
}
70+
71+
struct WeightedEdge {
72+
WeightT weight;
73+
std::uint64_t index;
74+
bool is_mutex;
75+
};
76+
77+
std::vector<WeightedEdge> edge_order;
78+
edge_order.reserve(number_of_edges + number_of_mutex);
79+
for (std::size_t index = 0; index < number_of_edges; ++index) {
80+
edge_order.push_back(
81+
WeightedEdge{edge_costs[index], static_cast<std::uint64_t>(index), false}
82+
);
83+
}
84+
for (std::size_t index = 0; index < number_of_mutex; ++index) {
85+
edge_order.push_back(
86+
WeightedEdge{mutex_costs[index], static_cast<std::uint64_t>(index), true}
87+
);
88+
}
89+
90+
std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) {
91+
if (first.weight != second.weight) {
92+
return first.weight > second.weight;
93+
}
94+
if (first.is_mutex != second.is_mutex) {
95+
return !first.is_mutex;
96+
}
97+
return first.index < second.index;
98+
});
99+
100+
bioimage_cpp::detail::UnionFind sets(number_of_nodes);
101+
MutexStorage mutexes(number_of_nodes);
102+
103+
for (const auto &edge : edge_order) {
104+
std::uint64_t u;
105+
std::uint64_t v;
106+
if (edge.is_mutex) {
107+
const auto &pair = mutex_uvs[static_cast<std::size_t>(edge.index)];
108+
u = pair[0];
109+
v = pair[1];
110+
} else {
111+
const auto uv = graph.uv(edge.index);
112+
u = uv.first;
113+
v = uv.second;
114+
}
115+
if (u == v) {
116+
continue;
117+
}
118+
119+
const auto root_u = sets.find(u);
120+
const auto root_v = sets.find(v);
121+
if (root_u == root_v) {
122+
continue;
123+
}
124+
125+
if (edge.is_mutex) {
126+
insert_mutex(root_u, root_v, mutexes);
127+
} else {
128+
if (check_mutex(root_u, root_v, mutexes)) {
129+
continue;
130+
}
131+
const auto new_root = sets.unite_roots(root_u, root_v);
132+
const auto old_root = (new_root == root_u) ? root_v : root_u;
133+
merge_mutexes(old_root, new_root, mutexes);
134+
}
135+
}
136+
137+
std::vector<std::uint64_t> roots(number_of_nodes);
138+
for (std::size_t node = 0; node < number_of_nodes; ++node) {
139+
roots[node] = sets.find(static_cast<std::uint64_t>(node));
140+
}
141+
return bioimage_cpp::detail::dense_relabel(roots);
142+
}
143+
144+
} // namespace bioimage_cpp::graph

0 commit comments

Comments
 (0)