Added box and icosphere primitives (nerfstudio-project#132)

* Added box and icosphere primitives

* Nits

---------

Co-authored-by: Jonathan <[email protected]>
Co-authored-by: Brent Yi <[email protected]>

examples/12_click_meshes.py

@@ -6,8 +6,6 @@
 import time
 
 import matplotlib
-import numpy as onp
-import trimesh.creation
 
 import viser
 
@@ -43,33 +41,33 @@ def add_swappable_mesh(i: int, j: int) -> None:
 
         def create_mesh(counter: int) -> None:
             if counter == 0:
-                mesh = trimesh.creation.box((0.5, 0.5, 0.5))
-            elif counter == 1:
-                mesh = trimesh.creation.box((0.5, 0.5, 0.5))
+                color = (0.8, 0.8, 0.8)
             else:
-                mesh = trimesh.creation.icosphere(subdivisions=2, radius=0.4)
-
-            colors = colormap(
-                (i * grid_shape[1] + j + onp.random.rand(mesh.vertices.shape[0]))
-                / (grid_shape[0] * grid_shape[1])
-            )
-            if counter != 0:
-                assert mesh.visual is not None
-                mesh.visual.vertex_colors = colors
-
-            handle = server.add_mesh_trimesh(
-                name=f"/sphere_{i}_{j}",
-                mesh=mesh,
-                position=(i, j, 0.0),
-            )
+                index = (i * grid_shape[1] + j) / (grid_shape[0] * grid_shape[1])
+                color = colormap(index)[:3]
+
+            if counter in (0, 1):
+                handle = server.add_box(
+                    name=f"/sphere_{i}_{j}",
+                    position=(i, j, 0.0),
+                    color=color,
+                    dimensions=(0.5, 0.5, 0.5),
+                )
+            else:
+                handle = server.add_icosphere(
+                    name=f"/sphere_{i}_{j}",
+                    radius=0.4,
+                    color=color,
+                    position=(i, j, 0.0),
+                )
 
             @handle.on_click
             def _(_) -> None:
                 x_value.value = i
                 y_value.value = j
 
-                # The new mesh will replace the old one because the names (/sphere_{i}_{j}) are
-                # the same.
+                # The new mesh will replace the old one because the names
+                # /sphere_{i}_{j} are the same.
                 create_mesh((counter + 1) % 3)
 
         create_mesh(0)

src/viser/_message_api.py

@@ -33,6 +33,7 @@
 import numpy as onp
 import numpy.typing as onpt
 import trimesh
+import trimesh.creation
 import trimesh.exchange
 import trimesh.visual
 from typing_extensions import Literal, ParamSpec, TypeAlias, assert_never
@@ -632,6 +633,7 @@ def add_mesh_simple(
         wireframe: bool = False,
         opacity: Optional[float] = None,
         material: Literal["standard", "toon3", "toon5"] = "standard",
+        flat_shading: bool = True,
         side: Literal["front", "back", "double"] = "front",
         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),
@@ -649,6 +651,8 @@ def add_mesh_simple(
             wireframe: Boolean indicating if the mesh should be rendered as a wireframe.
             opacity: Opacity of the mesh. None means opaque.
             material: Material type of the mesh ('standard', 'toon3', 'toon5').
+            flat_shading: Whether to do flat shading. Set to False to apply smooth
+                shading.
             side: Side of the surface to render ('front', 'back', 'double').
             wxyz: Quaternion rotation to parent frame from local frame (R_pl).
             position: Translation from parent frame to local frame (t_pl).
@@ -668,6 +672,7 @@ def add_mesh_simple(
                 vertex_colors=None,
                 wireframe=wireframe,
                 opacity=opacity,
+                flat_shading=flat_shading,
                 side=side,
                 material=material,
             )
@@ -711,6 +716,82 @@ def add_mesh_trimesh(
                 visible=visible,
             )
 
+    def add_box(
+        self,
+        name: str,
+        color: RgbTupleOrArray,
+        dimensions: Tuple[float, float, float] | onp.ndarray = (1.0, 1.0, 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 box to the scene.
+
+        Args:
+            name: A scene tree name. Names in the format of /parent/child can be used to
+                define a kinematic tree.
+            color: Color of the box as an RGB tuple.
+            dimensions: Dimensions of the box (x, y, z).
+            wxyz: Quaternion rotation to parent frame from local frame (R_pl).
+            position: Translation from parent frame to local frame (t_pl).
+            visible: Whether or not this box is initially visible.
+
+        Returns:
+            Handle for manipulating scene node.
+        """
+        mesh = trimesh.creation.box(dimensions)
+
+        return self.add_mesh_simple(
+            name=name,
+            vertices=mesh.vertices,
+            faces=mesh.faces,
+            color=color,
+            flat_shading=True,
+            position=position,
+            wxyz=wxyz,
+            visible=visible,
+        )
+
+    def add_icosphere(
+        self,
+        name: str,
+        radius: float,
+        color: RgbTupleOrArray,
+        subdivisions: int = 3,
+        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 an icosphere to the scene.
+
+        Args:
+            name: A scene tree name. Names in the format of /parent/child can be used to
+                define a kinematic tree.
+            radius: Radius of the icosphere.
+            color: Color of the icosphere as an RGB tuple.
+            subdivisions: Number of subdivisions to use when creating the icosphere.
+            wxyz: Quaternion rotation to parent frame from local frame (R_pl).
+            position: Translation from parent frame to local frame (t_pl).
+            visible: Whether or not this icosphere is initially visible.
+
+        Returns:
+            Handle for manipulating scene node.
+        """
+        mesh = trimesh.creation.icosphere(subdivisions=subdivisions, radius=radius)
+
+        # We use add_mesh_simple() because it lets us do smooth shading;
+        # add_mesh_trimesh() currently does not.
+        return self.add_mesh_simple(
+            name=name,
+            vertices=mesh.vertices,
+            faces=mesh.faces,
+            color=color,
+            flat_shading=False,
+            position=position,
+            wxyz=wxyz,
+            visible=visible,
+        )
+
     def set_background_image(
         self,
         image: onp.ndarray,

src/viser/_messages.py

@@ -185,6 +185,7 @@ class MeshMessage(Message):
 
     wireframe: bool
     opacity: Optional[float]
+    flat_shading: bool
     side: Literal["front", "back", "double"]
     material: Literal["standard", "toon3", "toon5"]
 

src/viser/client/src/WebsocketInterface.tsx

@@ -261,6 +261,7 @@ function useMessageHandler() {
           wireframe: message.wireframe,
           transparent: message.opacity !== null,
           opacity: message.opacity ?? 1.0,
+          flatShading: message.flat_shading,
           side: {
             front: THREE.FrontSide,
             back: THREE.BackSide,

src/viser/client/src/WebsocketMessages.tsx

@@ -144,6 +144,7 @@ export interface MeshMessage {
   vertex_colors: Uint8Array | null;
   wireframe: boolean;
   opacity: number | null;
+  flat_shading: boolean;
   side: "front" | "back" | "double";
   material: "standard" | "toon3" | "toon5";
 }
@@ -346,7 +347,7 @@ export interface _GuiAddInputBase {
   hint: string | null;
   initial_value: any;
 }
-/** GuiAddButtonMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', initial_value: 'bool', color: "Optional[Literal[('dark', 'gray', 'red', 'pink', 'grape', 'violet', 'indigo', 'blue', 'cyan', 'green', 'lime', 'yellow', 'orange', 'teal')]]", icon_base64: 'Optional[str]')
+/** GuiAddButtonMessage(order: 'float', id: 'str', label: 'str', container_id: 'str', hint: 'Optional[str]', initial_value: 'bool', color: "Optional[Literal['dark', 'gray', 'red', 'pink', 'grape', 'violet', 'indigo', 'blue', 'cyan', 'green', 'lime', 'yellow', 'orange', 'teal']]", icon_base64: 'Optional[str]')
  *
  * (automatically generated)
  */