From e51df26a47397fb2c78c79a128d13425346dce97 Mon Sep 17 00:00:00 2001
From: Brent Yi <yibrenth@gmail.com>
Date: Wed, 18 Oct 2023 11:38:13 -0700
Subject: [PATCH] Grid support (#116)

* Add grid message

* Add plane parameterization comment
---
 examples/03_gui_callbacks.py                | 17 +++++++
 src/viser/_message_api.py                   | 40 +++++++++++++++-
 src/viser/_messages.py                      | 22 +++++++++
 src/viser/client/src/WebsocketInterface.tsx | 53 ++++++++++++++++++++-
 src/viser/client/src/WebsocketMessages.tsx  | 20 ++++++++
 5 files changed, 149 insertions(+), 3 deletions(-)

diff --git a/examples/03_gui_callbacks.py b/examples/03_gui_callbacks.py
index dba1795fb..bed333e25 100644
--- a/examples/03_gui_callbacks.py
+++ b/examples/03_gui_callbacks.py
@@ -20,6 +20,22 @@ def main() -> None:
 
     gui_reset_scene = server.add_gui_button("Reset Scene")
 
+    gui_plane = server.add_gui_dropdown(
+        "Grid plane", ("xz", "xy", "yx", "yz", "zx", "zy")
+    )
+
+    def update_plane() -> None:
+        server.add_grid(
+            "/grid",
+            width=10.0,
+            height=20.0,
+            width_segments=10,
+            height_segments=20,
+            plane=gui_plane.value,
+        )
+
+    gui_plane.on_update(lambda _: update_plane())
+
     with server.add_gui_folder("Control"):
         gui_show = server.add_gui_checkbox("Show Frame", initial_value=True)
         gui_axis = server.add_gui_dropdown("Axis", ("x", "y", "z"))
@@ -83,6 +99,7 @@ def _(_) -> None:
         draw_points()
 
     # Finally, let's add the initial frame + point cloud and just loop infinitely. :)
+    update_plane()
     draw_frame()
     draw_points()
     while True:
diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py
index 34fe8a6b9..d32312508 100644
--- a/src/viser/_message_api.py
+++ b/src/viser/_message_api.py
@@ -354,8 +354,7 @@ def add_frame(
         position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
         visible: bool = True,
     ) -> FrameHandle:
-        cast_vector(wxyz, length=4)
-        cast_vector(position, length=3)
+        """Add a coordinate frame to the scene."""
         self._queue(
             _messages.FrameMessage(
                 name=name,
@@ -366,6 +365,43 @@ def add_frame(
         )
         return FrameHandle._make(self, name, wxyz, position, visible)
 
+    def add_grid(
+        self,
+        name: str,
+        width: float = 10.0,
+        height: float = 10.0,
+        width_segments: int = 10,
+        height_segments: int = 10,
+        plane: Literal["xz", "xy", "yx", "yz", "zx", "zy"] = "xy",
+        cell_color: RgbTupleOrArray = (200, 200, 200),
+        cell_thickness: float = 1.0,
+        cell_size: float = 0.5,
+        section_color: RgbTupleOrArray = (140, 140, 140),
+        section_thickness: float = 1.0,
+        section_size: float = 1.0,
+        wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),
+        position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
+        visible: bool = True,
+    ) -> MeshHandle:
+        """Add a grid to the scene. Useful for visualizing things like ground planes."""
+        self._queue(
+            _messages.GridMessage(
+                name=name,
+                width=width,
+                height=height,
+                width_segments=width_segments,
+                height_segments=height_segments,
+                plane=plane,
+                cell_color=_encode_rgb(cell_color),
+                cell_thickness=cell_thickness,
+                cell_size=cell_size,
+                section_color=_encode_rgb(section_color),
+                section_thickness=section_thickness,
+                section_size=section_size,
+            )
+        )
+        return MeshHandle._make(self, name, wxyz, position, visible)
+
     def add_label(
         self,
         name: str,
diff --git a/src/viser/_messages.py b/src/viser/_messages.py
index 79b7c71d2..c827b4e18 100644
--- a/src/viser/_messages.py
+++ b/src/viser/_messages.py
@@ -107,6 +107,28 @@ class FrameMessage(Message):
     axes_radius: float = 0.025
 
 
+@dataclasses.dataclass
+class GridMessage(Message):
+    """Grid message. Helpful for visualizing things like ground planes."""
+
+    name: str
+
+    width: float
+    height: float
+    width_segments: int
+    height_segments: int
+
+    plane: Literal["xz", "xy", "yx", "yz", "zx", "zy"]
+
+    cell_color: int
+    cell_thickness: float
+    cell_size: float
+
+    section_color: int
+    section_thickness: float
+    section_size: float
+
+
 @dataclasses.dataclass
 class LabelMessage(Message):
     """Add a 2D label to the scene."""
diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx
index 151e7a772..652849d98 100644
--- a/src/viser/client/src/WebsocketInterface.tsx
+++ b/src/viser/client/src/WebsocketInterface.tsx
@@ -1,5 +1,5 @@
 import AwaitLock from "await-lock";
-import { CatmullRomLine, CubicBezierLine } from "@react-three/drei";
+import { CatmullRomLine, CubicBezierLine, Grid } from "@react-three/drei";
 import { unpack } from "msgpackr";
 
 import React, { useContext } from "react";
@@ -120,6 +120,57 @@ function useMessageHandler() {
         return;
       }
 
+      case "GridMessage": {
+        addSceneNodeMakeParents(
+          new SceneNode<THREE.Group>(message.name, (ref) => (
+            <group ref={ref}>
+              <Grid
+                args={[
+                  message.width,
+                  message.height,
+                  message.width_segments,
+                  message.height_segments,
+                ]}
+                side={THREE.DoubleSide}
+                cellColor={message.cell_color}
+                cellThickness={message.cell_thickness}
+                cellSize={message.cell_size}
+                sectionColor={message.section_color}
+                sectionThickness={message.section_thickness}
+                sectionSize={message.section_size}
+                rotation={
+                  // There's redundancy here when we set the side to
+                  // THREE.DoubleSide, where xy and yx should be the same.
+                  // 
+                  // But it makes sense to keep this parameterization because
+                  // specifying planes by xy seems more natural than the normal
+                  // direction (z, +z, or -z), and it opens the possibility of
+                  // rendering only FrontSide or BackSide grids in the future.
+                  //
+                  // If we add support for FrontSide or BackSide, we should
+                  // double-check that the normal directions from each of these
+                  // rotations match the right-hand rule!
+                  message.plane == "xz"
+                    ? new THREE.Euler(0.0, 0.0, 0.0)
+                    : message.plane == "xy"
+                    ? new THREE.Euler(Math.PI / 2.0, 0.0, 0.0)
+                    : message.plane == "yx"
+                    ? new THREE.Euler(0.0, Math.PI / 2.0, Math.PI / 2.0)
+                    : message.plane == "yz"
+                    ? new THREE.Euler(0.0, 0.0, Math.PI / 2.0)
+                    : message.plane == "zx"
+                    ? new THREE.Euler(0.0, Math.PI / 2.0, 0.0)
+                    : message.plane == "zy"
+                    ? new THREE.Euler(-Math.PI / 2.0, 0.0, -Math.PI / 2.0)
+                    : undefined
+                }
+              />
+            </group>
+          )),
+        );
+        return;
+      }
+
       // Add a point cloud.
       case "PointCloudMessage": {
         const geometry = new THREE.BufferGeometry();
diff --git a/src/viser/client/src/WebsocketMessages.tsx b/src/viser/client/src/WebsocketMessages.tsx
index 07e7d53de..1d43da2a5 100644
--- a/src/viser/client/src/WebsocketMessages.tsx
+++ b/src/viser/client/src/WebsocketMessages.tsx
@@ -75,6 +75,25 @@ export interface FrameMessage {
   axes_length: number;
   axes_radius: number;
 }
+/** Grid message. Helpful for visualizing things like ground planes.
+ *
+ * (automatically generated)
+ */
+export interface GridMessage {
+  type: "GridMessage";
+  name: string;
+  width: number;
+  height: number;
+  width_segments: number;
+  height_segments: number;
+  plane: "xz" | "xy" | "yx" | "yz" | "zx" | "zy";
+  cell_color: number;
+  cell_thickness: number;
+  cell_size: number;
+  section_color: number;
+  section_thickness: number;
+  section_size: number;
+}
 /** Add a 2D label to the scene.
  *
  * (automatically generated)
@@ -657,6 +676,7 @@ export type Message =
   | CameraFrustumMessage
   | GlbMessage
   | FrameMessage
+  | GridMessage
   | LabelMessage
   | Gui3DMessage
   | PointCloudMessage