From 23d2bf31ad5ad1c47d39cba8463c5222e4f9be3b Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Tue, 23 Jan 2024 00:05:21 -0800 Subject: [PATCH] `add_frames_batched()` => `add_batched_axes()` --- src/viser/_message_api.py | 57 +++++++++++++-------- src/viser/_messages.py | 11 ++-- src/viser/_scene_handles.py | 2 +- src/viser/client/src/ThreeAssets.tsx | 8 +-- src/viser/client/src/WebsocketInterface.tsx | 11 ++-- src/viser/client/src/WebsocketMessages.tsx | 13 ++--- 6 files changed, 57 insertions(+), 45 deletions(-) diff --git a/src/viser/_message_api.py b/src/viser/_message_api.py index 6afdd21a9..b03169e3e 100644 --- a/src/viser/_message_api.py +++ b/src/viser/_message_api.py @@ -38,9 +38,9 @@ from . import _messages, infra, theme from . import transforms as tf from ._scene_handles import ( + BatchedAxesHandle, CameraFrustumHandle, FrameHandle, - FramesBatchedHandle, GlbHandle, Gui3dContainerHandle, ImageHandle, @@ -553,8 +553,14 @@ def add_frame( ) -> FrameHandle: """Add a coordinate frame to the scene. - This method is used for adding a visual representation of a coordinate frame, - which can help in understanding the orientation and position of objects in 3D space. + This method is used for adding a visual representation of a coordinate + frame, which can help in understanding the orientation and position of + objects in 3D space. + + For cases where we want to visualize many coordinate frames, like + trajectories containing thousands or tens of thousands of frames, + batching and calling `add_batched_axes()` may be a better choice than calling + `add_frame()` in a loop. Args: name: A scene tree name. Names in the format of /parent/child can be used to @@ -579,53 +585,60 @@ def add_frame( ) return FrameHandle._make(self, name, wxyz, position, visible) - def add_frames_batched( + def add_batched_axes( self, name: str, - wxyzs_batched: Tuple[Tuple[float, float, float, float], ...] | onp.ndarray, - positions_batched: Tuple[Tuple[float, float, float], ...] | onp.ndarray, + batched_wxyzs: Tuple[Tuple[float, float, float, float], ...] | onp.ndarray, + batched_positions: Tuple[Tuple[float, float, float], ...] | onp.ndarray, axes_length: float = 0.5, axes_radius: float = 0.025, 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, - ) -> FramesBatchedHandle: - """Visualize a batched set of coordinate frames. + ) -> BatchedAxesHandle: + """Visualize batched sets of coordinate frame axes. + + The functionality of `add_batched_axes()` overlaps significantly with + `add_frame()` when `show_axes=True`. The primary difference is that + `add_batched_axes()` supports multiple axes via the `wxyzs_batched` + (shape Nx4) and `positions_batched` (shape Nx3) arguments. - This method is useful for adding visual representation of many - coordinate frames (tens or hundreds of thousands), particularly when - `add_frame()` is too slow. + Axes that are batched and rendered via a single call to + `add_batched_axes()` are instanced on the client; this will be much + faster to render than `add_frame()` called in a loop. Args: name: A scene tree name. Names in the format of /parent/child can be used to define a kinematic tree. - wxyzs_batched: Float array of shape (N,4). - positions_batched: Float array of shape (N,3). + batched_wxyzs: Float array of shape (N,4). + batched_positions: Float array of shape (N,3). axes_length: Length of each axis. axes_radius: Radius of each axis. wxyz: Quaternion rotation to parent frame from local frame (R_pl). + This will be applied to all axes. position: Translation to parent frame from local frame (t_pl). + This will be applied to all axes. visible: Whether or not this scene node is initially visible. Returns: Handle for manipulating scene node. """ - wxyzs_batched = onp.asarray(wxyzs_batched) - positions_batched = onp.asarray(positions_batched) + batched_wxyzs = onp.asarray(batched_wxyzs) + batched_positions = onp.asarray(batched_positions) - num_frames = wxyzs_batched.shape[0] - assert wxyzs_batched.shape == (num_frames, 4) - assert positions_batched.shape == (num_frames, 3) + num_axes = batched_wxyzs.shape[0] + assert batched_wxyzs.shape == (num_axes, 4) + assert batched_positions.shape == (num_axes, 3) self._queue( - _messages.FrameBatchedMessage( + _messages.BatchedAxesMessage( name=name, - wxyzs_batched=wxyzs_batched.astype(onp.float32), - positions_batched=positions_batched.astype(onp.float32), + wxyzs_batched=batched_wxyzs.astype(onp.float32), + positions_batched=batched_positions.astype(onp.float32), axes_length=axes_length, axes_radius=axes_radius, ) ) - return FramesBatchedHandle._make(self, name, wxyz, position, visible) + return BatchedAxesHandle._make(self, name, wxyz, position, visible) def add_grid( self, diff --git a/src/viser/_messages.py b/src/viser/_messages.py index dbceabfd6..9cf8cd411 100644 --- a/src/viser/_messages.py +++ b/src/viser/_messages.py @@ -96,10 +96,7 @@ class GlbMessage(Message): @dataclasses.dataclass class FrameMessage(Message): - """Coordinate frame message. - - Position and orientation should follow a `T_parent_local` convention, which - corresponds to the R matrix and t vector in `p_parent = [R | t] p_local`.""" + """Coordinate frame message.""" name: str show_axes: bool = True @@ -108,10 +105,10 @@ class FrameMessage(Message): @dataclasses.dataclass -class FrameBatchedMessage(Message): - """Batched coordinate frames message. +class BatchedAxesMessage(Message): + """Batched axes message. - Position and orientation should follow a `T_parent_local` convention, which + Positions and orientations should follow a `T_parent_local` convention, which corresponds to the R matrix and t vector in `p_parent = [R | t] p_local`.""" name: str diff --git a/src/viser/_scene_handles.py b/src/viser/_scene_handles.py index cb6e78764..89fe3086c 100644 --- a/src/viser/_scene_handles.py +++ b/src/viser/_scene_handles.py @@ -187,7 +187,7 @@ class PointCloudHandle(SceneNodeHandle): @dataclasses.dataclass -class FramesBatchedHandle(SceneNodeHandle): +class BatchedAxesHandle(_ClickableSceneNodeHandle): """Handle for batched coordinate frames.""" diff --git a/src/viser/client/src/ThreeAssets.tsx b/src/viser/client/src/ThreeAssets.tsx index 6ce892919..24513c4ec 100644 --- a/src/viser/client/src/ThreeAssets.tsx +++ b/src/viser/client/src/ThreeAssets.tsx @@ -300,7 +300,7 @@ export const CoordinateFrame = React.forwardRef< }); /** Helper for adding batched/instanced coordinate frames as scene nodes. */ -export const CoordinateFrameBatched = React.forwardRef< +export const InstancedAxes = React.forwardRef< THREE.Group, { wxyzsBatched: Float32Array; @@ -308,7 +308,7 @@ export const CoordinateFrameBatched = React.forwardRef< axes_length?: number; axes_radius?: number; } ->(function CoordinateFrameBatched( +>(function InstancedAxes( { wxyzsBatched: instance_wxyzs, positionsBatched: instance_positions, @@ -393,7 +393,9 @@ export const CoordinateFrameBatched = React.forwardRef< + > + + ); }); diff --git a/src/viser/client/src/WebsocketInterface.tsx b/src/viser/client/src/WebsocketInterface.tsx index 95f72c629..bc627f516 100644 --- a/src/viser/client/src/WebsocketInterface.tsx +++ b/src/viser/client/src/WebsocketInterface.tsx @@ -13,7 +13,7 @@ import { syncSearchParamServer } from "./SearchParamsUtils"; import { CameraFrustum, CoordinateFrame, - CoordinateFrameBatched as CoordinateFramesBatched, + InstancedAxes, GlbAsset, OutlinesIfHovered, PointCloud, @@ -154,11 +154,14 @@ function useMessageHandler() { return; } - // Add a coordinate frame. - case "FrameBatchedMessage": { + // Add axes to visualize. + case "BatchedAxesMessage": { addSceneNodeMakeParents( new SceneNode(message.name, (ref) => ( -