RayClicks integration (nerfstudio-project#106)

* initial rayclick imp

* Example code bug

* Move scene pointer logic to diff file, add pointer on/off logic

* Updated names + example

* Remove pointer listener if removed all callbacks

* Formatting changes

* Tweaks

* Docs + backwards compatibility

* Update cursor for scene clicks

* Add events.md

---------

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

docs/source/events.md

@@ -0,0 +1,23 @@
+# Events
+
+We define a small set of event types, which are passed to callback functions
+when events like clicks or GUI updates are triggered.
+
+<!-- prettier-ignore-start -->
+
+.. autoapiclass:: viser.ScenePointerEvent
+   :members:
+   :undoc-members:
+   :inherited-members:
+
+.. autoapiclass:: viser.SceneNodePointerEvent
+   :members:
+   :undoc-members:
+   :inherited-members:
+
+.. autoapiclass:: viser.GuiEvent
+   :members:
+   :undoc-members:
+   :inherited-members:
+
+<!-- prettier-ignore-end -->

docs/source/gui_handles.md

@@ -50,9 +50,4 @@ connected clients. When a GUI element is added to a client (for example, via
    :undoc-members:
    :inherited-members:
 
-.. autoapiclass:: viser.GuiEvent
-   :members:
-   :undoc-members:
-   :inherited-members:
-
 <!-- prettier-ignore-end -->

docs/source/index.md

@@ -66,6 +66,7 @@ URL (default: `http://localhost:8080`).
    ./client_handles.md
    ./gui_handles.md
    ./scene_handles.md
+   ./events.md
    ./icons.md
 
 

docs/source/scene_handles.md

@@ -25,6 +25,11 @@ connected clients. When a scene node is added to a client (for example, via
    :undoc-members:
    :inherited-members:
 
+.. autoapiclass:: viser.GlbHandle
+   :members:
+   :undoc-members:
+   :inherited-members:
+
 .. autoapiclass:: viser.Gui3dContainerHandle
    :members:
    :undoc-members:
@@ -55,9 +60,4 @@ connected clients. When a scene node is added to a client (for example, via
    :undoc-members:
    :inherited-members:
 
-.. autoapiclass:: viser.ClickEvent
-   :members:
-   :undoc-members:
-   :inherited-members:
-
 <!-- prettier-ignore-end -->

examples/06_mesh.py

@@ -25,13 +25,13 @@
     name="/simple",
     vertices=vertices,
     faces=faces,
-    wxyz=tf.SO3.exp(onp.array([onp.pi / 2, 0.0, 0.0])).wxyz,
+    wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,
     position=(0.0, 0.0, 0.0),
 )
 server.add_mesh_trimesh(
     name="/trimesh",
     mesh=mesh.smoothed(),
-    wxyz=tf.SO3.exp(onp.array([onp.pi / 2, 0.0, 0.0])).wxyz,
+    wxyz=tf.SO3.from_x_radians(onp.pi / 2).wxyz,
     position=(0.0, 5.0, 0.0),
 )
 

src/viser/__init__.py

@@ -11,7 +11,6 @@
 from ._gui_handles import GuiTabHandle as GuiTabHandle
 from ._icons_enum import Icon as Icon
 from ._scene_handles import CameraFrustumHandle as CameraFrustumHandle
-from ._scene_handles import ClickEvent as ClickEvent
 from ._scene_handles import FrameHandle as FrameHandle
 from ._scene_handles import GlbHandle as GlbHandle
 from ._scene_handles import Gui3dContainerHandle as Gui3dContainerHandle
@@ -20,6 +19,8 @@
 from ._scene_handles import MeshHandle as MeshHandle
 from ._scene_handles import PointCloudHandle as PointCloudHandle
 from ._scene_handles import SceneNodeHandle as SceneNodeHandle
+from ._scene_handles import SceneNodePointerEvent as SceneNodePointerEvent
+from ._scene_handles import ScenePointerEvent as ScenePointerEvent
 from ._scene_handles import TransformControlsHandle as TransformControlsHandle
 from ._viser import CameraHandle as CameraHandle
 from ._viser import ClientHandle as ClientHandle
@@ -28,3 +29,4 @@
 if not TYPE_CHECKING:
     # Backwards compatibility.
     GuiHandle = GuiInputHandle
+    ClickEvent = SceneNodePointerEvent

src/viser/_message_api.py

@@ -14,7 +14,17 @@
 import queue
 import threading
 import time
-from typing import TYPE_CHECKING, Dict, Optional, Tuple, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import imageio.v3 as iio
 import numpy as onp
@@ -27,7 +37,6 @@
 from . import _messages, infra, theme
 from ._scene_handles import (
     CameraFrustumHandle,
-    ClickEvent,
     FrameHandle,
     GlbHandle,
     Gui3dContainerHandle,
@@ -36,11 +45,14 @@
     MeshHandle,
     PointCloudHandle,
     SceneNodeHandle,
+    SceneNodePointerEvent,
+    ScenePointerEvent,
     TransformControlsHandle,
     _TransformControlsState,
 )
 
 if TYPE_CHECKING:
+    from ._viser import ClientHandle
     from .infra import ClientId
 
 
@@ -138,13 +150,21 @@ def __init__(self, handler: infra.MessageHandler) -> None:
         ] = {}
         self._handle_from_node_name: Dict[str, SceneNodeHandle] = {}
 
+        # Callbacks for scene pointer events -- by default don't enable them.
+        self._scene_pointer_cb: List[Callable[[ScenePointerEvent], None]] = []
+        self._scene_pointer_enabled = False
+
         handler.register_handler(
             _messages.TransformControlsUpdateMessage,
             self._handle_transform_controls_updates,
         )
         handler.register_handler(
-            _messages.SceneNodeClickedMessage,
-            self._handle_click_updates,
+            _messages.SceneNodeClickMessage,
+            self._handle_node_click_updates,
+        )
+        handler.register_handler(
+            _messages.ScenePointerMessage,
+            self._handle_scene_pointer_updates,
         )
 
         self._atomic_lock = threading.Lock()
@@ -598,6 +618,25 @@ def _queue_unsafe(self, message: _messages.Message) -> None:
         """Abstract method for sending messages."""
         ...
 
+    def _get_client_handle(self, client_id: ClientId) -> ClientHandle:
+        """Private helper for getting a client handle from its ID."""
+        # Avoid circular imports.
+        from ._viser import ClientHandle, ViserServer
+
+        # Implementation-wise, note that MessageApi is never directly instantiated.
+        # Instead, it serves as a mixin/base class for either ViserServer, which
+        # maintains a registry of connected clients, or ClientHandle, which should
+        # only ever be dealing with its own client_id.
+        if isinstance(self, ViserServer):
+            # TODO: there's a potential race condition here when the client disconnects.
+            # This probably applies to multiple other parts of the code, we should
+            # revisit all of the cases where we index into connected_clients.
+            return self._state.connected_clients[client_id]
+        else:
+            assert isinstance(self, ClientHandle)
+            assert client_id == self.client_id
+            return self
+
     def _handle_transform_controls_updates(
         self, client_id: ClientId, message: _messages.TransformControlsUpdateMessage
     ) -> None:
@@ -617,17 +656,62 @@ def _handle_transform_controls_updates(
         if handle._impl_aux.sync_cb is not None:
             handle._impl_aux.sync_cb(client_id, handle)
 
-    def _handle_click_updates(
-        self, client_id: ClientId, message: _messages.SceneNodeClickedMessage
+    def _handle_node_click_updates(
+        self, client_id: ClientId, message: _messages.SceneNodeClickMessage
     ) -> None:
         """Callback for handling click messages."""
         handle = self._handle_from_node_name.get(message.name, None)
         if handle is None or handle._impl.click_cb is None:
             return
         for cb in handle._impl.click_cb:
-            event = ClickEvent(client_id=client_id, target=handle)
+            event = SceneNodePointerEvent(
+                client=self._get_client_handle(client_id),
+                client_id=client_id,
+                event="click",
+                target=handle,
+                ray_origin=message.ray_origin,
+                ray_direction=message.ray_direction,
+            )
             cb(event)  # type: ignore
 
+    def _handle_scene_pointer_updates(
+        self, client_id: ClientId, message: _messages.ScenePointerMessage
+    ):
+        """Callback for handling click messages."""
+        for cb in self._scene_pointer_cb:
+            event = ScenePointerEvent(
+                client=self._get_client_handle(client_id),
+                client_id=client_id,
+                event=message.event_type,
+                ray_origin=message.ray_origin,
+                ray_direction=message.ray_direction,
+            )
+            cb(event)
+
+    def on_scene_click(
+        self,
+        func: Callable[[ScenePointerEvent], None],
+    ) -> Callable[[ScenePointerEvent], None]:
+        """Add a callback for scene pointer events."""
+        self._scene_pointer_cb.append(func)
+
+        # Notify client of a new listener. This can help the client determine whether
+        # or not click events should still be sent; note that we have no way of knowing
+        # here because both server and client handles manage their own callbacks.
+        self._queue(_messages.ScenePointerCallbackInfoMessage(count=1))
+        return func
+
+    def remove_scene_click_callback(
+        self,
+        func: Callable[[ScenePointerEvent], None],
+    ) -> None:
+        """Check for the function handle in the list of callbacks and remove it."""
+        if func in self._scene_pointer_cb:
+            self._scene_pointer_cb.remove(func)
+
+        # Notify client that the listener has been removed.
+        self._queue(_messages.ScenePointerCallbackInfoMessage(count=-1))
+
     def add_3d_gui_container(
         self,
         name: str,

src/viser/_messages.py

@@ -50,6 +50,26 @@ class ViewerCameraMessage(Message):
     up_direction: Tuple[float, float, float]
 
 
+@dataclasses.dataclass
+class ScenePointerMessage(Message):
+    """Message for a raycast-like pointer in the scene.
+    origin is the viewing camera position, in world coordinates.
+    direction is the vector if a ray is projected from the camera through the clicked pixel,
+    """
+
+    # Later we can add `double_click`, `move`, `down`, `up`, etc.
+    event_type: Literal["click"]
+    ray_origin: Tuple[float, float, float]
+    ray_direction: Tuple[float, float, float]
+
+
+@dataclasses.dataclass
+class ScenePointerCallbackInfoMessage(Message):
+    """Message to enable/disable scene pointer"""
+
+    count: int
+
+
 @dataclasses.dataclass
 class CameraFrustumMessage(Message):
     """Variant of CameraMessage used for visualizing camera frustums.
@@ -277,10 +297,12 @@ class SetSceneNodeClickableMessage(Message):
 
 
 @dataclasses.dataclass
-class SceneNodeClickedMessage(Message):
+class SceneNodeClickMessage(Message):
     """Message for clicked objects."""
 
     name: str
+    ray_origin: Tuple[float, float, float]
+    ray_direction: Tuple[float, float, float]
 
 
 @dataclasses.dataclass

src/viser/_scene_handles.py

@@ -12,6 +12,7 @@
     Dict,
     Generic,
     List,
+    Literal,
     Optional,
     Tuple,
     Type,
@@ -26,6 +27,17 @@
     from ._gui_api import GuiApi
     from ._gui_handles import SupportsRemoveProtocol
     from ._message_api import ClientId, MessageApi
+    from ._viser import ClientHandle
+
+
+@dataclasses.dataclass(frozen=True)
+class ScenePointerEvent:
+    client: ClientHandle
+    client_id: ClientId
+    # Later we can add `double_click`, `move`, `down`, `up`, etc
+    event: Literal["click"]
+    ray_origin: Tuple[float, float, float]
+    ray_direction: Tuple[float, float, float]
 
 
 TSceneNodeHandle = TypeVar("TSceneNodeHandle", bound="SceneNodeHandle")
@@ -43,7 +55,9 @@ class _SceneNodeHandleState:
     )
     visible: bool = True
     # TODO: we should remove SceneNodeHandle as an argument here.
-    click_cb: Optional[List[Callable[[ClickEvent[SceneNodeHandle]], None]]] = None
+    click_cb: Optional[
+        List[Callable[[SceneNodePointerEvent[SceneNodeHandle]], None]]
+    ] = None
 
 
 @dataclasses.dataclass
@@ -121,16 +135,21 @@ def remove(self) -> None:
 
 
 @dataclasses.dataclass(frozen=True)
-class ClickEvent(Generic[TSceneNodeHandle]):
+class SceneNodePointerEvent(Generic[TSceneNodeHandle]):
+    client: ClientHandle
     client_id: ClientId
+    event: Literal["click"]
     target: TSceneNodeHandle
+    ray_origin: Tuple[float, float, float]
+    ray_direction: Tuple[float, float, float]
 
 
 @dataclasses.dataclass
 class _ClickableSceneNodeHandle(SceneNodeHandle):
     def on_click(
-        self: TSceneNodeHandle, func: Callable[[ClickEvent[TSceneNodeHandle]], None]
-    ) -> Callable[[ClickEvent[TSceneNodeHandle]], None]:
+        self: TSceneNodeHandle,
+        func: Callable[[SceneNodePointerEvent[TSceneNodeHandle]], None],
+    ) -> Callable[[SceneNodePointerEvent[TSceneNodeHandle]], None]:
         """Attach a callback for when a scene node is clicked.
 
         TODO:

src/viser/client/src/App.tsx

@@ -15,6 +15,7 @@ import {
 import { BlendFunction, KernelSize } from "postprocessing";
 
 import { SynchronizedCameraControls } from "./CameraControls";
+import { ScenePointerControls } from "./ScenePointerControls";
 import { Box, MantineProvider, MediaQuery } from "@mantine/core";
 import React from "react";
 import { SceneNodeThreeObject, UseSceneTree } from "./SceneTree";
@@ -67,6 +68,7 @@ export type ViewerContextContents = {
     "ready" | "triggered" | "pause" | "in_progress"
   >;
   getRenderRequest: React.MutableRefObject<null | GetRenderRequestMessage>;
+  scenePointerCallbackCount: React.MutableRefObject<number>;
 };
 export const ViewerContext = React.createContext<null | ViewerContextContents>(
   null,
@@ -108,6 +110,7 @@ function ViewerRoot() {
     messageQueueRef: React.useRef([]),
     getRenderRequestState: React.useRef("ready"),
     getRenderRequest: React.useRef(null),
+    scenePointerCallbackCount: React.useRef(0),
   };
 
   return (
@@ -180,6 +183,7 @@ function ViewerCanvas({ children }: { children: React.ReactNode }) {
       <AdaptiveEvents />
       <SceneContextSetter />
       <SynchronizedCameraControls />
+      <ScenePointerControls />
       <Selection>
         <SceneNodeThreeObject name="" parent={null} />
         <EffectComposer enabled={true} autoClear={false}>