Refactor ScenePointerCallbacks for viser compatibility (#3036)

* Update to match viser pr

* Bump viser version

* `done` => `removed` in docs

* Update viewer_elements.py

* Fix

* Point function signature overload

* Update viewer_control.md

* type ignore on callback

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>

docs/developer_guides/viewer/viewer_control.md

@@ -73,7 +73,7 @@ To do this, register a callback using `register_pointer_cb()`.
 
 You can also use `unregister_pointer_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.
 
-Note that the viewer can only listen to *one* scene pointer callback at a time. If you register a new callback, the old one will be unregistered! Be warned that if the callback includes GUI state changes (e.g., re-enabling a disabled button), they may be lost. You can ensure that the GUI state is restored by providing a `done_cb` function that will be called after 1) the callback is run, or 2) the callback is unregistered w/o execution.
+Note that the viewer can only listen to *one* scene pointer callback at a time. If you register a new callback, the old one will be unregistered! Be warned that if the callback includes GUI state changes (e.g., re-enabling a disabled button), they may be lost. You can ensure that the GUI state is restored by providing a `removed_cb` function that will be called after the callback is removed.
 
 ```python
 from nerfstudio.viewer.viewer_elements import ViewerControl,ViewerClick
@@ -104,19 +104,22 @@ class MyModel(nn.Module):  # must inherit from nn.Module
 
         # Or make a button that, once pressed, listens to clicks in the viewer.
         # Here, the button is disabled while it is listening to clicks.
+        # The button will become enabled again if either:
+        # - the callback is removed in `pointer_click_cb`, with the `unregister...`, or
+        # - the callback is overridden by the viewer (to listen to another callback).
         def button_cb(button: ViewerButton):
             def pointer_click_cb(click: ViewerClick):
                 ...
-
-            def pointer_click_cb_done():
-                self.viewer_button.set_disabled(False)
                 self.viewer_control.unregister_pointer_cb()
 
+            def pointer_click_removed_cb():
+                self.viewer_button.set_disabled(False)
+
             self.viewer_button.set_disabled(True)
             self.viewer_control.register_pointer_cb(
                 "click",
                 cb=pointer_click_cb,
-                done_cb=pointer_click_cb_done
+                removed_cb=pointer_click_removed_cb
             )
         self.viewer_button = ViewerButton(name="Click on Scene", cb_hook=button_cb)
 ```

nerfstudio/viewer/viewer_elements.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 import warnings
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generic, List, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Generic, List, Literal, Optional, Tuple, Union, overload
 
 import numpy as np
 import torch
@@ -166,8 +166,29 @@ class ViewerControl:
         CONSOLE.log("`register_click_cb` is deprecated, use `register_pointer_cb` instead.")
         self.register_pointer_cb("click", cb)
 
+    @overload
     def register_pointer_cb(
-        self, event_type: Literal["click", "rect-select"], cb: Callable, done_cb: Optional[Callable] = None
+        self,
+        event_type: Literal["click"],
+        cb: Callable[[ViewerClick], None],
+        removed_cb: Optional[Callable[[], None]] = None,
+    ):
+        ...
+
+    @overload
+    def register_pointer_cb(
+        self,
+        event_type: Literal["rect-select"],
+        cb: Callable[[ViewerRectSelect], None],
+        removed_cb: Optional[Callable[[], None]] = None,
+    ):
+        ...
+
+    def register_pointer_cb(
+        self,
+        event_type: Literal["click", "rect-select"],
+        cb: Callable[[ViewerClick], None] | Callable[[ViewerRectSelect], None],
+        removed_cb: Optional[Callable[[], None]] = None,
     ):
         """
         Add a callback which will be called when a scene pointer event is detected in the viewer.
@@ -180,6 +201,7 @@ class ViewerControl:
 
         Args:
             cb: The callback to call when a click or a rect-select is detected.
+            removed_cb: The callback to run when the pointer event is removed.
         """
         from nerfstudio.viewer.viewer import VISER_NERFSTUDIO_SCALE_RATIO
 
@@ -203,7 +225,7 @@ class ViewerControl:
             else:
                 raise ValueError(f"Unknown event type: {scene_pointer_msg.event_type}")
 
-            cb(pointer_event)
+            cb(pointer_event)  # type: ignore
 
         cb_overriden = False
         with warnings.catch_warnings(record=True) as w:
@@ -219,8 +241,8 @@ class ViewerControl:
             )
 
         # If there exists a cleanup callback after the pointer event is done, register it.
-        if done_cb:
-            self.viser_server.on_scene_pointer_done(done_cb)
+        if removed_cb is not None:
+            self.viser_server.on_scene_pointer_removed(removed_cb)
 
     def unregister_click_cb(self, cb: Optional[Callable] = None):
         """Deprecated, use unregister_pointer_cb instead. `cb` is ignored."""

pyproject.toml

@@ -57,7 +57,7 @@ dependencies = [
     "torchvision>=0.14.1",
     "torchmetrics[image]>=1.0.1",
     "typing_extensions>=4.4.0",
-    "viser==0.1.26",
+    "viser==0.1.27",
     "nuscenes-devkit>=1.1.1",
     "wandb>=0.13.3",
     "xatlas",