add grid snap

Author: Milotrince

genesis/ext/pyrender/interaction/plugins/mesh_selector.py

@@ -58,6 +58,29 @@ def __init__(
 
         self.raycaster: ViewerRaycaster = ViewerRaycaster(self.scene)
 
+    def _snap_to_grid(self, position: Vec3) -> Vec3:
+        """
+        Snap a position to the grid based on grid_snap settings.
+        
+        Parameters
+        ----------
+        position : Vec3
+            The position to snap.
+            
+        Returns
+        -------
+        Vec3
+            The snapped position.
+        """
+        snap_x, snap_y, snap_z = self.options.grid_snap
+        
+        # Snap each axis if the snap value is non-negative
+        x = round(position.x / snap_x) * snap_x if snap_x >= 0 else position.x
+        y = round(position.y / snap_y) * snap_y if snap_y >= 0 else position.y
+        z = round(position.z / snap_z) * snap_z if snap_z >= 0 else position.z
+            
+        return Vec3.from_xyz(x, y, z)
+
     @override
     def on_mouse_motion(self, x: int, y: int, dx: int, dy: int) -> EVENT_HANDLE_STATE:
         super().on_mouse_motion(x, y, dx, dy)
@@ -80,6 +103,9 @@ def on_mouse_press(self, x: int, y: int, button: int, modifiers: int) -> EVENT_H
                 local_pos = pose.inverse_transform_point(world_pos)
                 local_normal = pose.inverse_transform_direction(world_normal)
 
+                # Apply grid snapping to local position
+                local_pos = self._snap_to_grid(local_pos)
+
                 selected_point = SelectedPoint(
                     link=link,
                     local_position=local_pos,
@@ -103,14 +129,15 @@ def on_draw(self) -> None:
 
             closest_hit = self.raycaster.cast_ray(mouse_ray.origin.v, mouse_ray.direction.v)
             if closest_hit.is_hit:
+                snap_pos = self._snap_to_grid(closest_hit.position)
                 # Draw hover preview
                 self.scene.draw_debug_sphere(
-                    closest_hit.position.v,
+                    snap_pos.v,
                     self.options.sphere_radius,
                     self.options.hover_color,
                 )
                 self.scene.draw_debug_arrow(
-                    closest_hit.position.v,
+                    snap_pos.v,
                     closest_hit.normal.v * 0.1,
                     self.options.sphere_radius / 2,
                     self.options.hover_color,

genesis/options/viewer_plugins.py

@@ -31,9 +31,13 @@ class MeshPointSelectorPlugin(ViewerPlugin):
         The color of the sphere used to visualize selected points.
     hover_color : tuple
         The color of the sphere used to visualize the point and normal when hovering over a mesh.
+    grid_snap : tuple[float, float, float]
+        Grid snap spacing for each axis (x, y, z). Any negative value disables snapping for that axis.
+        Default is (-1.0, -1.0, -1.0) which means no snapping.
     """
 
     sphere_radius: float = 0.005
-    sphere_color: tuple = (1.0, 0.0, 0.0, 1.0)
-    hover_color: tuple = (1.0, 0.0, 0.0, 1.0)
+    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"

genesis/options/vis.py

@@ -3,6 +3,7 @@
 import genesis as gs
 
 from .options import Options
+from .viewer_plugins import ViewerPlugin
 
 
 class ViewerOptions(Options):
@@ -33,17 +34,19 @@ class ViewerOptions(Options):
         The up vector of the camera's extrinsic pose.
     camera_fov : float
         The field of view (in degrees) of the camera.
+    viewer_plugin : ViewerPluginOptions | None
+        Optional viewer plugin that adds interactive functionality to the viewer.
     """
 
-    res: Optional[tuple] = None
-    run_in_thread: Optional[bool] = None
+    res: tuple | None = None
+    run_in_thread: bool | None = None
     refresh_rate: int = 60
-    max_FPS: Optional[int] = 60
+    max_FPS: int | None = 60
     camera_pos: tuple = (3.5, 0.5, 2.5)
     camera_lookat: tuple = (0.0, 0.0, 0.5)
     camera_up: tuple = (0.0, 0.0, 1.0)
     camera_fov: float = 40
-    enable_interaction: bool = False
+    viewer_plugin: ViewerPlugin | None = None
 
 
 class VisOptions(Options):
@@ -139,7 +142,7 @@ def __init__(self, **data):
                 f"Unsupported `render_particle_as`: {self.render_particle_as}, must be one of ['sphere', 'tet']"
             )
 
-        if not self.n_rendered_envs is None:
+        if self.n_rendered_envs is not None:
             gs.logger.warning(
                 "Viewer option 'n_rendered_envs' is deprecated and will be removed in future release. Please use "
                 "'rendered_envs_idx' instead."