Refactor primitive picking filter base class

examples/scene/face_picking.py

@@ -14,10 +14,10 @@
 * --mesh - Path to a mesh file (OBJ/OBJ.GZ) [optional]
 
 Controls:
+* p  - Toggle face picking view - shows the colors encoding face ID
+* r  - Clear painted faces
 * s  - Cycle shading modes (None, 'flat', 'smooth')
 * w  - Toggle wireframe
-* p  - Toggle face picking view - shows the colors encoding face ID
-* c  - Clear painted faces
 """
 import argparse
 import itertools
@@ -114,7 +114,7 @@ def on_mouse_move(event):
     mesh.update_gl_state(blend=not face_picking_filter.enabled)
 
     # unpack the face index from the color in the center pixel
-    face_idx = (picking_render.view(np.uint32) - 1)[1, 1]
+    face_idx = (picking_render.view(np.uint32) - 1)[1, 1, 0]
 
     if face_idx > 0 and face_idx < len(face_colors):
         # this may be less safe, but it's faster than set_data
@@ -124,19 +124,21 @@ def on_mouse_move(event):
 
 @canvas.events.key_press.connect
 def on_key_press(event):
-    if event.key == 'c':
+    if event.key == 'p':
+        face_picking_filter.enabled = not face_picking_filter.enabled
+        mesh.update_gl_state(blend=not face_picking_filter.enabled)
+        mesh.update()
+
+    if event.key == 'r':
         mesh.set_data(vertices, faces, face_colors=face_colors)
+
     if event.key == 's':
         shading_filter.shading = next(shading)
         mesh.update()
-    elif event.key == 'w':
+
+    if event.key == 'w':
         wireframe_filter.enabled = not wireframe_filter.enabled
         mesh.update()
-    elif event.key == 'p':
-        # toggle face picking view
-        face_picking_filter.enabled = not face_picking_filter.enabled
-        mesh.update_gl_state(blend=not face_picking_filter.enabled)
-        mesh.update()
 
 
 canvas.show()

examples/scene/marker_picking.py

@@ -8,46 +8,62 @@
 Picking Markers
 ===============
 
-Demonstrates how to identify (pick) markers. Click markers to change their edge
-color.
+Demonstrates how to identify (pick) markers. Hover markers to change their
+symbol and color.
 
 Controls:
 * p  - Toggle picking view - shows the colors encoding marker ID
-* s  - Cycle marker symbols
-* c  - Clear picked markers
+* r  - Reset marker symbols and colors
 """
-import itertools
-
+import random
+import time
 import numpy as np
 from scipy.constants import golden as GOLDEN
 
 from vispy import app, scene
 from vispy.scene.visuals import Markers
 from vispy.visuals.filters import MarkerPickingFilter
 
-symbols = itertools.cycle(Markers._symbol_shader_values.keys())
-
-
-MARKER_SIZE = 0.0125
-EDGE_WDITH = MARKER_SIZE / 10
-
 canvas = scene.SceneCanvas(keys='interactive', bgcolor='black')
 view = canvas.central_widget.add_view(camera="panzoom")
 view.camera.rect = (-1, -1, 2, 2)
+
+# floret pattern
 n = 10_000
-radius = np.linspace(0, 0.9, n)**0.6
-theta = np.arange(n) * GOLDEN * np.pi
+radius = np.linspace(0, 0.9, n)**0.6  # prevent extreme density at center
+theta = np.arange(n) * GOLDEN
 pos = np.column_stack([radius * np.cos(theta), radius * np.sin(theta)])
