Add scene.add_line_segments()

docs/source/examples/18_splines.rst → docs/source/examples/18_lines.rst

@@ -1,11 +1,11 @@
 .. Comment: this file is automatically generated by `update_example_docs.py`.
    It should not be modified manually.
 
-Splines
+Lines
 ==========================================
 
 
-Make a ball with some random splines.
+Make a ball with some random line segments and splines.
 
 
 
@@ -16,17 +16,35 @@ Make a ball with some random splines.
         import time
 
         import numpy as np
-
         import viser
 
 
         def main() -> None:
             server = viser.ViserServer()
+
+            # Line segments.
+            #
+            # This will be much faster than creating separate scene objects for
+            # individual line segments or splines.
+            N = 2000
+            points = np.random.normal(size=(N, 2, 3)) * 3.0
+            colors = np.random.randint(0, 255, size=(N, 2, 3))
+            server.scene.add_line_segments(
+                f"/line_segments",
+                points=points,
+                colors=colors,
+                line_width=3.0,
+            )
+
+            # Spline helpers.
+            #
+            # If many lines are needed, it'll be more efficient to batch them in
+            # `add_line_segments()`.
             for i in range(10):
-                positions = np.random.normal(size=(30, 3)) * 3.0
+                points = np.random.normal(size=(30, 3)) * 3.0
                 server.scene.add_spline_catmull_rom(
-                    f"/catmull_{i}",
-                    positions,
+                    f"/catmull/{i}",
+                    positions=points,
                     tension=0.5,
                     line_width=3.0,
                     color=np.random.uniform(size=3),
@@ -35,9 +53,9 @@ Make a ball with some random splines.
 
                 control_points = np.random.normal(size=(30 * 2 - 2, 3)) * 3.0
                 server.scene.add_spline_cubic_bezier(
-                    f"/cubic_bezier_{i}",
-                    positions,
-                    control_points,
+                    f"/cubic_bezier/{i}",
+                    positions=points,
+                    control_points=control_points,
                     line_width=3.0,
                     color=np.random.uniform(size=3),
                     segments=100,

examples/18_lines.py

@@ -0,0 +1,59 @@
+"""Lines
+
+Make a ball with some random line segments and splines.
+"""
+
+import time
+
+import numpy as np
+import viser
+
+
+def main() -> None:
+    server = viser.ViserServer()
+
+    # Line segments.
+    #
+    # This will be much faster than creating separate scene objects for
+    # individual line segments or splines.
+    N = 2000
+    points = np.random.normal(size=(N, 2, 3)) * 3.0
+    colors = np.random.randint(0, 255, size=(N, 2, 3))
+    server.scene.add_line_segments(
+        f"/line_segments",
+        points=points,
+        colors=colors,
+        line_width=3.0,
+    )
+
+    # Spline helpers.
+    #
+    # If many lines are needed, it'll be more efficient to batch them in
+    # `add_line_segments()`.
+    for i in range(10):
+        points = np.random.normal(size=(30, 3)) * 3.0
+        server.scene.add_spline_catmull_rom(
+            f"/catmull/{i}",
+            positions=points,
+            tension=0.5,
+            line_width=3.0,
+            color=np.random.uniform(size=3),
+            segments=100,
+        )
+
+        control_points = np.random.normal(size=(30 * 2 - 2, 3)) * 3.0
+        server.scene.add_spline_cubic_bezier(
+            f"/cubic_bezier/{i}",
+            positions=points,
+            control_points=control_points,
+            line_width=3.0,
+            color=np.random.uniform(size=3),
+            segments=100,
+        )
+
+    while True:
+        time.sleep(10.0)
+
+
+if __name__ == "__main__":
+    main()

src/viser/_messages.py

@@ -1183,6 +1183,26 @@ class ThemeConfigurationMessage(Message):
     colors: Optional[Tuple[str, str, str, str, str, str, str, str, str, str]]
 
 
+@dataclasses.dataclass
+class LineSegmentsMessage(Message, tag="SceneNodeMessage"):
+    """Message from server->client carrying line segments information."""
+
+    name: str
+    props: LineSegmentsProps
+
+
+@dataclasses.dataclass
+class LineSegmentsProps:
+    points: npt.NDArray[np.float32]
+    """A numpy array of shape (N, 2, 3) containing a batched set of line
+    segments. Synchronized automatically when assigned."""
+    line_width: float
+    """Width of the lines. Synchronized automatically when assigned."""
+    colors: npt.NDArray[np.uint8]
+    """Numpy array of shape (N, 2, 3) containing a color for each point.
+    Synchronized automatically when assigned."""
+
+
 @dataclasses.dataclass
 class CatmullRomSplineMessage(Message, tag="SceneNodeMessage"):
     """Message from server->client carrying Catmull-Rom spline information."""
@@ -1193,6 +1213,7 @@ class CatmullRomSplineMessage(Message, tag="SceneNodeMessage"):
 
 @dataclasses.dataclass
 class CatmullRomSplineProps:
+    # TODO: consider renaming positions to points and using numpy arrays for consistency with LineSegmentsProps.
     positions: Tuple[Tuple[float, float, float], ...]
     """A tuple of 3D positions (x, y, z) defining the spline's path. Synchronized automatically when assigned."""
     curve_type: Literal["centripetal", "chordal", "catmullrom"]

src/viser/_scene_api.py

@@ -28,6 +28,7 @@
     HemisphereLightHandle,
     ImageHandle,
     LabelHandle,
+    LineSegmentsHandle,
     MeshHandle,
     MeshSkinnedBoneHandle,
     MeshSkinnedHandle,
@@ -562,9 +563,58 @@ def add_glb(
         message = _messages.GlbMessage(name, _messages.GlbProps(glb_data, scale))
         return GlbHandle._make(self, message, name, wxyz, position, visible)
 
+    def add_line_segments(
+        self,
+        name: str,
+        points: np.ndarray,
+        colors: np.ndarray | tuple[float, float, float],
+        line_width: float = 1,
+        wxyz: tuple[float, float, float, float] | np.ndarray = (1.0, 0.0, 0.0, 0.0),
+        position: tuple[float, float, float] | np.ndarray = (0.0, 0.0, 0.0),
+        visible: bool = True,
+    ) -> LineSegmentsHandle:
+        """Add line segments to the scene.
+
+        Args:
+            name: A scene tree name. Names in the format of /parent/child can
+                be used to define a kinematic tree.
+            points: A numpy array of shape (N, 2, 3) defining start/end points
+                for each of N line segments.
+            colors: Colors of points. Should have shape (N, 2, 3) or be
+                broadcastable to it.
+            line_width: Width of the lines.
+            wxyz: Quaternion rotation to parent frame from local frame (R_pl).
+            position: Translation to parent frame from local frame (t_pl).
+            visible: Whether or not these line segments are initially visible.
+
+        Returns:
+            Handle for manipulating scene node.
+        """
+        points_array = np.asarray(points, dtype=np.float32)
+        if (
+            points_array.shape[-1] != 3
+            or points_array.ndim != 3
+            or points_array.shape[1] != 2
+        ):
+            raise ValueError("Points should have shape (N, 2, 3) for N line segments.")
+
+        colors_array = colors_to_uint8(np.asarray(colors))
+        colors_array = np.broadcast_to(colors_array, points_array.shape)
+
+        message = _messages.LineSegmentsMessage(
+            name=name,
+            props=_messages.LineSegmentsProps(
+                points=points_array,
+                colors=colors_array,
+                line_width=line_width,
+            ),
+        )
+        return LineSegmentsHandle._make(self, message, name, wxyz, position, visible)
+
     def add_spline_catmull_rom(
         self,
         name: str,
+        # The naming inconsistency here compared to add_line_segments is unfortunate...
         positions: tuple[tuple[float, float, float], ...] | np.ndarray,
         curve_type: Literal["centripetal", "chordal", "catmullrom"] = "centripetal",
         tension: float = 0.5,
@@ -581,6 +631,9 @@ def add_spline_catmull_rom(
         This method creates a spline based on a set of positions and interpolates
         them using the Catmull-Rom algorithm. This can be used to create smooth curves.
 
+        If many splines are needed, it'll be more efficient to batch them in
+        :meth:`add_line_segments()`.
+
         Args:
             name: A scene tree name. Names in the format of /parent/child can be used to
                 define a kinematic tree.
@@ -637,6 +690,9 @@ def add_spline_cubic_bezier(
         positions and control points. It is useful for creating complex, smooth,
         curving shapes.
 
+        If many splines are needed, it'll be more efficient to batch them in
+        :meth:`add_line_segments()`.
+
         Args:
             name: A scene tree name. Names in the format of /parent/child can be used to
                 define a kinematic tree.

src/viser/_scene_handles.py

@@ -502,6 +502,14 @@ class GridHandle(
     """Handle for grid objects."""
 
 
+class LineSegmentsHandle(
+    SceneNodeHandle,
+    _messages.LineSegmentsProps,
+    _OverridableScenePropApi if not TYPE_CHECKING else object,
+):
+    """Handle for line segments objects."""
+
+
 class SplineCatmullRomHandle(
     SceneNodeHandle,
     _messages.CatmullRomSplineProps,

src/viser/client/src/SceneTree.tsx

@@ -2,6 +2,7 @@ import {
   CatmullRomLine,
   CubicBezierLine,
   Grid,
+  Line,
   PivotControls,
   useCursor,
 } from "@react-three/drei";
@@ -418,6 +419,52 @@ function useObjectFactory(message: SceneNodeMessage | undefined): {
         ),
       };
     }
+    case "LineSegmentsMessage": {
+      return {
+        makeObject: (ref) => {
+          // The array conversion here isn't very efficient. We go from buffer
+          // => TypeArray => Javascript Array, then back to buffers in drei's
+          // <Line /> abstraction.
+          const pointsArray = new Float32Array(
+            message.props.points.buffer.slice(
+              message.props.points.byteOffset,
+              message.props.points.byteOffset + message.props.points.byteLength,
+            ),
+          );
+          const colorArray = new Uint8Array(
+            message.props.colors.buffer.slice(
+              message.props.colors.byteOffset,
+              message.props.colors.byteOffset + message.props.colors.byteLength,
+            ),
+          );
+          return (
+            <group ref={ref}>
+              <Line
+                points={Array.from(
+                  { length: pointsArray.length / 3 },
+                  (_, i) => [
+                    pointsArray[i * 3],
+                    pointsArray[i * 3 + 1],
+                    pointsArray[i * 3 + 2],
+                  ],
+                )}
+                color="white"
+                lineWidth={message.props.line_width}
+                vertexColors={Array.from(
+                  { length: colorArray.length / 3 },
+                  (_, i) => [
+                    colorArray[i * 3] / 255,
+                    colorArray[i * 3 + 1] / 255,
+                    colorArray[i * 3 + 2] / 255,
+                  ],
+                )}
+                segments={true}
+              />
+            </group>
+          );
+        },
+      };
+    }
     case "CatmullRomSplineMessage": {
       return {
         makeObject: (ref) => {