From 71de07163a9642016c9bde1d3705e3fb62babd69 Mon Sep 17 00:00:00 2001 From: Chung Min Kim Date: Mon, 25 Mar 2024 12:48:49 -0700 Subject: [PATCH] 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 --- .../developer_guides/viewer/viewer_control.md | 62 +++++++--- nerfstudio/viewer/viewer_elements.py | 106 +++++++++++++----- pyproject.toml | 2 +- 3 files changed, 130 insertions(+), 40 deletions(-) diff --git a/docs/developer_guides/viewer/viewer_control.md b/docs/developer_guides/viewer/viewer_control.md index 7490a3f67..7a1bb5381 100644 --- a/docs/developer_guides/viewer/viewer_control.md +++ b/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 diff --git a/nerfstudio/viewer/viewer_elements.py b/nerfstudio/viewer/viewer_elements.py index c9851acb3..ec057acb7 100644 --- a/nerfstudio/viewer/viewer_elements.py +++ b/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): diff --git a/pyproject.toml b/pyproject.toml index af9ebc393..a181b5c11 100644 --- a/pyproject.toml +++ b/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",