From 5865ecf3236e47818974354a793b5d94dbc0c272 Mon Sep 17 00:00:00 2001
From: Chung Min Kim <chungmin99@gmail.com>
Date: Wed, 29 Nov 2023 14:48:58 -0800
Subject: [PATCH] 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>
---
 .../developer_guides/viewer/viewer_control.md | 14 ++++++--
 nerfstudio/viewer_beta/viewer_elements.py     | 33 ++++++++++++++++---
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/docs/developer_guides/viewer/viewer_control.md b/docs/developer_guides/viewer/viewer_control.md
index e6a84d3d7..0bb33294c 100644
--- a/docs/developer_guides/viewer/viewer_control.md
+++ b/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()`.
 
diff --git a/nerfstudio/viewer_beta/viewer_elements.py b/nerfstudio/viewer_beta/viewer_elements.py
index 82a8b3d44..e8de855e0 100644
--- a/nerfstudio/viewer_beta/viewer_elements.py
+++ b/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):