Add primitive picking filters for Mesh and Markers visuals

examples/scene/face_picking.py

@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*-
+# vispy: gallery 30
+# -----------------------------------------------------------------------------
+# Copyright (c) Vispy Development Team. All Rights Reserved.
+# Distributed under the (new) BSD License. See LICENSE.txt for more info.
+# -----------------------------------------------------------------------------
+"""
+Picking Faces from a Mesh
+=========================
+
+Demonstrates how to identify (pick) individual faces on a mesh.
+
+Arguments:
+* --mesh - Path to a mesh file (OBJ/OBJ.GZ) [optional]
+
+Controls:
+* 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
+
+import numpy as np
+
+from vispy import app, scene
+from vispy.io import read_mesh, load_data_file
+from vispy.scene.visuals import Mesh
+from vispy.scene import transforms
+from vispy.visuals.filters import ShadingFilter, WireframeFilter, FacePickingFilter
+
+
+parser = argparse.ArgumentParser()
+default_mesh = load_data_file('orig/triceratops.obj.gz')
+parser.add_argument('--mesh', default=default_mesh)
+args, _ = parser.parse_known_args()
+
+vertices, faces, _normals, _texcoords = read_mesh(args.mesh)
+
+canvas = scene.SceneCanvas(keys='interactive', bgcolor='white')
+view = canvas.central_widget.add_view()
+
+view.camera = 'arcball'
+view.camera.depth_value = 1e3
+
+# Create a colored `MeshVisual`.
+face_colors = np.tile((0.5, 0.0, 0.5, 1.0), (len(faces), 1))
+mesh = Mesh(
+    vertices,
+    faces,
+    face_colors=face_colors.copy()
+)
+mesh.transform = transforms.MatrixTransform()
+mesh.transform.rotate(90, (1, 0, 0))
+mesh.transform.rotate(-45, (0, 0, 1))
+view.add(mesh)
+
+# Use filters to affect the rendering of the mesh.
+wireframe_filter = WireframeFilter()
+shading_filter = ShadingFilter()
+face_picking_filter = FacePickingFilter()
+mesh.attach(wireframe_filter)
+mesh.attach(shading_filter)
+mesh.attach(face_picking_filter)
+
+
+def attach_headlight(view):
+    light_dir = (0, 1, 0, 0)
+    shading_filter.light_dir = light_dir[:3]
+    initial_light_dir = view.camera.transform.imap(light_dir)
+
+    @view.scene.transform.changed.connect
+    def on_transform_change(event):
+        transform = view.camera.transform
+        shading_filter.light_dir = transform.map(initial_light_dir)[:3]
+
+
+attach_headlight(view)
+
+shading = itertools.cycle(("flat", "smooth", None))
+shading_filter.shading = next(shading)
+
+
+@canvas.events.mouse_move.connect
+def on_mouse_move(event):
+    restore_state = not face_picking_filter.enabled
+
+    face_picking_filter.enabled = True
+    mesh.update_gl_state(blend=False)
+    picking_render = canvas.render(bgcolor=(0, 0, 0, 0), alpha=True)
+
+    if restore_state:
+        face_picking_filter.enabled = False
+
+    mesh.update_gl_state(blend=not face_picking_filter.enabled)
+
+    ids = picking_render.view(np.uint32) - 1
+    # account for hidpi screens - canvas and render may have different size
+    cols, rows = canvas.size
+    col, row = event.pos
+    col = int(col / cols * (ids.shape[1] - 1))
+    row = int(row / rows * (ids.shape[0] - 1))
+
+    # color the hovered face on the mesh
+    if ids[row, col] in range(1, len(face_colors)):
+        # this may be less safe, but it's faster
+        mesh.mesh_data._face_colors_indexed_by_faces[ids[row, col]] = (0, 1, 0, 1)
+        mesh.mesh_data_changed()
+
+
+@canvas.events.key_press.connect
+def on_key_press(event):
+    if event.key == 'c':
+        mesh.set_data(vertices, faces, face_colors=face_colors)
+    if event.key == 's':
+        shading_filter.shading = next(shading)
+        mesh.update()
+    elif 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()
+
+
+if __name__ == "__main__":
+    print(__doc__)
+    app.run()

vispy/visuals/filters/__init__.py

@@ -6,4 +6,4 @@
 from .clipper import Clipper  # noqa
 from .color import Alpha, ColorFilter, IsolineFilter, ZColormapFilter  # noqa
 from .picking import PickingFilter  # noqa
-from .mesh import TextureFilter, ShadingFilter, InstancedShadingFilter, WireframeFilter  # noqa
+from .mesh import TextureFilter, ShadingFilter, InstancedShadingFilter, WireframeFilter, FacePickingFilter  # noqa

vispy/visuals/filters/mesh.py

@@ -766,3 +766,85 @@ def _attach(self, visual):
     def _detach(self, visual):
         visual.events.data_updated.disconnect(self.on_mesh_data_updated)
         super()._detach(visual)
+
+
+class FacePickingFilter(Filter):
+    """Filter used to color mesh faces by a picking ID.
+
+    Note that the ID color uses the alpha channel, so this may not be used
+    with blending enabled.
+
+    Examples
+    --------
+    See
+    `examples/scene/face_picking.py
+    <https://github.com/vispy/vispy/blob/main/examples/scene/face_picking.py>`_
+    example script.
+    """
+
+    def __init__(self):
+        vfunc = Function("""\
+            varying vec4 v_face_picking_color;
+            void prepare_face_picking() {
+                v_face_picking_color = $ids;
+            }
+        """)
+        ffunc = Function("""\
+            varying vec4 v_face_picking_color;
+            void face_picking_filter() {
+                if ($enabled != 1) {
+                    return;
+                }
+                gl_FragColor = v_face_picking_color;
+            }
+        """)
+
+        self._ids = VertexBuffer(np.zeros((0, 4), dtype=np.float32))
+        vfunc['ids'] = self._ids
+        super().__init__(vcode=vfunc, fcode=ffunc)
+        self._n_faces = 0
+        self.enabled = False
+
+    @property
+    def enabled(self):
+        return self._enabled
+
+    @enabled.setter
+    def enabled(self, e):
+        self._enabled = bool(e)
+        self.fshader['enabled'] = int(self._enabled)
+        self._update_data()
+
+    def _update_data(self):
+        if not self.attached:
+            return
+        if self._visual.mesh_data.is_empty():
+            n_faces = 0
+        else:
+            n_faces = len(self._visual.mesh_data.get_faces())
+
+        # we only care about the number of faces changing
+        if self._n_faces == n_faces:
+            return
+        self._n_faces = n_faces
+
+        # pack the face ids into a color buffer
+        # TODO: consider doing the bit-packing in the shader
+        ids = np.arange(
+            1, n_faces + 1,
+            dtype=np.uint32
+        ).view(np.uint8).reshape(n_faces, 4)
+        ids = np.divide(ids, 255, dtype=np.float32)
+        self._ids.set_data(np.repeat(ids, 3, axis=0))
+
+    def on_mesh_data_updated(self, event):
+        self._update_data()
+
+    def _attach(self, visual):
+        super()._attach(visual)
+        visual.events.data_updated.connect(self.on_mesh_data_updated)
+        self._update_data()
+
+    def _detach(self, visual):
+        visual.events.data_updated.disconnect(self.on_mesh_data_updated)
+        super()._detach(visual)

vispy/visuals/filters/tests/test_face_picking_filter.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*- 
+# Copyright (c) Vispy Development Team. All Rights Reserved.
+# Distributed under the (new) BSD License. See LICENSE.txt for more info.
+import numpy as np
+
+from vispy.geometry import create_plane
+from vispy.scene.visuals import Mesh
+from vispy.testing import TestingCanvas
+from vispy.visuals.filters import FacePickingFilter
+
+
+def test_empty_mesh_face_picking():
+    mesh = Mesh()
+    filter = FacePickingFilter()
+    mesh.attach(filter)
+    filter.enabled = True
+
+
+def test_face_picking():
+    vertices, faces, _ = create_plane(125, 125)
+    vertices = vertices["position"]
+    vertices[:, :2] += 125 / 2
+    mesh = Mesh(vertices=vertices, faces=faces)
+    filter = FacePickingFilter()
+    mesh.attach(filter)
+
+    with TestingCanvas(size=(125, 125)) as c:
+        view = c.central_widget.add_view()
+        view.add(mesh)
+        filter.enabled = True
+        mesh.update_gl_state(blend=False)
+        picking_render = c.render(bgcolor=(0, 0, 0, 0), alpha=True)
+
+    # the plane is made up of two triangles and nearly fills the view
+    # pick one point on each triangle
+    assert np.allclose(
+        picking_render[125 // 2, 125 // 4],
+        (1, 0, 0, 0),
+    )
+    assert np.allclose(
+        picking_render[125 // 2, 3 * 125 // 4],
+        (2, 0, 0, 0),
+    )