[FEATURE] Add example viewer plugins. (#2357)

Co-authored-by: Alexis DUBURCQ <alexis.duburcq@gmail.com>

Author: Milotrince

examples/viewer_plugin/mesh_point_selector.py

@@ -0,0 +1,289 @@
+import csv
+import os
+from typing import TYPE_CHECKING, NamedTuple
+
+import numpy as np
+from typing_extensions import override
+
+import genesis as gs
+import genesis.utils.geom as gu
+import genesis.vis.keybindings as kb
+from genesis.utils.misc import tensor_to_array
+from genesis.vis.viewer_plugins import EVENT_HANDLE_STATE, EVENT_HANDLED, RaycasterViewerPlugin
+
+if TYPE_CHECKING:
+    from genesis.engine.entities.rigid_entity import RigidLink
+    from genesis.engine.scene import Scene
+    from genesis.ext.pyrender.node import Node
+
+
+class SelectedPoint(NamedTuple):
+    """
+    Represents a selected point on a rigid mesh surface.
+
+    Attributes
+    ----------
+    link : RigidLink
+        The rigid link that the point belongs to.
+    local_position : np.ndarray, shape (3,)
+        The position of the point in the link's local coordinate frame.
+    local_normal : np.ndarray, shape (3,)
+        The surface normal at the point in the link's local coordinate frame.
+    """
+
+    link: "RigidLink"
+    local_position: np.ndarray  # shape (3,)
+    local_normal: np.ndarray  # shape (3,)
+
+
+class MeshPointSelectorPlugin(RaycasterViewerPlugin):
+    """
+    Interactive viewer plugin that enables using mouse clicks to select points on rigid meshes.
+    Selected points are stored in local coordinates relative to their link's frame.
+    """
+
+    def __init__(
+        self,
+        sphere_radius: float = 0.005,
+        sphere_color: tuple = (0.1, 0.3, 1.0, 1.0),
+        hover_color: tuple = (0.3, 0.5, 1.0, 1.0),
+        grid_snap: tuple[float, float, float] = (-1.0, -1.0, -1.0),
+        output_file: str = "selected_points.csv",
+    ) -> None:
+        super().__init__()
+        self.sphere_radius = sphere_radius
+        self.sphere_color = sphere_color
+        self.hover_color = hover_color
+        self.grid_snap = grid_snap
+        self.output_file = output_file
+
+        self.selected_points: dict[int, SelectedPoint] = {}
+        self._prev_mouse_pos: tuple[int, int] = (0, 0)
+
+    def build(self, viewer, camera: "Node", scene: "Scene"):
+        super().build(viewer, camera, scene)
+        self._prev_mouse_pos: tuple[int, int] = (self.viewer._viewport_size[0] // 2, self.viewer._viewport_size[1] // 2)
+
+    def _get_pos_hash(self, pos: np.ndarray) -> int:
+        """
+        Generate a hash for a given position to use as a unique identifier.
+
+        Parameters
+        ----------
+        pos : np.ndarray, shape (3,)
+            The position to hash.
+
+        Returns
+        -------
+        int
+            The hash of the position.
+        """
+        return hash((round(pos[0], 6), round(pos[1], 6), round(pos[2], 6)))
+
+    def _snap_to_grid(self, point: np.ndarray) -> np.ndarray:
+        """
+        Snap a point to the grid based on grid_snap settings.
+
+        Parameters
+        ----------
+        point : np.ndarray, shape (3,)
+            The point to snap.
+
+        Returns
+        -------
+        np.ndarray, shape (3,)
+            The point snapped to the grid.
+        """
+        grid_snap = np.array(self.grid_snap)
+        # Snap each axis if the snap value is non-negative
+        return np.where(grid_snap >= 0, np.round(point / grid_snap) * grid_snap, point)
+
+    @override
+    def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STATE:
+        self._prev_mouse_pos = (x, y)
+
+    @override
+    def on_mouse_press(self, x: int, y: int, button: int, modifiers: int) -> EVENT_HANDLE_STATE:
+        if button == 1:  # left click
+            ray = self._screen_position_to_ray(x, y)
+            ray_hit = self._raycaster.cast(*ray)
+
+            if ray_hit is not None and ray_hit.geom:
+                link = ray_hit.geom.link
+                world_pos = ray_hit.position
+                world_normal = ray_hit.normal
+
+                # Get link pose
+                link_pos = tensor_to_array(link.get_pos())
+                link_quat = tensor_to_array(link.get_quat())
+
+                local_pos = gu.inv_transform_by_trans_quat(world_pos, link_pos, link_quat)
+                local_normal = gu.inv_transform_by_quat(world_normal, link_quat)
+
+                # Apply grid snapping to local position
+                local_pos = self._snap_to_grid(local_pos)
+
+                pos_hash = self._get_pos_hash(local_pos)
+                if pos_hash in self.selected_points:
+                    # Deselect point if already selected
+                    del self.selected_points[pos_hash]
+                else:
+                    selected_point = SelectedPoint(link, local_pos, local_normal)
+                    self.selected_points[pos_hash] = selected_point
+
+                return EVENT_HANDLED
+        return None
+
+    @override
+    def on_draw(self) -> None:
+        super().on_draw()
+        if self.scene._visualizer is not None and self.scene._visualizer.is_built:
+            self.scene.clear_debug_objects()
+            mouse_ray = self._screen_position_to_ray(*self._prev_mouse_pos)
+
+            closest_hit = self._raycaster.cast(*mouse_ray)
+            if closest_hit is not None:
+                snap_pos = self._snap_to_grid(closest_hit.position)
+
+                # Draw hover preview
+                self.scene.draw_debug_sphere(
+                    snap_pos,
+                    self.sphere_radius,
+                    self.hover_color,
+                )
+                self.scene.draw_debug_arrow(
+                    snap_pos,
+                    tuple(n * 0.05 for n in closest_hit.normal),
+                    self.sphere_radius / 2,
+                    self.hover_color,
+                )
+
+            if self.selected_points:
+                world_positions = []
+                for point in self.selected_points.values():
+                    link_pos = tensor_to_array(point.link.get_pos())
+                    link_quat = tensor_to_array(point.link.get_quat())
+                    local_pos_arr = np.array(point.local_position, dtype=np.float32)
+                    current_world_pos = gu.transform_by_trans_quat(local_pos_arr, link_pos, link_quat)
+                    world_positions.append(current_world_pos)
+
+                if len(world_positions) == 1:
+                    self.scene.draw_debug_sphere(
+                        world_positions[0],
+                        self.sphere_radius,
+                        self.sphere_color,
+                    )
+                else:
+                    positions_array = np.array(world_positions)
+                    self.scene.draw_debug_spheres(positions_array, self.sphere_radius, self.sphere_color)
+
+    @override
+    def on_close(self) -> None:
+        super().on_close()
+
+        if not self.selected_points:
+            print("[MeshPointSelectorPlugin] No points selected.")
+            return
+
+        output_file = self.output_file
+        try:
+            with open(output_file, "w", newline="") as csvfile:
+                writer = csv.writer(csvfile)
+
+                writer.writerow(
+                    [
+                        "point_idx",
+                        "link_idx",
+                        "local_pos_x",
+                        "local_pos_y",
+                        "local_pos_z",
+                        "local_normal_x",
+                        "local_normal_y",
+                        "local_normal_z",
+                    ]
+                )
+
+                for i, point in enumerate(self.selected_points.values(), 1):
+                    writer.writerow(
+                        [
+                            i,
+                            point.link.idx,
+                            point.local_position[0],
+                            point.local_position[1],
+                            point.local_position[2],
+                            point.local_normal[0],
+                            point.local_normal[1],
+                            point.local_normal[2],
+                        ]
+                    )
+
+            gs.logger.info(
+                f"[MeshPointSelectorPlugin] Wrote {len(self.selected_points)} selected points to '{output_file}'"
+            )
+
+        except Exception as e:
+            gs.logger.error(f"[MeshPointSelectorPlugin] Error writing to '{output_file}': {e}")
+
+
+if __name__ == "__main__":
+    gs.init(backend=gs.gpu)
+
+    scene = gs.Scene(
+        sim_options=gs.options.SimOptions(
+            gravity=(0.0, 0.0, 0.0),
+        ),
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(0.6, 0.6, 0.6),
+            camera_lookat=(0.0, 0.0, 0.2),
+            camera_fov=40,
+        ),
+        vis_options=gs.options.VisOptions(
+            show_world_frame=True,
+        ),
+        profiling_options=gs.options.ProfilingOptions(
+            show_FPS=False,
+        ),
+        show_viewer=True,
+    )
+
+    hand = scene.add_entity(
+        morph=gs.morphs.URDF(
+            file="urdf/shadow_hand/shadow_hand.urdf",
+            collision=True,
+            pos=(0.0, 0.0, 0.0),
+            euler=(0.0, 0.0, 180.0),
+            fixed=True,
+            merge_fixed_links=False,
+        ),
+    )
+
+    scene.viewer.add_plugin(
+        MeshPointSelectorPlugin(
+            sphere_radius=0.004,
+            grid_snap=(-1.0, 0.01, 0.01),
+            output_file="selected_points.csv",
+        )
+    )
+
+    scene.build()
+
+    is_running = True
+
+    def stop():
+        global is_running
+        is_running = False
+
+    scene.viewer.register_keybinds(
+        kb.Keybind("quit", kb.Key.ESCAPE, kb.KeyAction.PRESS, callback=stop),
+    )
+
+    try:
+        while is_running:
+            scene.step()
+
+            if "PYTEST_VERSION" in os.environ:
+                break
+    except KeyboardInterrupt:
+        gs.logger.info("Simulation interrupted, exiting.")
+    finally:
+        gs.logger.info("Simulation finished.")

examples/viewer_plugin/mouse_interaction.py

@@ -0,0 +1,73 @@
+import argparse
+import math
+import os
+
+import genesis as gs
+import genesis.vis.keybindings as kb
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Mouse interaction viewer plugin example.")
+    parser.add_argument("--use_force", action="store_true", help="Apply spring forces instead of setting position")
+    args = parser.parse_args()
+
+    gs.init(backend=gs.gpu)
+
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(3.5, 0.0, 2.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+            camera_fov=40,
+        ),
+        profiling_options=gs.options.ProfilingOptions(
+            show_FPS=False,
+        ),
+        show_viewer=True,
+    )
+
+    scene.add_entity(gs.morphs.Plane())
+
+    sphere = scene.add_entity(
+        morph=gs.morphs.Sphere(
+            pos=(-0.3, -0.3, 0),
+            radius=0.1,
+        ),
+    )
+    for i in range(6):
+        angle = i * (2 * math.pi / 6)
+        radius = 0.5 + i * 0.1
+        cube = scene.add_entity(
+            morph=gs.morphs.Box(
+                pos=(radius * math.cos(angle), radius * math.sin(angle), 0.1 + i * 0.1),
+                size=(0.2, 0.2, 0.2),
+            ),
+        )
+
+    scene.viewer.add_plugin(
+        gs.vis.viewer_plugins.MouseInteractionPlugin(
+            use_force=args.use_force,
+            color=(0.1, 0.6, 0.8, 0.6),
+        )
+    )
+
+    scene.build()
+
+    is_running = True
+
+    def stop():
+        global is_running
+        is_running = False
+
+    scene.viewer.register_keybinds(
+        kb.Keybind("quit", kb.Key.ESCAPE, kb.KeyAction.PRESS, callback=stop),
+    )
+
+    try:
+        while is_running:
+            scene.step()
+
+            if "PYTEST_VERSION" in os.environ:
+                break
+    except KeyboardInterrupt:
+        gs.logger.info("Simulation interrupted, exiting.")
+    finally:
+        gs.logger.info("Simulation finished.")