diff --git a/development/segmentation/_mutex_watershed_equivalence.py b/development/segmentation/_mutex_watershed_equivalence.py index ea51b3f..fbcb6fa 100644 --- a/development/segmentation/_mutex_watershed_equivalence.py +++ b/development/segmentation/_mutex_watershed_equivalence.py @@ -65,7 +65,7 @@ def run_affogato_reference( offsets: list[tuple[int, ...]], number_of_attractive_channels: int, ) -> np.ndarray: - from elf.segmentation.mutex_watershed import mutex_watershed + from affogato.segmentation import compute_mws_segmentation spatial_ndim = len(offsets[0]) if number_of_attractive_channels != spatial_ndim: @@ -74,7 +74,12 @@ def run_affogato_reference( f"spatial axis, got ndim={spatial_ndim}, attractive channels=" f"{number_of_attractive_channels}" ) - return mutex_watershed(affinities.copy(), offsets, strides=[1] * spatial_ndim) + affs = affinities.copy() + affs[:spatial_ndim] *= -1 + affs[:spatial_ndim] += 1 + return compute_mws_segmentation(affs, offsets, + number_of_attractive_channels=number_of_attractive_channels, + strides=[1] * spatial_ndim) def _load_validation_metrics(): @@ -126,21 +131,41 @@ def compare_segmentations( return metrics -def time_function( - run: Callable[[np.ndarray, list[tuple[int, ...]], int], np.ndarray], +def time_functions_interleaved( + first: Callable[[np.ndarray, list[tuple[int, ...]], int], np.ndarray], + second: 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): +) -> tuple[list[float], np.ndarray, list[float], np.ndarray]: + def timed_call(run): start = perf_counter() result = run(affinities, offsets, number_of_attractive_channels) - timings.append(perf_counter() - start) - assert result is not None - return timings, result + return perf_counter() - start, result + + # Warm up imports, extension initialization, and first allocation paths + # outside the timed section. + first(affinities, offsets, number_of_attractive_channels) + second(affinities, offsets, number_of_attractive_channels) + + first_timings = [] + second_timings = [] + first_result = None + second_result = None + for repeat in range(repeats): + if repeat % 2 == 0: + first_time, first_result = timed_call(first) + second_time, second_result = timed_call(second) + else: + second_time, second_result = timed_call(second) + first_time, first_result = timed_call(first) + first_timings.append(first_time) + second_timings.append(second_time) + + assert first_result is not None + assert second_result is not None + return first_timings, first_result, second_timings, second_result def print_report( @@ -193,11 +218,13 @@ def run_check( 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 + ref_timings, ref_seg, bic_timings, bic_seg = time_functions_interleaved( + run_affogato_reference, + run_bioimage_cpp, + affs, + used_offsets, + attractive_channels, + repeats, ) metrics = compare_segmentations(bic_seg, ref_seg) print_report( diff --git a/development/segmentation/check_mutex_watershed_2d.py b/development/segmentation/check_mutex_watershed_2d.py index 8960639..22f5c5a 100644 --- a/development/segmentation/check_mutex_watershed_2d.py +++ b/development/segmentation/check_mutex_watershed_2d.py @@ -21,7 +21,7 @@ def main() -> None: type=int, nargs=2, metavar=("Y", "X"), - default=(128, 128), + default=(512, 512), help="Spatial crop shape for the 2D check.", ) args = parser.parse_args() diff --git a/development/segmentation/check_mutex_watershed_3d.py b/development/segmentation/check_mutex_watershed_3d.py index 45ae843..3d65e2b 100644 --- a/development/segmentation/check_mutex_watershed_3d.py +++ b/development/segmentation/check_mutex_watershed_3d.py @@ -15,7 +15,7 @@ def main() -> None: type=int, nargs=3, metavar=("Z", "Y", "X"), - default=(6, 96, 96), + default=(6, 512, 512), help="Spatial crop shape for the 3D check.", ) args = parser.parse_args() diff --git a/examples/segmentation/mutex_watershed.py b/examples/segmentation/mutex_watershed.py index 7477e48..a47aa2f 100644 --- a/examples/segmentation/mutex_watershed.py +++ b/examples/segmentation/mutex_watershed.py @@ -18,15 +18,6 @@ def mws_bic(affinities, offsets): 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] @@ -42,12 +33,7 @@ def main(): 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) + segmentation = mws_bic(affinities, offsets) with open_file(data_path, "r") as f: raw = f["raw"][0] if check_2d else f["raw"][:] diff --git a/include/bioimage_cpp/segmentation/mutex_watershed.hxx b/include/bioimage_cpp/segmentation/mutex_watershed.hxx index b96eb71..47cf445 100644 --- a/include/bioimage_cpp/segmentation/mutex_watershed.hxx +++ b/include/bioimage_cpp/segmentation/mutex_watershed.hxx @@ -8,7 +8,6 @@ #include #include #include -#include #include #include @@ -45,6 +44,12 @@ private: using MutexStorage = std::vector>; +template +struct WeightedGridEdge { + T weight; + std::uint64_t id; +}; + inline bool check_mutex( const std::uint64_t first, const std::uint64_t second, @@ -97,22 +102,6 @@ inline std::vector c_order_strides(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, @@ -178,36 +167,33 @@ void mutex_watershed_grid( } 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; + std::vector> edge_order; + edge_order.reserve(static_cast(number_of_edges)); + for (std::uint64_t edge_id = 0; edge_id < number_of_edges; ++edge_id) { + if (valid_edges.data[edge_id] != 0) { + edge_order.push_back(WeightedGridEdge{affinities.data[edge_id], edge_id}); } - return first_weight > second_weight; + } + std::sort(edge_order.begin(), edge_order.end(), [](const auto &first, const auto &second) { + if (first.weight == second.weight) { + return first.id < second.id; + } + 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) { - if (valid_edges.data[edge_id] == 0) { - continue; - } - + for (const auto &edge : edge_order) { + const auto edge_id = edge.id; 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)) { + if (root_u == root_v) { continue; } @@ -215,21 +201,25 @@ void mutex_watershed_grid( if (is_mutex_edge) { insert_mutex(root_u, root_v, mutexes); } else { + if (check_mutex(root_u, root_v, mutexes)) { + continue; + } 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; + std::vector root_labels(static_cast(number_of_nodes), 0); + std::uint64_t next_label = 1; 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; + auto &label = root_labels[static_cast(root)]; + if (label == 0) { + label = next_label; + ++next_label; } - out.data[node] = found->second; + out.data[node] = label; } } diff --git a/pyproject.toml b/pyproject.toml index f3522ba..40aece5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ test = ["pytest"] [tool.scikit-build] cmake.version = ">=3.21" +cmake.build-type = "Release" wheel.packages = ["src/bioimage_cpp"] [tool.pytest.ini_options]