Implement PolygonSelector

almarklein · almarklein · commit 46ccae36c620 · 2025-05-26T22:58:40.000+02:00
diff --git a/docs/source/api/graphics/ImageGraphic.rst b/docs/source/api/graphics/ImageGraphic.rst
@@ -48,6 +48,7 @@ Methods
     ImageGraphic.add_linear_region_selector
     ImageGraphic.add_linear_selector
     ImageGraphic.add_rectangle_selector
+    ImageGraphic.add_polygon_selector
     ImageGraphic.clear_event_handlers
     ImageGraphic.remove_event_handler
     ImageGraphic.reset_vmin_vmax
diff --git a/docs/source/api/graphics/LineCollection.rst b/docs/source/api/graphics/LineCollection.rst
@@ -53,6 +53,7 @@ Methods
     LineCollection.add_linear_region_selector
     LineCollection.add_linear_selector
     LineCollection.add_rectangle_selector
+    LineCollection.add_polygon_selector
     LineCollection.clear_event_handlers
     LineCollection.remove_event_handler
     LineCollection.remove_graphic
diff --git a/docs/source/api/graphics/LineGraphic.rst b/docs/source/api/graphics/LineGraphic.rst
@@ -47,6 +47,7 @@ Methods
     LineGraphic.add_linear_region_selector
     LineGraphic.add_linear_selector
     LineGraphic.add_rectangle_selector
+    LineGraphic.add_polygon_selector
     LineGraphic.clear_event_handlers
     LineGraphic.remove_event_handler
     LineGraphic.rotate
diff --git a/docs/source/api/graphics/LineStack.rst b/docs/source/api/graphics/LineStack.rst
@@ -53,6 +53,7 @@ Methods
     LineStack.add_linear_region_selector
     LineStack.add_linear_selector
     LineStack.add_rectangle_selector
+    LineStack.add_polygon_selector
     LineStack.clear_event_handlers
     LineStack.remove_event_handler
     LineStack.remove_graphic
diff --git a/fastplotlib/graphics/features/_selection_features.py b/fastplotlib/graphics/features/_selection_features.py
@@ -1,6 +1,7 @@
 from typing import Sequence
 
 import numpy as np
+import pygfx as gfx
 
 from ...utils import mesh_masks
 from ._base import GraphicFeature, GraphicFeatureEvent, block_reentrance
@@ -340,3 +341,96 @@ def set_value(self, selector, value: Sequence[float]):
 
         # calls any events
         self._call_event_handlers(event)
+
+
+class PolygonSelectionFeature(GraphicFeature):
+    event_info_spec = [
+        {
+            "dict key": "value",
+            "type": "np.ndarray",
+            "description": "new array of points that represents the polygon selection",
+        },
+    ]
+
+    event_extra_attrs = [
+        {
+            "attribute": "get_selected_indices",
+            "type": "callable",
+            "description": "returns indices under the selector",
+        },
+        {
+            "attribute": "get_selected_data",
+            "type": "callable",
+            "description": "returns data under the selector",
+        },
+    ]
+
+    def __init__(
+        self,
+        value: Sequence[tuple[float]],
+        limits: tuple[float, float, float, float],
+    ):
+        super().__init__()
+
+        self._limits = limits
+        self._value = np.asarray(value).reshape(-1, 3).astype(float)
+
+    @property
+    def value(self) -> np.ndarray[float]:
+        """
+        The array of the polygon, in data space
+        """
+        return self._value
+
+    @block_reentrance
+    def set_value(self, selector, value: Sequence[tuple[float]]):
+        """
+        Set the selection of the rectangle selector.
+
+        Parameters
+        ----------
+        selector: PolygonSelector
+
+        value: array
+            new values (3D points) of the selection
+        """
+
+        value = np.asarray(value, dtype=np.float32)
+
+        if not value.shape[1] == 3:
+            raise TypeError(
+                "Selection must be an array, tuple, list, or sequence of the shape Nx3."
+            )
+
+        # # clip values if they are beyond the limits
+        # value[:, 0] = value[:2].clip(self._limits[0], self._limits[1])
+        # # clip y
+        # value[:, 1] = value[2:].clip(self._limits[2], self._limits[3])
+
+        self._value = value
+
+        # TODO: Update the fill mesh
+        # selector.fill.geometry.positions = ...
+
+        edge_geometry = selector.edge.geometry
+
+        # Need larger buffer?
+        if len(value) > edge_geometry.positions.nitems:
+            arr = np.zeros((edge_geometry.positions.nitems * 2, 3), np.float32)
+            edge_geometry.positions = gfx.Buffer(arr)
+
+        edge_geometry.positions.data[: len(value)] = value
+        edge_geometry.positions.draw_range = 0, len(value)
+        edge_geometry.positions.update_full()
+
+        # send event
+        if len(self._event_handlers) < 1:
+            return
+
+        event = GraphicFeatureEvent("selection", {"value": self.value})
+
+        event.get_selected_indices = selector.get_selected_indices
+        event.get_selected_data = selector.get_selected_data
+
+        # calls any events
+        self._call_event_handlers(event)
diff --git a/fastplotlib/graphics/image.py b/fastplotlib/graphics/image.py
@@ -5,7 +5,7 @@
 
 from ..utils import quick_min_max
 from ._base import Graphic
