fastplotlib · kushalkolar · Jul 11, 2023 · Jul 11, 2023 · Jul 11, 2023
diff --git a/fastplotlib/graphics/selectors/__init__.py b/fastplotlib/graphics/selectors/__init__.py
@@ -1,10 +1,12 @@
 from ._linear import LinearSelector
 from ._linear_region import LinearRegionSelector
+from ._polygon import PolygonSelector
 
 from ._sync import Synchronizer
 
 __all__ = [
     "LinearSelector",
     "LinearRegionSelector",
+    "PolygonSelector",
     "Synchronizer",
 ]
diff --git a/fastplotlib/graphics/selectors/_polygon.py b/fastplotlib/graphics/selectors/_polygon.py
@@ -0,0 +1,138 @@
+from typing import *
+
+import numpy as np
+
+import pygfx
+
+from ._base_selector import BaseSelector, MoveInfo
+from .._base import Graphic
+
+
+class PolygonSelector(Graphic, BaseSelector):
+    def __init__(
+        self,
+        edge_color="magenta",
+        edge_width: float = 3,
+        parent: Graphic = None,
+        name: str = None,
+    ):
+        Graphic.__init__(self, name=name)
+
+        self.parent = parent
+
+        group = pygfx.Group()
+
+        self._set_world_object(group)
+
+        self.edge_color = edge_color
+        self.edge_width = edge_width
+
+        self._move_info: MoveInfo = None
+
+        self._current_mode = None
+
+    def get_vertices(self) -> np.ndarray:
+        """Get the vertices for the polygon"""
+        vertices = list()
+        for child in self.world_object.children:
+            vertices.append(child.geometry.positions.data[:, :2])
+
+        return np.vstack(vertices)
+
+    def _add_plot_area_hook(self, plot_area):
+        self._plot_area = plot_area
+
+        # click to add new segment
+        self._plot_area.renderer.add_event_handler(self._add_segment, "click")
+
+        # pointer move to change endpoint of segment
+        self._plot_area.renderer.add_event_handler(self._move_segment_endpoint, "pointer_move")
+
+        # click to finish existing segment
+        self._plot_area.renderer.add_event_handler(self._finish_segment, "click")
+
+        # double click to finish polygon
+        self._plot_area.renderer.add_event_handler(self._finish_polygon, "double_click")
+
+        self.position_z = len(self._plot_area) + 10
+
+    def _add_segment(self, ev):
+        """After click event, adds a new line segment"""
+        self._current_mode = "add"
+
+        last_position = self._plot_area.map_screen_to_world(ev)
+        self._move_info = MoveInfo(last_position=last_position, source=None)
+
+        # line with same position for start and end until mouse moves
+        data = np.array([last_position, last_position])
+
+        new_line = pygfx.Line(
+            geometry=pygfx.Geometry(positions=data.astype(np.float32)),
+            material=pygfx.LineMaterial(thickness=self.edge_width, color=pygfx.Color(self.edge_color))
+        )
+
+        self.world_object.add(new_line)
+
+    def _move_segment_endpoint(self, ev):
+        """After mouse pointer move event, moves endpoint of current line segment"""
+        if self._move_info is None:
+            return
+        self._current_mode = "move"
+
+        world_pos = self._plot_area.map_screen_to_world(ev)
+
+        if world_pos is None:
+            return
+
+        # change endpoint
+        self.world_object.children[-1].geometry.positions.data[1] = np.array([world_pos]).astype(np.float32)
+        self.world_object.children[-1].geometry.positions.update_range()
+
+    def _finish_segment(self, ev):
+        """After click event, ends a line segment"""
+        # should start a new segment
+        if self._move_info is None:
+            return
+
+        # since both _add_segment and _finish_segment use the "click" callback
+        # this is to block _finish_segment right after a _add_segment call
+        if self._current_mode == "add":
+            return
+
+        # just make move info None so that _move_segment_endpoint is not called
+        # and _add_segment gets triggered for "click"
+        self._move_info = None
+
+        self._current_mode = "finish-segment"
+
+    def _finish_polygon(self, ev):
+        """finishes the polygon, disconnects events"""
+        world_pos = self._plot_area.map_screen_to_world(ev)
+
+        if world_pos is None:
+            return
+
+        # make new line to connect first and last vertices
+        data = np.vstack([
+            world_pos,
+            self.world_object.children[0].geometry.positions.data[0]
+        ])
+
+        print(data)
+
+        new_line = pygfx.Line(
+            geometry=pygfx.Geometry(positions=data.astype(np.float32)),
+            material=pygfx.LineMaterial(thickness=self.edge_width, color=pygfx.Color(self.edge_color))
+        )
+
+        self.world_object.add(new_line)
+
+        handlers = {
+            self._add_segment: "click",
+            self._move_segment_endpoint: "pointer_move",
+            self._finish_segment: "click",
+            self._finish_polygon: "double_click"
+        }
+
+        for handler, event in handlers.items():
+            self._plot_area.renderer.remove_event_handler(handler, event)
diff --git a/fastplotlib/layouts/_base.py b/fastplotlib/layouts/_base.py
@@ -246,7 +246,7 @@ def add_graphic(self, graphic: Graphic, center: bool = True):
         """
         self._add_or_insert_graphic(graphic=graphic, center=center, action="add")
 
-        graphic.position_z = len(self._graphics)
+        graphic.position_z = len(self)
 
     def insert_graphic(
         self,
@@ -505,10 +505,15 @@ def __getitem__(self, name: str):
         graphic_names = list()
         for g in self.graphics:
             graphic_names.append(g.name)
+
+        selector_names = list()
         for s in self.selectors:
-            graphic_names.append(s.name)
+            selector_names.append(s.name)
+
         raise IndexError(
-            f"no graphic of given name, the current graphics are:\n {graphic_names}"
+            f"No graphic or selector of given name.\n"
+            f"The current graphics are:\n {graphic_names}\n"
+            f"The current selectors are:\n {selector_names}"
         )
 
     def __str__(self):
@@ -529,3 +534,6 @@ def __repr__(self):
             f"\t{newline.join(graphic.__repr__() for graphic in self.graphics)}"
             f"\n"
         )
+
+    def __len__(self) -> int:
+        return len(self._graphics) + len(self.selectors)
diff --git a/fastplotlib/layouts/_plot.py b/fastplotlib/layouts/_plot.py
@@ -11,6 +11,7 @@
 
 from ._subplot import Subplot
 from ._record_mixin import RecordMixin
+from ..graphics.selectors import PolygonSelector
 
 
 class Plot(Subplot, RecordMixin):
@@ -211,6 +212,15 @@ def __init__(self, plot: Plot):
             layout=Layout(width="auto"),
             tooltip="flip",
         )
+
+        self.add_polygon_button = Button(
+            value=False,
+            disabled=False,
+            icon="draw-polygon",
+            layout=Layout(width="auto"),
+            tooltip="add PolygonSelector"
+        )
+
         self.record_button = ToggleButton(
             value=False,
             disabled=False,
@@ -226,6 +236,7 @@ def __init__(self, plot: Plot):
                 self.panzoom_controller_button,
                 self.maintain_aspect_button,
                 self.flip_camera_button,
+                self.add_polygon_button,
                 self.record_button,
             ]
         )
@@ -235,6 +246,7 @@ def __init__(self, plot: Plot):
         self.center_scene_button.on_click(self.center_scene)
         self.maintain_aspect_button.observe(self.maintain_aspect, "value")
         self.flip_camera_button.on_click(self.flip_camera)
+        self.add_polygon_button.on_click(self.add_polygon)
         self.record_button.observe(self.record_plot, "value")
 
     def auto_scale(self, obj):
@@ -252,6 +264,11 @@ def maintain_aspect(self, obj):
     def flip_camera(self, obj):
         self.plot.camera.world.scale_y *= -1
 
+    def add_polygon(self, obj):
+        ps = PolygonSelector(edge_width=3, edge_color="magenta")
+
+        self.plot.add_graphic(ps, center=False)
+
     def record_plot(self, obj):
         if self.record_button.value:
             try: