Enh: add IncrementalWatershedCut for interactive seeded watershed

doc/source/python/watershed.rst

@@ -9,7 +9,11 @@ Labelisation watershed
 
     labelisation_watershed
     labelisation_seeded_watershed
+    IncrementalWatershedCut
 
 .. autofunction:: higra.labelisation_watershed
 
-.. autofunction:: higra.labelisation_seeded_watershed
+.. autofunction:: higra.labelisation_seeded_watershed
+
+.. autoclass:: higra.IncrementalWatershedCut
+    :members:

higra/algo/py_watershed.cpp

@@ -62,6 +62,38 @@ namespace py_watershed {
 
         add_type_overloads<def_labelisation_watershed<hg::ugraph>, HG_TEMPLATE_NUMERIC_TYPES>(m, "");
         add_type_overloads<def_labelisation_seeded_watershed<hg::ugraph>, HG_TEMPLATE_NUMERIC_TYPES>(m, "");
+
+        auto cls = py::class_<hg::incremental_watershed_cut>(m, "IncrementalWatershedCut",
+                "Incremental seeded watershed cut based on the binary partition tree.\n\n"
+                "Provides efficient computation of seeded watershed cuts in an interactive "
+                "segmentation setting, where seeds are added and removed incrementally.");
+
+        cls.def(py::init<const hg::tree &, const hg::ugraph &>(),
+                py::arg("bpt"),
+                py::arg("mst"),
+                "Create an incremental watershed cut object from a binary partition tree and its MST.");
+
+        cls.def("_add_seeds",
+                [](hg::incremental_watershed_cut &self,
+                   const pyarray<hg::index_t> &seed_vertices,
+                   const pyarray<hg::index_t> &seed_labels) {
+                    self.add_seeds(seed_vertices, seed_labels);
+                },
+                "Add seeds with given labels.",
+                py::arg("seed_vertices"),
+                py::arg("seed_labels"));
+
+        cls.def("_remove_seeds",
+                [](hg::incremental_watershed_cut &self,
+                   const pyarray<hg::index_t> &seed_vertices) {
+                    self.remove_seeds(seed_vertices);
+                },
+                "Remove seeds.",
+                py::arg("seed_vertices"));
+
+        cls.def("_get_labeling",
+                &hg::incremental_watershed_cut::get_labeling,
+                "Compute and return the current vertex labeling.");
     }
 }
 

higra/algo/watershed.py

@@ -73,3 +73,103 @@ def labelisation_seeded_watershed(graph, edge_weights, vertex_seeds, background_
 
     labels = hg.delinearize_vertex_weights(labels, graph)
     return labels
+
+
+class IncrementalWatershedCut:
+    """
+    Incremental seeded watershed cut based on the binary partition tree.
+
+    This class provides an efficient way to compute seeded watershed cuts
+    in an interactive segmentation setting, where seeds are added and removed
+    incrementally. Instead of recomputing the full watershed from scratch at
+    each interaction, only the affected regions are updated.
+
+    The algorithm maintains a canonical BPT and a visitCount array to identify
+    watershed edges. The labeling is obtained by BFS on the MST forest.
+
+    :Complexity:
+
+    Construction is dominated by the canonical BPT and MST construction in
+    :math:`\\mathcal{O}(n \\log n)` with :math:`n` the number of edges in the
+    input graph (dominated by the edge sort).
+
+    :func:`add_seeds` on a batch of :math:`K` seeds runs in
+    :math:`\\mathcal{O}(K \\cdot d + S)` where :math:`d` is the height of the
+    BPT (:math:`\\mathcal{O}(\\log N)` for balanced trees, :math:`\\mathcal{O}(N)`
+    worst case, with :math:`N` the number of vertices) and :math:`S` is the
+    total size of the :math:`K` MST-forest components relabeled in the second
+    pass. By the :math:`\\text{visit\\_count} == 2` invariant these :math:`K`
+    components are pairwise disjoint, hence :math:`S \\leq N`.
+
+    :func:`remove_seeds` on a batch of :math:`K` seeds runs in
+    :math:`\\mathcal{O}(K \\cdot d + S')` where :math:`d` is as above and
+    :math:`S'` is the total work performed by the BFS calls during the de-cut
+    and relabel phase. When the de-cuts of the batch touch mostly disjoint
+    regions, :math:`S' \\leq N`; in the worst case of cascading merges within
+    the same batch (each removal extends a growing super-component),
+    :math:`S'` can grow to :math:`\\mathcal{O}(K \\cdot N)`.
+
+    :func:`get_labeling` is :math:`\\mathcal{O}(1)` (the labeling is
+    maintained incrementally).
+
+    Reference:
+
+        Q. Lebon, J. Lefevre, J. Cousty, B. Perret.
+        `Interactive Segmentation With Incremental Watershed Cuts <https://hal.science/hal-04069187v1>`_.
+        CIARP 2023.
+
+    :param graph: input graph (must be connected)
+    :param edge_weights: Weights on the edges of the graph
+    """
+
+    def __init__(self, graph, edge_weights):
+        tree, _ = hg.bpt_canonical(graph, edge_weights, compute_mst=True)
+        mst = hg.CptBinaryHierarchy.get_mst(tree)
+        self._impl = hg.cpp.IncrementalWatershedCut(tree, mst)
+        self._graph = graph
+
+    def add_seeds(self, seed_vertices, seed_labels):
+        """
+        Add seeds to the current watershed cut.
+
+        Each seed is defined by a vertex index and a label. Two seeds cannot share
+        the same vertex but can share the same label (resulting in merged regions
+        in the output labeling). Labels must be non-zero (0 is reserved for
+        unlabeled/background vertices).
+
+        Re-adding a seed on a vertex that is already a seed (regardless of the
+        label) raises an exception. To change the label of an existing seed,
+        call :func:`remove_seeds` first and then :func:`add_seeds` with the new
+        label.
+
+        :param seed_vertices: 1d array of seed vertex indices
+        :param seed_labels: 1d array of seed labels (same size as seed_vertices)
+        """
+        seed_vertices = np.asarray(seed_vertices, dtype=np.int64).ravel()
+        seed_labels = np.asarray(seed_labels, dtype=np.int64).ravel()
+        if seed_vertices.size != seed_labels.size:
+            raise ValueError("seed_vertices and seed_labels must have the same size.")
+        self._impl._add_seeds(seed_vertices, seed_labels)
+
+    def remove_seeds(self, seed_vertices):
+        """
+        Remove seeds from the current watershed cut.
+
+        :param seed_vertices: 1d array of seed vertex indices to remove
+        """
+        seed_vertices = np.asarray(seed_vertices, dtype=np.int64).ravel()
+        if seed_vertices.size == 0:
+            return
+        self._impl._remove_seeds(seed_vertices)
+
+    def get_labeling(self):
+        """
+        Compute and return the current vertex labeling.
+
+        Vertices with no seed in their component are labeled 0 (background).
+
+        :return: A labeling of the graph vertices
+        """
+        labels = self._impl._get_labeling()
+        labels = hg.delinearize_vertex_weights(labels, self._graph)
+        return labels