diff --git a/docs/developer_guides/viewer/viewer_control.md b/docs/developer_guides/viewer/viewer_control.md index 7a1bb5381..268030521 100644 --- a/docs/developer_guides/viewer/viewer_control.md +++ b/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) ``` diff --git a/nerfstudio/viewer/viewer_elements.py b/nerfstudio/viewer/viewer_elements.py index ec057acb7..654503a3c 100644 --- a/nerfstudio/viewer/viewer_elements.py +++ b/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.""" diff --git a/pyproject.toml b/pyproject.toml index e906d2442..5ac3f71b8 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.26", + "viser==0.1.27", "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", "xatlas",