Update MWS checks and examples, speed it up (#5)

Author: constantinpape

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(

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()

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()

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

include/bioimage_cpp/segmentation/mutex_watershed.hxx

@@ -8,7 +8,6 @@
 #include <numeric>
 #include <stdexcept>
 #include <string>
-#include <unordered_map>
 #include <unordered_set>
 #include <vector>
 
@@ -45,6 +44,12 @@ private:
 
 using MutexStorage = std::vector<std::unordered_set<std::uint64_t>>;
 
+template <class T>
+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<std::ptrdiff_t> c_order_strides(const std::vector<std::ptrdif
     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,
@@ -178,58 +167,59 @@ void mutex_watershed_grid(
     }
 
     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;
+    std::vector<WeightedGridEdge<T>> edge_order;
+    edge_order.reserve(static_cast<std::size_t>(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<T>{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<std::size_t>(number_of_nodes));
     MutexStorage mutexes(static_cast<std::size_t>(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<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)) {
+        if (root_u == root_v) {
             continue;
         }
 
         const bool is_mutex_edge = channel >= number_of_attractive_channels;
         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<std::uint64_t, std::uint64_t> label_map;
+    std::vector<std::uint64_t> root_labels(static_cast<std::size_t>(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<std::uint64_t>(label_map.size() + 1);
-            found = label_map.emplace(root, next_label).first;
+        auto &label = root_labels[static_cast<std::size_t>(root)];
+        if (label == 0) {
+            label = next_label;
+            ++next_label;
         }
-        out.data[node] = found->second;
+        out.data[node] = label;
     }
 }
 

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]