-from .selectors import LinearSelector, LinearRegionSelector, RectangleSelector
+from .selectors import LinearSelector, LinearRegionSelector, RectangleSelector, PolygonSelector
 from .features import (
     TextureArray,
     ImageCmap,
@@ -437,3 +437,46 @@ def add_rectangle_selector(
         selector.offset = selector.offset + (0.0, 0.0, self.offset[-1] + 1)
 
         return selector
+
+    def add_polygon_selector(
+        self,
+        selection: List[tuple[float, float]] = None,
+        fill_color=(0, 0, 0.35, 0.2),
+        **kwargs,
+    ) -> PolygonSelector:
+        """
+        Add a :class:`.PolygonSelector`.
+
+        Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them
+        from a plot area just like any other ``Graphic``.
+
+        Parameters
+        ----------
+        selection: List of positions, optional
+            initial points for the polygon
+
+        """
+        # default selection is 25% of the diagonal
+        if selection is None:
+            diagonal = math.sqrt(
+                self._data.value.shape[0] ** 2 + self._data.value.shape[1] ** 2
+            )
+
+            selection = (0, int(diagonal / 4), 0, int(diagonal / 4))
+
+        # min/max limits are image shape
+        # rows are ys, columns are xs
+        limits = (0, self._data.value.shape[1], 0, self._data.value.shape[0])
+
+        selector = PolygonSelector(
+            fill_color=fill_color,
+            parent=self,
+            **kwargs,
+        )
+
+        self._plot_area.add_graphic(selector, center=False)
+
+        # place above this graphic
+        selector.offset = selector.offset + (0.0, 0.0, self.offset[-1] + 1)
+
+        return selector
diff --git a/fastplotlib/graphics/line.py b/fastplotlib/graphics/line.py
@@ -5,7 +5,7 @@
 import pygfx
 
 from ._positions_base import PositionsGraphic
-from .selectors import LinearRegionSelector, LinearSelector, RectangleSelector
+from .selectors import LinearRegionSelector, LinearSelector, RectangleSelector, PolygonSelector
 from .features import (
     Thickness,
     VertexPositions,
@@ -245,7 +245,7 @@ def add_rectangle_selector(
         self,
         selection: tuple[float, float, float, float] = None,
         **kwargs,
-    ) -> RectangleSelector:
+    ) ->  RectangleSelector:
         """
         Add a :class:`.RectangleSelector`.
 
@@ -288,6 +288,31 @@ def add_rectangle_selector(
 
         return selector
 
+    def add_polygon_selector(
+        self,
+        selection: List[tuple[float, float]] = None,
+        **kwargs,
+    ) -> PolygonSelector:
+        """
+        Add a :class:`.PolygonSelector`.
+
+        Selectors are just ``Graphic`` objects, so you can manage, remove, or delete them from a
+        plot area just like any other ``Graphic``.
+
+        Parameters
+        ----------
+        selection: (float, float, float, float), optional
+            initial (xmin, xmax, ymin, ymax) of the selection
+        """
+        selector = PolygonSelector(
+            parent=self,
+            **kwargs,
+        )
+
+        self._plot_area.add_graphic(selector, center=False)
+
+        return selector
+
     # TODO: this method is a bit of a mess, can refactor later
     def _get_linear_selector_init_args(
         self, axis: str, padding
diff --git a/fastplotlib/graphics/line_collection.py b/fastplotlib/graphics/line_collection.py
@@ -7,7 +7,7 @@
 from ..utils import parse_cmap_values
 from ._collection_base import CollectionIndexer, GraphicCollection, CollectionFeature
 from .line import LineGraphic
-from .selectors import LinearRegionSelector, LinearSelector, RectangleSelector
+from .selectors import LinearRegionSelector, LinearSelector, RectangleSelector, PolygonSelector
 
 
 class _LineCollectionProperties:
@@ -486,6 +486,36 @@ def add_rectangle_selector(
 
         return selector
 
+    def add_polygon_selector(
+        self,
+        selection: List[tuple[float, float]] = None,
+        **kwargs,
+    ) -> PolygonSelector:
+        """
+        Add a :class:`.PolygonSelector`. Selectors are just ``Graphic`` objects, so you can manage,
+        remove, or delete them from a plot area just like any other ``Graphic``.
+
+        Parameters
+        ----------
+        selection: List of positions, optional
+            initial points for the polygon
+        """
+        bbox = self.world_object.get_world_bounding_box()
+
+
+        if selection is not None:
+            selection = []  # TODO: fill selection
+
+
+        selector = PolygonSelector(
+            parent=self,
+            **kwargs,
+        )
+
+        self._plot_area.add_graphic(selector, center=False)
+
+        return selector
+
     def _get_linear_selector_init_args(self, axis, padding):
         # use bbox to get size and center
         bbox = self.world_object.get_world_bounding_box()
diff --git a/fastplotlib/graphics/selectors/_polygon.py b/fastplotlib/graphics/selectors/_polygon.py
