Fix hierarchy watershed to reproduce the same behaviour as on BPT

pylene/include/mln/morpho/alphatree.hpp

@@ -39,11 +39,11 @@ namespace mln::morpho
   {
     /// \brief Canvas for the edges in the alphatree. Using different data
     /// structures related to the type of the edges.
-    template <typename P, typename N, typename W>
+    template <typename P, typename N, typename W, bool HQ>
     class alphatree_edges;
 
-    template <typename P, typename N, typename W>
-    requires(std::is_integral_v<W>&& std::is_unsigned_v<W> && sizeof(W) <= 2) class alphatree_edges<P, N, W>
+    template <typename P, typename N, typename W, bool HQ>
+    requires(std::is_integral_v<W>&& std::is_unsigned_v<W> && sizeof(W) <= 2 && HQ) class alphatree_edges<P, N, W, HQ>
     {
     public:
       void                push(int dir, W w, P p) { m_cont.insert(dir, w, p); }
@@ -64,11 +64,14 @@ namespace mln::morpho
       W w;
     };
 
-    template <typename P, typename N, typename W>
+    template <typename P, typename N, typename W, bool HQ>
     class alphatree_edges
     {
     public:
-      void                push(int dir, W w, P p) { m_cont.push_back({p, p + cn.after_offsets()[dir], w}); }
+      void push(int dir, W w, P p) { m_cont.push_back({p, p + cn.after_offsets()[dir], w}); }
+
+      void push(P p, P q, W w) { m_cont.push_back({p, q, w}); }
+
       std::tuple<P, P, W> pop()
       {
         assert(m_current < m_cont.size());
@@ -81,10 +84,13 @@ namespace mln::morpho
         assert(m_current < m_cont.size());
         return m_cont[m_current].w;
       }
+
       bool empty() const { return m_cont.size() == m_current; }
+
       void on_finish_insert()
       {
-        std::sort(m_cont.begin(), m_cont.end(), [](const edge_t<P, W>& a, const edge_t<P, W>& b) { return a.w < b.w; });
+        std::stable_sort(m_cont.begin(), m_cont.end(),
+                         [](const edge_t<P, W>& a, const edge_t<P, W>& b) { return a.w < b.w; });
       }
 
     private:
@@ -115,8 +121,6 @@ namespace mln::morpho
     template <class E, class J>
     void alphatree_compute_flatzones(E& edges, J zpar)
     {
-      canvas::impl::union_find_init_par(zpar);
-
       while (!edges.empty() && edges.top() == 0)
       {
         const auto [p, q, w] = edges.pop();
@@ -312,11 +316,11 @@ namespace mln::morpho
       return {std::move(t), std::move(node_map)};
     }
 
-    template <class I, class N, class F,
+    template <class I, class N, class F, bool HQ = true,
               class M = edge_t<image_point_t<I>, std::invoke_result_t<F, image_value_t<I>, image_value_t<I>>>>
     std::pair<component_tree<std::invoke_result_t<F, image_value_t<I>, image_value_t<I>>>, image_ch_value_t<I, int>> //
-    __alphatree(I input, N nbh, F distance, bool canonize_tree = true, std::vector<M>* mst = nullptr,
-                std::size_t* nb_leaves = nullptr)
+    __alphatree(I input, N nbh, F distance, bool canonize_tree = true, bool compute_flatzones = true,
+                std::vector<M>* mst = nullptr)
     {
       static_assert(mln::is_a<I, mln::details::Image>());
       static_assert(mln::is_a<N, mln::details::Neighborhood>());
@@ -330,23 +334,22 @@ namespace mln::morpho
       static_assert(std::is_same<M, edge_t<P, W>>());
 
       // 1. Get the list of edges
-      auto edges = alphatree_edges<P, N, W>();
+      auto edges = alphatree_edges<P, N, W, HQ>();
       internal::alphatree_compute_edges(std::move(input), std::move(nbh), std::move(distance), edges);
 
       std::size_t              flatzones_count;
       image_ch_value_t<I, int> node_map = imchvalue<int>(input).set_init_value(-1);
       {
         image_ch_value_t<I, P> zpar = imchvalue<P>(input);
-        // 2. Compute flat zone of the image
-        internal::alphatree_compute_flatzones(edges, zpar);
+        canvas::impl::union_find_init_par(zpar);
+
+        if (compute_flatzones)
+          internal::alphatree_compute_flatzones(edges, zpar);
 
         // 3. Compute a node_id for each flat zone
         flatzones_count = internal::alphatree_create_nodemap(node_map, zpar);
       }
 
-      if (nb_leaves != nullptr)
-        *nb_leaves = flatzones_count;
-
       return alphatree_from_graph<W>(edges, node_map, flatzones_count, canonize_tree, mst);
     }
   } // namespace internal

pylene/include/mln/morpho/watershed_hierarchy.hpp

@@ -5,7 +5,7 @@
 
 namespace mln::morpho
 {
-  // TODO Add the attribute parameter
+  // FIXME Pass a std::function instead of an accumulator ?
 
   /// Compute the watershed hierarchy of an image
   ///
@@ -37,7 +37,7 @@ namespace mln::morpho
       {
         int parent = tree.parent[i];
 
-        if (i <= nb_leaves && tree.values[parent] != tree.values[i])
+        if (i <= (n - nb_leaves - 1) && tree.values[parent] != tree.values[i])
           res[i] = attribute[i];
 
         res[parent] = std::max(res[parent], res[i]);
@@ -47,10 +47,10 @@ namespace mln::morpho
       return res;
     }
 
-    template <typename I, typename W, typename E, typename N, typename A>
+    template <typename I, typename N, typename W, typename E, typename A>
     std::pair<component_tree<A>, image_ch_value_t<I, int>>
     watershed(const component_tree<W>& tree, image_ch_value_t<I, int> node_map, std::vector<E> mst,
-              const std::vector<A>& attribute, N nbh, std::size_t nb_leaves)
+              const std::vector<A>& attribute, std::size_t nb_leaves)
     {
       auto computed_attribute = get_computed_attribute(tree, attribute, nb_leaves);
 
@@ -64,26 +64,16 @@ namespace mln::morpho
         min_computed_attributes[parent] = std::min<A>(min_computed_attributes[parent], computed_attribute[i]);
       }
 
-      auto edges = alphatree_edges<image_point_t<I>, N, A>();
+      auto edges = alphatree_edges<image_point_t<I>, N, A, false>();
       for (std::size_t i = 0; i < mst.size(); ++i)
       {
         auto [p, q, _] = mst[i];
         A new_weight   = min_computed_attributes[nb_leaves - i - 2];
-
-        int dir = 0;
-        for (auto nb : nbh.after(p))
-        {
-          if (q == nb)
-            break;
-
-          ++dir;
-        }
-
-        edges.push(dir, new_weight, p);
+        edges.push(p, q, new_weight);
       }
       edges.on_finish_insert();
 
-      return internal::alphatree_from_graph<A>(edges, node_map, nb_leaves, true);
+      return internal::alphatree_from_graph<A>(edges, node_map, nb_leaves, false);
     }
   } // namespace internal
 
