Add instance_index to click event for batched axes (#272)

* Add `instance_index` to click event for batched axes

* Fix pyright error

src/viser/_messages.py

@@ -10,7 +10,6 @@
     Callable,
     ClassVar,
     Dict,
-    List,
     Optional,
     Tuple,
     Type,
@@ -124,7 +123,7 @@ class ScenePointerMessage(Message):
     event_type: ScenePointerEventType
     ray_origin: Optional[Tuple[float, float, float]]
     ray_direction: Optional[Tuple[float, float, float]]
-    screen_pos: List[Tuple[float, float]]
+    screen_pos: Tuple[Tuple[float, float], ...]
 
 
 @dataclasses.dataclass
@@ -469,8 +468,11 @@ class SceneNodeClickMessage(Message):
     """Message for clicked objects."""
 
     name: str
+    instance_index: Optional[int]
+    """Instance index. Currently only used for batched axes."""
     ray_origin: Tuple[float, float, float]
     ray_direction: Tuple[float, float, float]
+    screen_pos: Tuple[float, float]
 
 
 @dataclasses.dataclass

src/viser/_scene_api.py

@@ -481,8 +481,8 @@ def add_frame(
 
         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.
+        batching and calling :meth:`add_batched_axes()` may be a better choice
+        than calling :meth:`add_frame()` in a loop.
 
         Args:
             name: A scene tree name. Names in the format of /parent/child can be used to
@@ -524,10 +524,11 @@ def add_batched_axes(
     ) -> 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.
+        The functionality of :meth:`add_batched_axes()` overlaps significantly
+        with :meth:`add_frame()` when `show_axes=True`. The primary difference
+        is that :meth:`add_batched_axes()` supports multiple axes via the
+        `wxyzs_batched` (shape Nx4) and `positions_batched` (shape Nx3)
+        arguments.
 
         Axes that are batched and rendered via a single call to
         `add_batched_axes()` are instanced on the client; this will be much
@@ -1323,6 +1324,8 @@ def _handle_node_click_updates(
                 target=handle,
                 ray_origin=message.ray_origin,
                 ray_direction=message.ray_direction,
+                screen_pos=message.screen_pos,
+                instance_index=message.instance_index,
             )
             cb(event)  # type: ignore
 

src/viser/_scene_handles.py

@@ -30,7 +30,7 @@ class ScenePointerEvent:
     """Origin of 3D ray corresponding to this click, in world coordinates."""
     ray_direction: tuple[float, float, float] | None
     """Direction of 3D ray corresponding to this click, in world coordinates."""
-    screen_pos: list[tuple[float, float]]
+    screen_pos: tuple[tuple[float, float], ...]
     """Screen position of the click on the screen (OpenCV image coordinates, 0 to 1).
     (0, 0) is the upper-left corner, (1, 1) is the bottom-right corner.
     For a box selection, this includes the min- and max- corners of the box."""
@@ -159,6 +159,11 @@ class SceneNodePointerEvent(Generic[TSceneNodeHandle]):
     """Origin of 3D ray corresponding to this click, in world coordinates."""
     ray_direction: tuple[float, float, float]
     """Direction of 3D ray corresponding to this click, in world coordinates."""
+    screen_pos: tuple[float, float]
+    """Screen position of the click on the screen (OpenCV image coordinates, 0 to 1).
+    (0, 0) is the upper-left corner, (1, 1) is the bottom-right corner."""
+    instance_index: int | None
+    """Instance ID of the clicked object, if applicable. Currently this is `None` for all objects except for the output of :meth:`SceneApi.add_batched_axes()`."""
 
 
 @dataclasses.dataclass

src/viser/client/src/MessageHandler.tsx

@@ -165,34 +165,43 @@ function useMessageHandler() {
       // Add axes to visualize.
       case "BatchedAxesMessage": {
         addSceneNodeMakeParents(
-          new SceneNode<THREE.Group>(message.name, (ref) => (
-            // Minor naming discrepancy: I think "batched" will be clearer to
-            // folks on the Python side, but instanced is somewhat more
-            // precise.
-            <InstancedAxes
-              ref={ref}
-              wxyzsBatched={
-                new Float32Array(
-                  message.wxyzs_batched.buffer.slice(
-                    message.wxyzs_batched.byteOffset,
-                    message.wxyzs_batched.byteOffset +
-                      message.wxyzs_batched.byteLength,
-                  ),
-                )
-              }
-              positionsBatched={
-                new Float32Array(
-                  message.positions_batched.buffer.slice(
-                    message.positions_batched.byteOffset,
-                    message.positions_batched.byteOffset +
-                      message.positions_batched.byteLength,
-                  ),
-                )
-              }
-              axes_length={message.axes_length}
-              axes_radius={message.axes_radius}
-            />
-          )),
+          new SceneNode<THREE.Group>(
+            message.name,
+            (ref) => (
+              // Minor naming discrepancy: I think "batched" will be clearer to
+              // folks on the Python side, but instanced is somewhat more
+              // precise.
+              <InstancedAxes
+                ref={ref}
+                wxyzsBatched={
+                  new Float32Array(
+                    message.wxyzs_batched.buffer.slice(
+                      message.wxyzs_batched.byteOffset,
+                      message.wxyzs_batched.byteOffset +
+                        message.wxyzs_batched.byteLength,
+                    ),
+                  )
+                }
+                positionsBatched={
+                  new Float32Array(
+                    message.positions_batched.buffer.slice(
+                      message.positions_batched.byteOffset,
+                      message.positions_batched.byteOffset +
+                        message.positions_batched.byteLength,
+                    ),
+                  )
+                }
+                axes_length={message.axes_length}
+                axes_radius={message.axes_radius}
+              />
+            ),
+            undefined,
+            undefined,
+            undefined,
+            // Compute click instance index from instance ID. Each visualized
+            // frame has 1 instance for each of 3 line segments.
+            (instanceId) => Math.floor(instanceId! / 3),
+          ),
         );
         return;
       }

src/viser/client/src/SceneTree.tsx

@@ -11,6 +11,7 @@ import { useSceneTreeState } from "./SceneTreeState";
 import { ErrorBoundary } from "react-error-boundary";
 import { rayToViserCoords } from "./WorldTransformUtils";
 import { HoverableContext } from "./ThreeAssets";
+import { opencvXyFromPointerXy } from "./ClickUtils";
 
 export type MakeObject<T extends THREE.Object3D = THREE.Object3D> = (
   ref: React.Ref<T>,
@@ -34,6 +35,10 @@ export class SceneNode<T extends THREE.Object3D = THREE.Object3D> {
      */
     public readonly unmountWhenInvisible?: boolean,
     public readonly everyFrameCallback?: () => void,
+    /** For click events on instanced nodes, like batched axes, we want to keep track of which. */
+    public readonly computeClickInstanceIndexFromInstanceId?: (
+      instanceId: number | undefined,
+    ) => number | null,
   ) {
     this.children = [];
     this.clickable = false;
@@ -137,6 +142,10 @@ export function SceneNodeThreeObject(props: {
   const everyFrameCallback = viewer.useSceneTree(
     (state) => state.nodeFromName[props.name]?.everyFrameCallback,
   );
+  const computeClickInstanceIndexFromInstanceId = viewer.useSceneTree(
+    (state) =>
+      state.nodeFromName[props.name]?.computeClickInstanceIndexFromInstanceId,
+  );
   const [unmount, setUnmount] = React.useState(false);
   const clickable =
     viewer.useSceneTree((state) => state.nodeFromName[props.name]?.clickable) ??
@@ -326,16 +335,30 @@ export function SceneNodeThreeObject(props: {
               if (state.dragging) return;
               // Convert ray to viser coordinates.
               const ray = rayToViserCoords(viewer, e.ray);
+
+              // Send OpenCV image coordinates to the server (normalized).
+              const canvasBbox =
+                viewer.canvasRef.current!.getBoundingClientRect();
+              const mouseVectorOpenCV = opencvXyFromPointerXy(viewer, [
+                e.clientX - canvasBbox.left,
+                e.clientY - canvasBbox.top,
+              ]);
+
               sendClicksThrottled({
                 type: "SceneNodeClickMessage",
                 name: props.name,
+                instance_index:
+                  computeClickInstanceIndexFromInstanceId === undefined
+                    ? null
+                    : computeClickInstanceIndexFromInstanceId(e.instanceId),
                 // Note that the threejs up is +Y, but we expose a +Z up.
                 ray_origin: [ray.origin.x, ray.origin.y, ray.origin.z],
                 ray_direction: [
                   ray.direction.x,
                   ray.direction.y,
                   ray.direction.z,
                 ],
+                screen_pos: [mouseVectorOpenCV.x, mouseVectorOpenCV.y],
               });
             }}
             onPointerOver={(e) => {

src/viser/client/src/WebsocketMessages.tsx

@@ -370,8 +370,10 @@ export interface SetSceneNodeClickableMessage {
 export interface SceneNodeClickMessage {
   type: "SceneNodeClickMessage";
   name: string;
+  instance_index: number | null;
   ray_origin: [number, number, number];
   ray_direction: [number, number, number];
+  screen_pos: [number, number];
 }
 /** Reset scene.
  *