Support rectangle selection in viewer (#3014)

* Support rectangle selection in viewer

* Update doc for click screen_pos

* Raise callback override warning from nerfstudio

* Bump viser version, rect-select only in >0.1.26

docs/developer_guides/viewer/viewer_control.md

@@ -62,8 +62,18 @@ class MyModel(nn.Module):  # Must inherit from nn.Module
         self.viewer_button = ViewerButton(name="Dummy Button",cb_hook=button_cb)
 ```
 
-## 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. 
+## Scene Pointer Callbacks
+We forward user interactions with the viewer to the `ViewerControl` object, which you can use to interact with the scene. 
+
+We currently support:
+ - `ViewerClick`: *single* clicks inside the viewer. 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. 
+ - `ViewerRectSelect`: drag to select a rectangle in the viewer screen. The rectangle is defined by two points (top-left and bottom-right corners) in normalized OpenCV screen coordinates.
+
+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.
 
 ```python
 from nerfstudio.viewer.viewer_elements import ViewerControl,ViewerClick
@@ -72,19 +82,43 @@ class MyModel(nn.Module):  # must inherit from nn.Module
     def __init__(self):
         # Must be a class variable
         self.viewer_control = ViewerControl()  # no arguments
-        def click_cb(click: ViewerClick):
-            print(f"Click at {click.origin} in direction {click.direction}")
-        self.viewer_control.register_click_cb(click_cb)
-```
+        
+        # Listen to clicks in the viewer...
+        def pointer_click_cb(click: ViewerClick):
+            print(f"Click at {click.origin} in direction {click.direction}, screen position {click.screen_pos}.")
+        self.viewer_control.register_pointer_cb("click", pointer_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)
+        # Listen to rectangle selections in the viewer...
+        def pointer_rect_cb(rect: ViewerRectSelect):
+            print(f"Rectangular selection from {rect.min_bounds} to {rect.max_bounds}.")
+        self.viewer_control.register_pointer_cb("click", pointer_rect_cb)
+
+        ... 
+        # Or make a button that, once pressed, listens to clicks in the viewer.
+        def button_cb(button: ViewerButton):
+            def pointer_click_cb(click: ViewerClick):
+                ...
+                self.viewer_control.unregister_pointer_cb()
+            self.viewer_control.register_pointer_cb("click", pointer_click_cb)
+        self.viewer_button = ViewerButton(name="Click on Scene", cb_hook=button_cb)
+
+        # Or make a button that, once pressed, listens to clicks in the viewer.
+        # Here, the button is disabled while it is listening to clicks.
+        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()
+
+            self.viewer_button.set_disabled(True)
+            self.viewer_control.register_pointer_cb(
+                "click",
+                cb=pointer_click_cb,
+                done_cb=pointer_click_cb_done
+            )
+        self.viewer_button = ViewerButton(name="Click on Scene", cb_hook=button_cb)
 ```
 
 ### Thread safety

nerfstudio/viewer/viewer_elements.py

@@ -18,9 +18,10 @@
 
 from __future__ import annotations
 
+import warnings
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Generic, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Generic, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -36,6 +37,7 @@ from viser import (
 )
 
 from nerfstudio.cameras.cameras import Cameras, CameraType
+from nerfstudio.utils.rich_utils import CONSOLE
 from nerfstudio.viewer.utils import CameraState, get_camera
 
 if TYPE_CHECKING:
@@ -60,6 +62,23 @@ class ViewerClick:
     The direction of the click if projected from the camera through the clicked pixel,
     in world coordinates
     """
+    screen_pos: Tuple[float, float]
+    """The screen position of the click in OpenCV screen coordinates, normalized to [0, 1]"""
+
+
+@dataclass
+class ViewerRectSelect:
+    """
+    Class representing a rectangle selection in the viewer (screen-space).
+
+    The screen coordinates follow OpenCV image coordinates, with the origin at the top-left corner,
+    but the bounds are also normalized to [0, 1] in both dimensions.
+    """
+
+    min_bounds: Tuple[float, float]
+    """The minimum bounds of the rectangle selection in screen coordinates."""
+    max_bounds: Tuple[float, float]
+    """The maximum bounds of the rectangle selection in screen coordinates."""
 
 
 class ViewerControl:
@@ -67,10 +86,6 @@ class ViewerControl:
     class for exposing non-gui controls of the viewer to the user
     """
 
-    def __init__(self):
-        # this should be a user-facing constructor, since it will be used inside the model/pipeline class
-        self._click_cbs = {}
-
     def _setup(self, viewer: Viewer):
         """
         Internal use only, setup the viewer control with the viewer state object
@@ -147,42 +162,83 @@ class ViewerControl:
         return get_camera(camera_state, img_height, img_width)
 
     def register_click_cb(self, cb: Callable):
+        """Deprecated, use register_pointer_cb instead."""
+        CONSOLE.log("`register_click_cb` is deprecated, use `register_pointer_cb` instead.")
+        self.register_pointer_cb("click", cb)
+
+    def register_pointer_cb(
+        self, event_type: Literal["click", "rect-select"], cb: Callable, done_cb: Optional[Callable] = None
+    ):
         """
-        Add a callback which will be called when a click is detected in the viewer.
+        Add a callback which will be called when a scene pointer event is detected in the viewer.
+        Scene pointer events include:
+        - "click": A click event, which includes the origin and direction of the click
+        - "rect-select": A rectangle selection event, which includes the screen bounds of the box selection
+
+        The callback should take a ViewerClick object as an argument if the event type is "click",
+        and a ViewerRectSelect object as an argument if the event type is "rect-select".
 
         Args:
-            cb: The callback to call when a click is detected.
-                The callback should take a ViewerClick object as an argument
+            cb: The callback to call when a click or a rect-select is detected.
         """
         from nerfstudio.viewer.viewer import VISER_NERFSTUDIO_SCALE_RATIO
 
         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
+            # Check that the event type is the same as the one we are interested in.
+            if scene_pointer_msg.event_type != event_type:
+                raise ValueError(f"Expected event type {event_type}, got {scene_pointer_msg.event_type}")
 
-            origin = tuple([x / VISER_NERFSTUDIO_SCALE_RATIO for x in origin])
-            assert len(origin) == 3
+            if scene_pointer_msg.event_type == "click":
+                origin = scene_pointer_msg.ray_origin
+                direction = scene_pointer_msg.ray_direction
+                screen_pos = scene_pointer_msg.screen_pos[0]
+                assert (origin is not None) and (
+                    direction is not None
+                ), "Origin and direction should not be None for click event."
+                origin = tuple([x / VISER_NERFSTUDIO_SCALE_RATIO for x in origin])
+                assert len(origin) == 3
+                pointer_event = ViewerClick(origin, direction, screen_pos)
+            elif scene_pointer_msg.event_type == "rect-select":
+                pointer_event = ViewerRectSelect(scene_pointer_msg.screen_pos[0], scene_pointer_msg.screen_pos[1])
+            else:
+                raise ValueError(f"Unknown event type: {scene_pointer_msg.event_type}")
 
-            click = ViewerClick(origin, direction)
-            cb(click)
+            cb(pointer_event)
 
-        self._click_cbs[cb] = wrapped_cb
-        self.viser_server.on_scene_click(wrapped_cb)
+        cb_overriden = False
+        with warnings.catch_warnings(record=True) as w:
+            # Register the callback with the viser server.
+            self.viser_server.on_scene_pointer(event_type=event_type)(wrapped_cb)
+            # If there exists a warning, it's because a callback was overriden.
+            cb_overriden = len(w) > 0
 
-    def unregister_click_cb(self, cb: Callable):
+        if cb_overriden:
+            warnings.warn(
+                "A ScenePointer callback has already been registered for this event type. "
+                "The new callback will override the existing one."
+            )
+
+        # 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)
+
+    def unregister_click_cb(self, cb: Optional[Callable] = None):
+        """Deprecated, use unregister_pointer_cb instead. `cb` is ignored."""
+        warnings.warn("`unregister_click_cb` is deprecated, use `unregister_pointer_cb` instead.")
+        if cb is not None:
+            # raise warning
+            warnings.warn("cb argument is ignored in unregister_click_cb.")
+
+        self.unregister_pointer_cb()
+
+    def unregister_pointer_cb(self):
         """
-        Remove a callback which will be called when a click is detected in the viewer.
+        Remove a callback which will be called, when a scene pointer event is detected in the viewer.
 
         Args:
             cb: The callback to remove
         """
-        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)
+        self.viser_server.remove_scene_pointer_callback()
 
     @property
     def server(self):

pyproject.toml

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