From f5ff9ae8b7e6f6444718ef382b06b1c262138ece Mon Sep 17 00:00:00 2001 From: brentyi Date: Sat, 2 Nov 2024 15:32:54 -0700 Subject: [PATCH] Add `scene.add_line_segments()` --- .../examples/{18_splines.rst => 18_lines.rst} | 36 ++++++++--- examples/18_lines.py | 59 +++++++++++++++++++ examples/18_splines.py | 41 ------------- src/viser/_messages.py | 21 +++++++ src/viser/_scene_api.py | 56 ++++++++++++++++++ src/viser/_scene_handles.py | 8 +++ src/viser/client/src/SceneTree.tsx | 47 +++++++++++++++ src/viser/client/src/WebsocketMessages.ts | 12 ++++ 8 files changed, 230 insertions(+), 50 deletions(-) rename docs/source/examples/{18_splines.rst => 18_lines.rst} (50%) create mode 100644 examples/18_lines.py delete mode 100644 examples/18_splines.py diff --git a/docs/source/examples/18_splines.rst b/docs/source/examples/18_lines.rst similarity index 50% rename from docs/source/examples/18_splines.rst rename to docs/source/examples/18_lines.rst index 16186249..275bbc67 100644 --- a/docs/source/examples/18_splines.rst +++ b/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, diff --git a/examples/18_lines.py b/examples/18_lines.py new file mode 100644 index 00000000..036d7afc --- /dev/null +++ b/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() diff --git a/examples/18_splines.py b/examples/18_splines.py deleted file mode 100644 index 8566dbb9..00000000 --- a/examples/18_splines.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Splines - -Make a ball with some random splines. -""" - -import time - -import numpy as np - -import viser - - -def main() -> None: - server = viser.ViserServer() - for i in range(10): - positions = np.random.normal(size=(30, 3)) * 3.0 - server.scene.add_spline_catmull_rom( - f"/catmull_{i}", - positions, - 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, - control_points, - line_width=3.0, - color=np.random.uniform(size=3), - segments=100, - ) - - while True: - time.sleep(10.0) - - -if __name__ == "__main__": - main() diff --git a/src/viser/_messages.py b/src/viser/_messages.py index 50c019f3..9922c0f2 100644 --- a/src/viser/_messages.py +++ b/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"] diff --git a/src/viser/_scene_api.py b/src/viser/_scene_api.py index 02c6d9b0..2c847a82 100644 --- a/src/viser/_scene_api.py +++ b/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. diff --git a/src/viser/_scene_handles.py b/src/viser/_scene_handles.py index a4c0dc00..280e2355 100644 --- a/src/viser/_scene_handles.py +++ b/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, diff --git a/src/viser/client/src/SceneTree.tsx b/src/viser/client/src/SceneTree.tsx index ceccfa93..91d10287 100644 --- a/src/viser/client/src/SceneTree.tsx +++ b/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 + // 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 ( + + [ + 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} + /> + + ); + }, + }; + } case "CatmullRomSplineMessage": { return { makeObject: (ref) => { diff --git a/src/viser/client/src/WebsocketMessages.ts b/src/viser/client/src/WebsocketMessages.ts index 3acdd771..66291459 100644 --- a/src/viser/client/src/WebsocketMessages.ts +++ b/src/viser/client/src/WebsocketMessages.ts @@ -994,6 +994,15 @@ export interface ThemeConfigurationMessage { ] | null; } +/** Message from server->client carrying line segments information. + * + * (automatically generated) + */ +export interface LineSegmentsMessage { + type: "LineSegmentsMessage"; + name: string; + props: { points: Uint8Array; line_width: number; colors: Uint8Array }; +} /** Message from server->client carrying Catmull-Rom spline information. * * (automatically generated) @@ -1191,6 +1200,7 @@ export type Message = | GuiUpdateMessage | SceneNodeUpdateMessage | ThemeConfigurationMessage + | LineSegmentsMessage | CatmullRomSplineMessage | CubicBezierSplineMessage | GaussianSplatsMessage @@ -1222,6 +1232,7 @@ export type SceneNodeMessage = | SkinnedMeshMessage | TransformControlsMessage | ImageMessage + | LineSegmentsMessage | CatmullRomSplineMessage | CubicBezierSplineMessage | GaussianSplatsMessage; @@ -1263,6 +1274,7 @@ const typeSetSceneNodeMessage = new Set([ "SkinnedMeshMessage", "TransformControlsMessage", "ImageMessage", + "LineSegmentsMessage", "CatmullRomSplineMessage", "CubicBezierSplineMessage", "GaussianSplatsMessage",