Skip to content

Commit 28eae5e

Browse files
constantinpapeconstantinpape
committed
Add RAG
1 parent dc7734b commit 28eae5e

4 files changed

Lines changed: 476 additions & 1 deletion

File tree

‎include/bioimage_cpp/graph/region_adjacency_graph.hxx‎

Lines changed: 259 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,259 @@
1+
#pragma once
2+
3+
#include "bioimage_cpp/array_view.hxx"
4+
#include "bioimage_cpp/graph/undirected_graph.hxx"
5+
6+
#include <algorithm>
7+
#include <cstddef>
8+
#include <cstdint>
9+
#include <limits>
10+
#include <numeric>
11+
#include <stdexcept>
12+
#include <string>
13+
#include <thread>
14+
#include <type_traits>
15+
#include <unordered_set>
16+
#include <utility>
17+
#include <vector>
18+
19+
namespace bioimage_cpp::graph {
20+
21+
class RegionAdjacencyGraph : public UndirectedGraph {
22+
public:
23+
using UndirectedGraph::UndirectedGraph;
24+
25+
explicit RegionAdjacencyGraph(
26+
const std::uint64_t number_of_nodes,
27+
std::vector<std::uint64_t> shape
28+
)
29+
: UndirectedGraph(number_of_nodes),
30+
shape_(std::move(shape)) {
31+
}
32+
33+
[[nodiscard]] const std::vector<std::uint64_t> &shape() const {
34+
return shape_;
35+
}
36+
37+
private:
38+
std::vector<std::uint64_t> shape_;
39+
};
40+
41+
namespace detail {
42+
43+
using Edge = UndirectedGraph::Edge;
44+
45+
struct EdgeHash {
46+
std::size_t operator()(const Edge &edge) const {
47+
const auto first = static_cast<std::size_t>(edge.first);
48+
const auto second = static_cast<std::size_t>(edge.second);
49+
return first ^ (second + 0x9e3779b97f4a7c15ULL + (first << 6U) + (first >> 2U));
50+
}
51+
};
52+
53+
inline Edge edge_key(std::uint64_t u, std::uint64_t v) {
54+
if (v < u) {
55+
std::swap(u, v);
56+
}
57+
return {u, v};
58+
}
59+
60+
template <class T>
61+
std::uint64_t checked_label_to_node(const T value) {
62+
if constexpr (std::is_signed_v<T>) {
63+
if (value < 0) {
64+
throw std::invalid_argument("labels must not contain negative values");
65+
}
66+
}
67+
return static_cast<std::uint64_t>(value);
68+
}
69+
70+
template <class T>
71+
std::uint64_t max_label(const ConstArrayView<T> &labels) {
72+
const auto number_of_pixels = static_cast<std::size_t>(std::accumulate(
73+
labels.shape.begin(),
74+
labels.shape.end(),
75+
std::ptrdiff_t{1},
76+
[](const std::ptrdiff_t a, const std::ptrdiff_t b) { return a * b; }
77+
));
78+
std::uint64_t max_value = 0;
79+
for (std::size_t index = 0; index < number_of_pixels; ++index) {
80+
max_value = std::max(max_value, checked_label_to_node(labels.data[index]));
81+
}
82+
return max_value;
83+
}
84+
85+
inline std::size_t normalize_thread_count(
86+
const std::size_t requested,
87+
const std::size_t number_of_work_items
88+
) {
89+
if (number_of_work_items == 0) {
90+
return 1;
91+
}
92+
std::size_t n_threads = requested;
93+
if (n_threads == 0) {
94+
n_threads = std::thread::hardware_concurrency();
95+
if (n_threads == 0) {
96+
n_threads = 1;
97+
}
98+
}
99+
return std::max<std::size_t>(1, std::min(n_threads, number_of_work_items));
100+
}
101+
102+
template <class T>
103+
void add_edge_if_different(
104+
const T *data,
105+
const std::size_t first,
106+
const std::size_t second,
107+
std::unordered_set<Edge, EdgeHash> &edges
108+
) {
109+
const auto u = checked_label_to_node(data[first]);
110+
const auto v = checked_label_to_node(data[second]);
111+
if (u != v) {
112+
edges.insert(edge_key(u, v));
113+
}
114+
}
115+
116+
template <class T>
117+
void scan_2d_chunk(
118+
const T *data,
119+
const std::size_t height,
120+
const std::size_t width,
121+
const std::size_t y_begin,
122+
const std::size_t y_end,
123+
std::unordered_set<Edge, EdgeHash> &edges
124+
) {
125+
for (std::size_t y = y_begin; y < y_end; ++y) {
126+
const auto row_offset = y * width;
127+
for (std::size_t x = 0; x < width; ++x) {
128+
const auto pixel = row_offset + x;
129+
if (x + 1 < width) {
130+
add_edge_if_different(data, pixel, pixel + 1, edges);
131+
}
132+
if (y + 1 < height) {
133+
add_edge_if_different(data, pixel, pixel + width, edges);
134+
}
135+
}
136+
}
137+
}
138+
139+
template <class T>
140+
void scan_3d_chunk(
141+
const T *data,
142+
const std::size_t depth,
143+
const std::size_t height,
144+
const std::size_t width,
145+
const std::size_t z_begin,
146+
const std::size_t z_end,
147+
std::unordered_set<Edge, EdgeHash> &edges
148+
) {
149+
const auto slice_size = height * width;
150+
for (std::size_t z = z_begin; z < z_end; ++z) {
151+
const auto slice_offset = z * slice_size;
152+
for (std::size_t y = 0; y < height; ++y) {
153+
const auto row_offset = slice_offset + y * width;
154+
for (std::size_t x = 0; x < width; ++x) {
155+
const auto pixel = row_offset + x;
156+
if (x + 1 < width) {
157+
add_edge_if_different(data, pixel, pixel + 1, edges);
158+
}
159+
if (y + 1 < height) {
160+
add_edge_if_different(data, pixel, pixel + width, edges);
161+
}
162+
if (z + 1 < depth) {
163+
add_edge_if_different(data, pixel, pixel + slice_size, edges);
164+
}
165+
}
166+
}
167+
}
168+
}
169+
170+
inline std::vector<Edge> merge_edge_sets(
171+
const std::vector<std::unordered_set<Edge, EdgeHash>> &per_thread_edges
172+
) {
173+
std::unordered_set<Edge, EdgeHash> merged;
174+
for (const auto &edges : per_thread_edges) {
175+
merged.insert(edges.begin(), edges.end());
176+
}
177+
178+
std::vector<Edge> sorted_edges(merged.begin(), merged.end());
179+
std::sort(sorted_edges.begin(), sorted_edges.end());
180+
return sorted_edges;
181+
}
182+
183+
} // namespace detail
184+
185+
template <class T>
186+
RegionAdjacencyGraph grid_region_adjacency_graph(
187+
const ConstArrayView<T> &labels,
188+
const std::size_t number_of_threads
189+
) {
190+
if (labels.ndim() != 2 && labels.ndim() != 3) {
191+
throw std::invalid_argument(
192+
"labels must be a 2D or 3D array, got ndim=" +
193+
std::to_string(labels.ndim())
194+
);
195+
}
196+
for (const auto axis_size : labels.shape) {
197+
if (axis_size <= 0) {
198+
throw std::invalid_argument("labels must not have empty dimensions");
199+
}
200+
}
201+
202+
const auto max_node = detail::max_label(labels);
203+
if (max_node == std::numeric_limits<std::uint64_t>::max()) {
204+
throw std::overflow_error("maximum label is too large");
205+
}
206+
207+
std::vector<std::uint64_t> shape;
208+
shape.reserve(labels.shape.size());
209+
for (const auto axis_size : labels.shape) {
210+
shape.push_back(static_cast<std::uint64_t>(axis_size));
211+
}
212+
213+
const auto work_items = static_cast<std::size_t>(labels.shape[0]);
214+
const auto n_threads = detail::normalize_thread_count(number_of_threads, work_items);
215+
std::vector<std::unordered_set<detail::Edge, detail::EdgeHash>> per_thread_edges(n_threads);
216+
std::vector<std::thread> threads;
217+
threads.reserve(n_threads > 0 ? n_threads - 1 : 0);
218+
219+
const auto run_chunk = [&](const std::size_t thread_id) {
220+
const auto begin = thread_id * work_items / n_threads;
221+
const auto end = (thread_id + 1) * work_items / n_threads;
222+
if (labels.ndim() == 2) {
223+
detail::scan_2d_chunk(
224+
labels.data,
225+
static_cast<std::size_t>(labels.shape[0]),
226+
static_cast<std::size_t>(labels.shape[1]),
227+
begin,
228+
end,
229+
per_thread_edges[thread_id]
230+
);
231+
} else {
232+
detail::scan_3d_chunk(
233+
labels.data,
234+
static_cast<std::size_t>(labels.shape[0]),
235+
static_cast<std::size_t>(labels.shape[1]),
236+
static_cast<std::size_t>(labels.shape[2]),
237+
begin,
238+
end,
239+
per_thread_edges[thread_id]
240+
);
241+
}
242+
};
243+
244+
for (std::size_t thread_id = 1; thread_id < n_threads; ++thread_id) {
245+
threads.emplace_back(run_chunk, thread_id);
246+
}
247+
run_chunk(0);
248+
for (auto &thread : threads) {
249+
thread.join();
250+
}
251+
252+
auto graph = RegionAdjacencyGraph(max_node + 1, std::move(shape));
253+
for (const auto edge : detail::merge_edge_sets(per_thread_edges)) {
254+
graph.insert_edge(edge.first, edge.second);
255+
}
256+
return graph;
257+
}
258+
259+
} // namespace bioimage_cpp::graph