-edge_color = np.ones((len(pos), 4), dtype=np.float32)
+
+COLORS = [
+    (1, 0, 0, 1),  # red
+    (1, 0.5, 0, 1),  # orange
+    (1, 1, 0, 1),  # yellow
+    (0, 1, 0, 1),  # green
+    (0, 0, 1, 1),  # blue
+    (0.29, 0, 0.51, 1),  # indigo
+    (0.93, 0.51, 0.93, 1),  # violet
+]
+
+colors = np.zeros((n, 4), dtype=np.float32)
+colors[:, 0] = 1  # red
+colors[:, -1] = 1  # alpha
+_colors = colors.copy()
+
+symbols = list(Markers._symbol_shader_values.keys())
+symbols_ring = dict(zip(symbols, symbols[1:]))
+symbols_ring[symbols[-1]] = symbols[0]
+
+EDGE_COLOR = "white"
+MARKER_SIZE = 0.0125
+EDGE_WDITH = MARKER_SIZE / 10
 
 markers = Markers(
     pos=pos,
-    edge_color=edge_color,
-    face_color="red",
+    edge_color=EDGE_COLOR,
+    face_color=colors,
     size=MARKER_SIZE,
     edge_width=EDGE_WDITH,
     scaling="scene",
-    symbol=next(symbols),
 )
 markers.update_gl_state(depth_test=True)
 view.add(markers)
@@ -63,8 +79,17 @@ def on_viewbox_change(event):
     markers.update_gl_state(blend=not picking_filter.enabled)
 
 
-@canvas.events.mouse_press.connect
-def on_mouse_press(event):
+throttle = time.monotonic()
+
+
+@canvas.events.mouse_move.connect
+def on_mouse_move(event):
+    global throttle
+    # throttle mouse events to 50ms
+    if time.monotonic() - throttle < 0.05:
+        return
+    throttle = time.monotonic()
+
     # adjust the event position for hidpi screens
     render_size = tuple(d * canvas.pixel_scale for d in canvas.size)
     x_pos = event.pos[0] * canvas.pixel_scale
@@ -84,40 +109,42 @@ def on_mouse_press(event):
     markers.update_gl_state(blend=not picking_filter.enabled)
 
     # unpack the face index from the color in the center pixel
-    face_idx = (picking_render.view(np.uint32) - 1)[2, 2]
+    marker_idx = (picking_render.view(np.uint32) - 1)[2, 2, 0]
 
-    if face_idx >= 0 and face_idx < len(pos):
-        edge_color[face_idx] = (0, 1, 0, 1)
+    if marker_idx >= 0 and marker_idx < len(pos):
+        new_symbols = list(markers.symbol)
+        new_symbol = symbols_ring[new_symbols[marker_idx]]
+        new_symbols[marker_idx] = new_symbol
+
+        colors[marker_idx] = random.choice(COLORS)
+        colors[marker_idx, -1] = 1  # alpha
         markers.set_data(
             pos=pos,
-            edge_color=edge_color,
-            face_color="red",
+            edge_color=EDGE_COLOR,
+            face_color=colors,
             size=MARKER_SIZE,
             edge_width=EDGE_WDITH,
-            symbol=markers.symbol,
+            symbol=new_symbols,
         )
 
 
 @canvas.events.key_press.connect
 def on_key_press(event):
-    if event.key == 'c':
-        edge_color = np.ones((len(pos), 4), dtype=np.float32)
-        markers.set_data(
-            pos=pos,
-            edge_color=edge_color,
-            face_color="red",
-            size=MARKER_SIZE,
-            edge_width=EDGE_WDITH,
-            symbol=markers.symbol,
-        )
     if event.key == 'p':
         # toggle face picking view
         picking_filter.enabled = not picking_filter.enabled
         markers.update_gl_state(blend=not picking_filter.enabled)
         markers.update()
-    if event.key == 's':
-        markers.symbol = next(symbols)
-        markers.update()
+    if event.key == 'r':
+        # reset marker symbols
+        colors = _colors
+        markers.set_data(
+            pos=pos,
+            edge_color=EDGE_COLOR,
+            face_color=colors,
+            size=MARKER_SIZE,
+            edge_width=EDGE_WDITH,
+        )
 
 
 canvas.show()

vispy/visuals/filters/__init__.py

@@ -2,7 +2,7 @@
 # Copyright (c) Vispy Development Team. All Rights Reserved.
 # Distributed under the (new) BSD License. See LICENSE.txt for more info.
 
