Integrate scene click capability from original viewer (#2509)

* Integrate scene click capability from original viewer

* filter non-click ScenePointerEvents from ClickEvent

* lint

* Put click origin+direction in nerfstudio world scale

* Update docs

* Update doc heading

---------

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>

docs/developer_guides/viewer/viewer_control.md

@@ -62,8 +62,8 @@ class MyModel(nn.Module):  # Must inherit from nn.Module
         self.viewer_button = ViewerButton(name="Dummy Button",cb_hook=button_cb)
 ```
 
-## Double-click Callbacks
-We forward *double* clicks inside the viewer to the ViewerControl object, which you can use to interact with the scene. To do this, register a callback using `register_click_cb()`. The click is defined to be a ray that starts at the camera origin and passes through the click point on the screen, in world coordinates. 
+## Scene Click Callbacks
+We forward *single* clicks inside the viewer to the ViewerControl object, which you can use to interact with the scene. To do this, register a callback using `register_click_cb()`. The click is defined to be a ray that starts at the camera origin and passes through the click point on the screen, in world coordinates. 
 
 ```python
 from nerfstudio.viewer.server.viewer_elements import ViewerControl,ViewerClick
@@ -77,6 +77,16 @@ class MyModel(nn.Module):  # must inherit from nn.Module
         self.viewer_control.register_click_cb(click_cb)
 ```
 
+You can also use `unregister_click_cb()` to remove callbacks that are no longer needed. A good example is a "Click on Scene" button, that when pressed, would register a callback that would wait for the next click, and then unregister itself.
+```python
+    ...
+    def button_cb(button: ViewerButton):
+        def click_cb(click: ViewerClick):
+            print(f"Click at {click.origin} in direction {click.direction}")
+            self.viewer_control.unregister_click_cb(click_cb)
+        self.viewer_control.register_click_cb(click_cb)
+```
+
 ### Thread safety
 Just like `ViewerElement` callbacks, click callbacks are asynchronous to training and can potentially interrupt a call to `get_outputs()`.
 

nerfstudio/viewer_beta/viewer_elements.py

@@ -31,6 +31,7 @@ from viser import (
     GuiButtonHandle,
     GuiDropdownHandle,
     GuiInputHandle,
+    ScenePointerEvent,
     ViserServer,
 )
 
@@ -68,7 +69,7 @@ class ViewerControl:
 
     def __init__(self):
         # this should be a user-facing constructor, since it will be used inside the model/pipeline class
-        self.click_cbs = []
+        self._click_cbs = {}
 
     def _setup(self, viewer: Viewer):
         """
@@ -151,13 +152,35 @@ class ViewerControl:
             cb: The callback to call when a click is detected.
                 The callback should take a ViewerClick object as an argument
         """
-        self.click_cbs.append(cb)
+        from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO
 
-    def on_click(self, msg):
+        def wrapped_cb(scene_pointer_msg: ScenePointerEvent):
+            # only call the callback if the event is a click
+            if scene_pointer_msg.event != "click":
+                return
+            origin = scene_pointer_msg.ray_origin
+            direction = scene_pointer_msg.ray_direction
+
+            origin = tuple([x / VISER_NERFSTUDIO_SCALE_RATIO for x in origin])
+            assert len(origin) == 3
+
+            click = ViewerClick(origin, direction)
+            cb(click)
+
+        self._click_cbs[cb] = wrapped_cb
+        self.viser_server.on_scene_click(wrapped_cb)
+
+    def unregister_click_cb(self, cb: Callable):
         """
-        Internal use only, register a click in the viewer which propagates to all self.click_cbs
+        Remove a callback which will be called when a click is detected in the viewer.
+
+        Args:
+            cb: The callback to remove
         """
-        raise NotImplementedError()
+        if cb not in self._click_cbs:
+            raise ValueError(f"Callback {cb} not registered, cannot remove")
+        self.viser_server.remove_scene_click_callback(self._click_cbs[cb])
+        self._click_cbs.pop(cb)
 
     @property
     def server(self):