‎src/bindings/graph.cxx‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "graph.hxx"
22

3+
#include "bioimage_cpp/array_view.hxx"
4+
#include "bioimage_cpp/graph/region_adjacency_graph.hxx"
35
#include "bioimage_cpp/graph/undirected_graph.hxx"
46

57
#include <nanobind/ndarray.h>
@@ -20,10 +22,14 @@ namespace bioimage_cpp::bindings {
2022
namespace {
2123

2224
using Graph = graph::UndirectedGraph;
25+
using Rag = graph::RegionAdjacencyGraph;
2326
using UInt64Array = nb::ndarray<nb::numpy, std::uint64_t, nb::c_contig>;
2427
using ConstUInt64Array = nb::ndarray<nb::numpy, const std::uint64_t, nb::c_contig>;
2528
using Int64Array = nb::ndarray<nb::numpy, std::int64_t, nb::c_contig>;
2629

30+
template <class T>
31+
using LabelArray = nb::ndarray<nb::numpy, const T, nb::c_contig>;
32+
2733
void require_uv_array(const ConstUInt64Array &uvs, const char *argument_name) {
2834
if (uvs.ndim() != 2 || uvs.shape(1) != 2) {
2935
throw std::invalid_argument(
@@ -198,6 +204,26 @@ UInt64Array graph_edges_from_node_list(const Graph &graph, ConstUInt64Array node
198204
return extracted.first;
199205
}
200206

207+
template <class T>
208+
Rag grid_region_adjacency_graph_t(
209+
LabelArray<T> labels,
210+
const std::size_t number_of_threads
211+
) {
212+
std::vector<std::ptrdiff_t> shape(labels.ndim());
213+
for (std::size_t axis = 0; axis < labels.ndim(); ++axis) {
214+
shape[axis] = static_cast<std::ptrdiff_t>(labels.shape(axis));
215+
}
216+
217+
ConstArrayView<T> labels_view{
218+
labels.data(),
219+
shape,
220+
{},
221+
};
222+
223+
nb::gil_scoped_release release;
224+
return graph::grid_region_adjacency_graph<T>(labels_view, number_of_threads);
225+
}
226+
201227
} // namespace
202228

203229
void bind_graph(nb::module_ &m) {
@@ -264,6 +290,38 @@ void bind_graph(nb::module_ &m) {
264290
nb::arg("nodes")
265291
)
266292
.def("edgesFromNodeList", &graph_edges_from_node_list, nb::arg("nodes"));
293+
294+
nb::class_<Rag, Graph>(m, "RegionAdjacencyGraph")
295+
.def_prop_ro("shape", &Rag::shape);
296+
297+
m.def(
298+
"_grid_region_adjacency_graph_uint32",
299+
&grid_region_adjacency_graph_t<std::uint32_t>,
300+
nb::arg("labels"),
301+
nb::arg("number_of_threads"),
302+
"Build a region adjacency graph for uint32 labels."
303+
);
304+
m.def(
305+
"_grid_region_adjacency_graph_uint64",
306+
&grid_region_adjacency_graph_t<std::uint64_t>,
307+
nb::arg("labels"),
308+
nb::arg("number_of_threads"),
309+
"Build a region adjacency graph for uint64 labels."
310+
);
311+
m.def(
312+
"_grid_region_adjacency_graph_int32",
313+
&grid_region_adjacency_graph_t<std::int32_t>,
314+
nb::arg("labels"),
315+
nb::arg("number_of_threads"),
316+
"Build a region adjacency graph for int32 labels."
317+
);
318+
m.def(
319+
"_grid_region_adjacency_graph_int64",
320+
&grid_region_adjacency_graph_t<std::int64_t>,
321+
nb::arg("labels"),
322+
nb::arg("number_of_threads"),
323+
"Build a region adjacency graph for int64 labels."
324+
);
267325
}
268326

269327
} // namespace bioimage_cpp::bindings

0 commit comments

Comments
 (0)