@@ -92,15 +82,13 @@ namespace mln::morpho
   watershed_hierarchy(I input, Accu acc, N nbh, F distance)
   {
     std::vector<internal::edge_t<image_point_t<I>, std::invoke_result_t<F, image_value_t<I>, image_value_t<I>>>> mst;
-    std::size_t nb_leaves;
-    auto [tree, nm] = internal::__alphatree(input, nbh, distance, false, &mst, &nb_leaves);
+    auto [tree, nm] = internal::__alphatree<I, N, F, false>(input, nbh, distance, false, false, &mst);
 
     auto attribute = tree.compute_attribute_on_points(nm, acc);
 
-    // FIXME Maybe an option in alpha tree function is better
     auto node_count = tree.parent.size();
     mln::for_each(nm, [node_count](int& id) { id = static_cast<int>(node_count) - id - 1; });
 
-    return internal::watershed<I>(tree, nm, mst, attribute, nbh, nb_leaves);
+    return internal::watershed<I, N>(tree, nm, mst, attribute, input.domain().size());
   }
 } // namespace mln::morpho
\ No newline at end of file

tests/morpho/alphatree.cpp

@@ -161,7 +161,8 @@ TEST(Morpho, AlphaTreeMST)
   std::vector<E> mst;
 
   auto [t, _] = mln::morpho::internal::__alphatree(
-      ima, mln::c4, [](const auto& a, const auto& b) -> W { return mln::functional::l2dist_t<>()(a, b); }, true, &mst);
+      ima, mln::c4, [](const auto& a, const auto& b) -> W { return mln::functional::l2dist_t<>()(a, b); }, true, true,
+      &mst);
 
   for (std::size_t i = 0; i < expected_mst.size(); ++i)
   {