-from .base_filter import Filter  # noqa
+from .base_filter import Filter, PrimitivePickingFilter  # noqa
 from .clipper import Clipper  # noqa
 from .color import Alpha, ColorFilter, IsolineFilter, ZColormapFilter  # noqa
 from .picking import PickingFilter  # noqa

vispy/visuals/filters/base_filter.py

@@ -2,6 +2,11 @@
 # Copyright (c) Vispy Development Team. All Rights Reserved.
 # Distributed under the (new) BSD License. See LICENSE.txt for more info.
 
+from abc import ABCMeta, abstractmethod
+
+import numpy as np
+
+from vispy.gloo import VertexBuffer
 from ..shaders import Function
 
 
@@ -120,3 +125,98 @@ def _detach(self, visual):
 
         self._attached = False
         self._visual = None
+
+
+class PrimitivePickingFilter(Filter, metaclass=ABCMeta):
+    """Abstract base class for Visual-specific filters to implement a
+    primitive-picking mode.
+
+    Subclasses must (and usually only need to) implement
+    :py:meth:`_update_id_colors`.
+    """
+
+    def __init__(self):
+        vfunc = Function("""\
+            varying vec4 v_marker_picking_color;
+            void prepare_marker_picking() {
+                v_marker_picking_color = $ids;
+            }
+        """)
+        ffunc = Function("""\
+            varying vec4 v_marker_picking_color;
+            void marker_picking_filter() {
+                if ($enabled != 1) {
+                    return;
+                }
+                gl_FragColor = v_marker_picking_color;
+            }
+        """)
+
+        self._id_colors = VertexBuffer(np.zeros((0, 4), dtype=np.float32))
+        vfunc['ids'] = self._id_colors
+        self._n_primitives = 0
+        super().__init__(vcode=vfunc, fcode=ffunc)
+        self.enabled = False
+
+    @abstractmethod
+    def _update_id_colors(self):
+        """Calculate the colors encoding the picking IDs for the visual.
+
+        Generally, this method should be implemented to:
+            1. Calculate the number of primitives in the visual, stored in
+            `self._n_primitives`.
+            2. Pack the picking IDs into a 4-component float array (RGBA
+            VertexBuffer), stored in `self._id_colors`.
+
+        As an example of packing IDs into a VertexBuffer, consider the following
+        implementation for the Mesh visual:
+        ```
+            # calculate the number of primitives
+            n_faces = len(self._visual.mesh_data.get_faces())
+            self._n_primitives = n_faces
+
+            # assign 32 bit IDs to each primitive
+            # starting from 1 reserves (0, 0, 0, 0) for the background
+            ids = np.arange(1, n_faces + 1,dtype=np.uint32)
+
+            # reinterpret as 8-bit RGBA and normalize colors into floats
+            id_colors = np.divide(
+                ids.view(np.uint8).reshape(n_faces, 4),
+                255,
+                dtype=np.float32
+            )
+
+            # store the colors in a VertexBuffer, repeating each color 3 times
+            # for each vertex in each triangle
+            self._id_colors.set_data(np.repeat(idid_colors, 3, axis=0))
+        ```
+
+        For performance, you may want to optimize this method to only update
+        the IDs when the data meaningfully changes - for example when the
+        number of primitives changes.
+        """
+        raise NotImplementedError(self)
+
+    def _on_data_updated(self, event=None):
+        if not self.attached:
+            return
+        self._update_id_colors()
+
+    @property
+    def enabled(self):
+        return self._enabled
+
+    @enabled.setter
+    def enabled(self, e):
+        self._enabled = bool(e)
+        self.fshader['enabled'] = int(self._enabled)
+        self._on_data_updated()
+
+    def _attach(self, visual):
+        super()._attach(visual)
+        visual.events.data_updated.connect(self._on_data_updated)
+        self._on_data_updated()
+
+    def _detach(self, visual):
+        visual.events.data_updated.disconnect(self._on_data_updated)
+        super()._detach(visual)