From 2ddb2631b5546a8caef16753a7204d28f4ee66b6 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Sat, 15 Apr 2023 01:14:38 -0700 Subject: [PATCH] Improve viewer element types (#1743) --- docs/developer_guides/viewer/custom_gui.md | 2 +- nerfstudio/viewer/server/viewer_elements.py | 119 +++++++++++--------- nerfstudio/viewer/viser/__init__.py | 1 + nerfstudio/viewer/viser/gui.py | 8 +- pyproject.toml | 2 +- 5 files changed, 76 insertions(+), 56 deletions(-) diff --git a/docs/developer_guides/viewer/custom_gui.md b/docs/developer_guides/viewer/custom_gui.md index 7ddd11a61..09275dac1 100644 --- a/docs/developer_guides/viewer/custom_gui.md +++ b/docs/developer_guides/viewer/custom_gui.md @@ -86,7 +86,7 @@ from nerfstudio.viewer.server.viewer_elements import * class MyModel(Model): def __init__(self): - self.a = ViewerButton(name="My Button", call_fn=self.handle_btn) + self.a = ViewerButton(name="My Button", cb_hook=self.handle_btn) self.b = ViewerNumber(name="Number", default_value=1.0) self.c = ViewerCheckbox(name="Checkbox", default_value=False) self.d = ViewerDropdown(name="Dropdown", default_value="A", options=["A", "B"]) diff --git a/nerfstudio/viewer/server/viewer_elements.py b/nerfstudio/viewer/server/viewer_elements.py index 6e602acc0..19d9e7588 100644 --- a/nerfstudio/viewer/server/viewer_elements.py +++ b/nerfstudio/viewer/server/viewer_elements.py @@ -17,14 +17,14 @@ from __future__ import annotations from abc import abstractmethod -from typing import Any, Callable, List, Optional, Tuple, TypeVar +from typing import Any, Callable, Generic, List, Optional, Tuple, TypeVar -from nerfstudio.viewer.viser import GuiHandle, ViserServer +from nerfstudio.viewer.viser import GuiHandle, GuiSelectHandle, ViserServer -IntOrFloat = TypeVar("IntOrFloat", int, float) +TValue = TypeVar("TValue") -class ViewerElement: +class ViewerElement(Generic[TValue]): """Base class for all viewer elements Args: @@ -32,9 +32,14 @@ class ViewerElement: disabled: If the element is disabled """ - def __init__(self, name: str, disabled: bool = False, cb_hook: Callable = lambda element: None): + def __init__( + self, + name: str, + disabled: bool = False, + cb_hook: Callable = lambda element: None, + ) -> None: self.name = name - self.gui_handle: Optional[GuiHandle] = None + self.gui_handle: Optional[GuiHandle[TValue]] = None self.disabled = disabled self.cb_hook = cb_hook @@ -48,18 +53,20 @@ class ViewerElement: """ ... - def remove(self): + def remove(self) -> None: """Removes the gui element from the viewer""" if self.gui_handle is not None: self.gui_handle.remove() self.gui_handle = None - def set_hidden(self, hidden: bool): + def set_hidden(self, hidden: bool) -> None: """Sets the hidden state of the gui element""" + assert self.gui_handle is not None self.gui_handle.set_hidden(hidden) - def set_disabled(self, disabled: bool): + def set_disabled(self, disabled: bool) -> None: """Sets the disabled state of the gui element""" + assert self.gui_handle is not None self.gui_handle.set_disabled(disabled) @abstractmethod @@ -68,17 +75,17 @@ class ViewerElement: ... -class ViewerButton(ViewerElement): +class ViewerButton(ViewerElement[bool]): """A button in the viewer Args: name: The name of the button - call_fn: The function to call when the button is pressed + cb_hook: The function to call when the button is pressed disabled: If the button is disabled """ - def __init__(self, name: str, call_fn: Callable, disabled: bool = False): - super().__init__(name, disabled=disabled, cb_hook=call_fn) + def __init__(self, name: str, cb_hook: Callable[[ViewerButton], Any], disabled: bool = False): + super().__init__(name, disabled=disabled, cb_hook=cb_hook) def _create_gui_handle(self, viser_server: ViserServer) -> None: self.gui_handle = viser_server.add_gui_button(self.name) @@ -87,14 +94,11 @@ class ViewerButton(ViewerElement): def install(self, viser_server: ViserServer) -> None: self._create_gui_handle(viser_server) - def call_fn(_: GuiHandle): # pylint: disable=unused-argument - self.cb_hook(self) - assert self.gui_handle is not None - self.gui_handle.on_update(call_fn) + self.gui_handle.on_update(lambda _: self.cb_hook(self)) -class ViewerParameter(ViewerElement): +class ViewerParameter(ViewerElement[TValue], Generic[TValue]): """A viewer element with state Args: @@ -104,9 +108,15 @@ class ViewerParameter(ViewerElement): cb_hook: Callback to call on update """ - def __init__(self, name: str, default_value: Any, disabled: bool = False, cb_hook: Callable = lambda element: None): + def __init__( + self, + name: str, + default_value: TValue, + disabled: bool = False, + cb_hook: Callable = lambda element: None, + ) -> None: super().__init__(name, disabled=disabled, cb_hook=cb_hook) - self.def_value = default_value + self.default_value = default_value def install(self, viser_server: ViserServer) -> None: """ @@ -117,32 +127,32 @@ class ViewerParameter(ViewerElement): """ self._create_gui_handle(viser_server) - def update_fn(_: GuiHandle): - self.cb_hook(self) - assert self.gui_handle is not None - self.gui_handle.on_update(update_fn) + self.gui_handle.on_update(lambda _: self.cb_hook(self)) @abstractmethod def _create_gui_handle(self, viser_server: ViserServer) -> None: ... @property - def value(self) -> Any: + def value(self) -> TValue: """Returns the current value of the viewer element""" if self.gui_handle is None: - return self.def_value + return self.default_value return self.gui_handle.get_value() @value.setter - def value(self, value: Any): + def value(self, value: TValue) -> None: if self.gui_handle is not None: self.gui_handle.set_value(value) else: - self.def_value = value + self.default_value = value -class ViewerSlider(ViewerParameter): +IntOrFloat = TypeVar("IntOrFloat", int, float) + + +class ViewerSlider(ViewerParameter[IntOrFloat], Generic[IntOrFloat]): """A slider in the viewer Args: @@ -163,7 +173,7 @@ class ViewerSlider(ViewerParameter): max_value: IntOrFloat, step: IntOrFloat = 0.1, disabled: bool = False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerSlider], Any] = lambda element: None, ): assert isinstance(default_value, (float, int)) super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) @@ -173,11 +183,11 @@ class ViewerSlider(ViewerParameter): def _create_gui_handle(self, viser_server: ViserServer) -> None: assert self.gui_handle is None, "gui_handle should be initialized once" - self.gui_handle = viser_server.add_gui_slider(self.name, self.min, self.max, self.step, self.def_value) + self.gui_handle = viser_server.add_gui_slider(self.name, self.min, self.max, self.step, self.default_value) self.gui_handle.set_disabled(self.disabled) -class ViewerText(ViewerParameter): +class ViewerText(ViewerParameter[str]): """A text field in the viewer Args: @@ -192,18 +202,18 @@ class ViewerText(ViewerParameter): name: str, default_value: str, disabled: bool = False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerText], Any] = lambda element: None, ): assert isinstance(default_value, str) super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) def _create_gui_handle(self, viser_server: ViserServer) -> None: assert self.gui_handle is None, "gui_handle should be initialized once" - self.gui_handle = viser_server.add_gui_text(self.name, self.def_value) + self.gui_handle = viser_server.add_gui_text(self.name, self.default_value) self.gui_handle.set_disabled(self.disabled) -class ViewerNumber(ViewerParameter): +class ViewerNumber(ViewerParameter[IntOrFloat], Generic[IntOrFloat]): """A number field in the viewer Args: @@ -213,22 +223,24 @@ class ViewerNumber(ViewerParameter): cb_hook: Callback to call on update """ + default_value: IntOrFloat + def __init__( self, name: str, default_value: IntOrFloat, disabled: bool = False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerNumber], Any] = lambda element: None, ): assert isinstance(default_value, (float, int)) super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) def _create_gui_handle(self, viser_server: ViserServer) -> None: assert self.gui_handle is None, "gui_handle should be initialized once" - self.gui_handle = viser_server.add_gui_number(self.name, self.def_value) + self.gui_handle = viser_server.add_gui_number(self.name, self.default_value) -class ViewerCheckbox(ViewerParameter): +class ViewerCheckbox(ViewerParameter[bool]): """A checkbox in the viewer Args: @@ -243,18 +255,18 @@ class ViewerCheckbox(ViewerParameter): name: str, default_value: bool, disabled: bool = False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerCheckbox], Any] = lambda element: None, ): assert isinstance(default_value, bool) super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) def _create_gui_handle(self, viser_server: ViserServer) -> None: assert self.gui_handle is None, "gui_handle should be initialized once" - self.gui_handle = viser_server.add_gui_checkbox(self.name, self.def_value) + self.gui_handle = viser_server.add_gui_checkbox(self.name, self.default_value) self.gui_handle.set_disabled(self.disabled) -class ViewerDropdown(ViewerParameter): +class ViewerDropdown(ViewerParameter[str]): """A dropdown in the viewer Args: @@ -265,13 +277,15 @@ class ViewerDropdown(ViewerParameter): cb_hook: Callback to call on update """ + gui_handle: Optional[GuiSelectHandle[str]] + def __init__( self, name: str, default_value: str, options: List[str], disabled: bool = False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerDropdown], Any] = lambda element: None, ): assert default_value in options super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) @@ -279,19 +293,22 @@ class ViewerDropdown(ViewerParameter): def _create_gui_handle(self, viser_server: ViserServer) -> None: assert self.gui_handle is None, "gui_handle should be initialized once" - self.gui_handle = viser_server.add_gui_select(self.name, self.options, self.def_value) + self.gui_handle = viser_server.add_gui_select(self.name, self.options, self.default_value) self.gui_handle.set_disabled(self.disabled) def set_options(self, new_options: List[str]) -> None: """ Sets the options of the dropdown, - new_options: The new options. if the current option isn't in the new options, the first option is selected + + Args: + new_options: The new options. If the current option isn't in the new options, the first option is selected. """ self.options = new_options - self.gui_handle.set_options(new_options) + if self.gui_handle is not None: + self.gui_handle.set_options(new_options) -class ViewerRGB(ViewerParameter): +class ViewerRGB(ViewerParameter[Tuple[int, int, int]]): """ An RGB color picker for the viewer """ @@ -301,17 +318,17 @@ class ViewerRGB(ViewerParameter): name, default_value: Tuple[int, int, int], disabled=False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerRGB], Any] = lambda element: None, ): assert len(default_value) == 3 super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) def _create_gui_handle(self, viser_server: ViserServer) -> None: - self.gui_handle = viser_server.add_gui_rgb(self.name, self.def_value) + self.gui_handle = viser_server.add_gui_rgb(self.name, self.default_value) self.gui_handle.set_disabled(self.disabled) -class ViewerVec3(ViewerParameter): +class ViewerVec3(ViewerParameter[Tuple[float, float, float]]): """ 3 number boxes in a row to input a vector """ @@ -322,12 +339,12 @@ class ViewerVec3(ViewerParameter): default_value: Tuple[float, float, float], step=0.1, disabled=False, - cb_hook: Callable = lambda element: None, + cb_hook: Callable[[ViewerVec3], Any] = lambda element: None, ): assert len(default_value) == 3 super().__init__(name, default_value, disabled=disabled, cb_hook=cb_hook) self.step = step def _create_gui_handle(self, viser_server: ViserServer) -> None: - self.gui_handle = viser_server.add_gui_vector3(self.name, self.def_value, self.step) + self.gui_handle = viser_server.add_gui_vector3(self.name, self.default_value, self.step) self.gui_handle.set_disabled(self.disabled) diff --git a/nerfstudio/viewer/viser/__init__.py b/nerfstudio/viewer/viser/__init__.py index 2a77fd123..9914442ed 100644 --- a/nerfstudio/viewer/viser/__init__.py +++ b/nerfstudio/viewer/viser/__init__.py @@ -16,5 +16,6 @@ # pylint: disable=useless-import-alias from .message_api import GuiHandle as GuiHandle +from .message_api import GuiSelectHandle as GuiSelectHandle from .messages import NerfstudioMessage as NerfstudioMessage from .server import ViserServer as ViserServer diff --git a/nerfstudio/viewer/viser/gui.py b/nerfstudio/viewer/viser/gui.py index 90d5b0edf..0e8e95e65 100644 --- a/nerfstudio/viewer/viser/gui.py +++ b/nerfstudio/viewer/viser/gui.py @@ -180,15 +180,17 @@ class GuiSelectHandle(GuiHandle[TString], Generic[TString]): inferred where possible when handles are instantiated; for the most flexibility, we can declare handles as `GuiHandle[str]`. """ - overwrite_value = False + + # Make sure initial value is in options. self._impl.leva_conf["options"] = options if self._impl.leva_conf["value"] not in options: self._impl.leva_conf["value"] = options[0] - overwrite_value = True + # Update options. self._impl.api._queue( GuiSetLevaConfMessage(self._impl.name, self._impl.leva_conf), ) - if overwrite_value: + # Make sure current value is in options. + if self.get_value() not in options: self.set_value(options[0]) diff --git a/pyproject.toml b/pyproject.toml index 4830047f5..e52f02f7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ dependencies = [ "torchmetrics[image]>=0.9.3", "torchtyping>=0.1.4", "torchvision>=0.13.0", - "viser==0.0.2", + "viser>=0.0.5", "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", "xatlas",