Add button groups to viewer (#1905)

* Add button groups to viewer

* Cleanup

* Add label

* Port some typing changes from viser, note that `disabled` doesn't work

* black

---------

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

nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx

@@ -6,6 +6,7 @@ import {
   useControls,
   folder,
   button,
+  buttonGroup,
 } from 'leva';
 import { useDispatch, useSelector } from 'react-redux';
 import { Box } from '@mui/material';
@@ -51,6 +52,20 @@ function CustomLeva() {
           value: true,
         });
       }, levaConf.settings);
+    } else if (levaConf.type === 'BUTTON_GROUP') {
+      // Add a button group.
+      leafFolder[key] = buttonGroup({
+        label: levaConf.label,
+        opts: levaConf.options.reduce((acc, opt) => {
+          acc[opt] = () =>
+            sendWebsocketMessage(viser_websocket, {
+              type: 'GuiUpdateMessage',
+              name: key,
+              value: opt,
+            });
+          return acc;
+        }, {}),
+      });
     } else {
       // Add any other kind of input.
       const throttledSender = makeThrottledMessageSender(viser_websocket, 25);
@@ -135,7 +150,10 @@ export function RenderControls() {
           "& input[type='checkbox']~label svg path": {
             stroke: '#222831 !important',
           },
-          '& button': { color: '#222831 !important', height: '2em' },
+          '& button': {
+            backgroundColor: '#393E46 !important',
+            height: '2em',
+          },
         }}
       >
         <Leva

nerfstudio/viewer/app/src/themes/leva_theme.json

@@ -1,6 +1,6 @@
 {
   "colors": {
-    "elevation1": "#292d39",
+    "elevation1": "#393E46",
     "elevation2": "#222831",
     "elevation3": "#393E46",
     "accent1": "#ffc640",

nerfstudio/viewer/server/control_panel.py

@@ -20,6 +20,7 @@ import torch
 
 from nerfstudio.utils.colormaps import ColormapOptions, Colormaps
 from nerfstudio.viewer.server.viewer_elements import (
+    ViewerButtonGroup,
     ViewerCheckbox,
     ViewerDropdown,
     ViewerElement,
@@ -55,8 +56,11 @@ class ControlPanel:
         self.viser_server = viser_server
         self._elements_by_tag: DefaultDict[str, List[ViewerElement]] = defaultdict(lambda: [])
 
-        self._train_speed = ViewerDropdown(
-            "Train Speed", "Balanced", ["Fast", "Balanced", "Slow"], cb_hook=lambda han: self._train_speed_cb()
+        self._train_speed = ViewerButtonGroup(
+            name="Train Speed  ",
+            default_value="Balanced",
+            options=["Slow", "Balanced", "Fast"],
+            cb_hook=lambda han: self._train_speed_cb(),
         )
         self._output_render = ViewerDropdown(
             "Output Render",
@@ -138,13 +142,13 @@ class ControlPanel:
 
     def _train_speed_cb(self) -> None:
         """Callback for when the train speed is changed"""
-        if self.train_speed == "Fast":
+        if self._train_speed.value == "Fast":
             self._train_util.value = 0.95
             self._max_res.value = 256
-        elif self.train_speed == "Balanced":
+        elif self._train_speed.value == "Balanced":
             self._train_util.value = 0.85
             self._max_res.value = 512
-        elif self.train_speed == "Slow":
+        elif self._train_speed.value == "Slow":
             self._train_util.value = 0.5
             self._max_res.value = 1024
 
@@ -204,11 +208,6 @@ class ControlPanel:
         """
         self._split_colormap.set_options(_get_colormap_options(dimensions, dtype))
 
-    @property
-    def train_speed(self) -> str:
-        """Returns the current train speed setting"""
-        return self._train_speed.value
-
     @property
     def output_render(self) -> str:
         """Returns the current output render"""

nerfstudio/viewer/server/viewer_elements.py

@@ -315,6 +315,35 @@ class ViewerDropdown(ViewerParameter[TString], Generic[TString]):
             self.gui_handle.set_options(new_options)
 
 
+class ViewerButtonGroup(ViewerParameter[TString], Generic[TString]):
+    """A button group in the viewer. Unlike other fields, cannot be disabled.
+
+    Args:
+        name: The name of the button group
+        default_value: The default value of the button group
+        options: The options of the button group
+        cb_hook: Callback to call on update
+    """
+
+    gui_handle: Optional[GuiHandle[TString]]
+    default_value: TString
+
+    def __init__(
+        self,
+        name: str,
+        default_value: TString,
+        options: List[TString],
+        cb_hook: Callable[[ViewerDropdown], Any] = lambda element: None,
+    ):
+        assert default_value in options
+        super().__init__(name, default_value, disabled=False, cb_hook=cb_hook)
+        self.options = options
+
+    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_button_group(self.name, self.options, self.default_value)  # type: ignore
+
+
 class ViewerRGB(ViewerParameter[Tuple[int, int, int]]):
     """
     An RGB color picker for the viewer

nerfstudio/viewer/viser/message_api.py

@@ -303,6 +303,47 @@ class MessageApi(abc.ABC):
             )._impl
         )
 
+    # Resolve type of value to a Literal whenever possible.
+    @overload
+    def add_gui_button_group(
+        self,
+        name: str,
+        options: List[TLiteralString],
+        initial_value: Optional[TLiteralString] = None,
+    ) -> GuiHandle[TLiteralString]:
+        ...
+
+    @overload
+    def add_gui_button_group(
+        self,
+        name: str,
+        options: List[str],
+        initial_value: Optional[str] = None,
+    ) -> GuiHandle[str]:
+        ...
+
+    def add_gui_button_group(
+        self,
+        name: str,
+        options: List[TLiteralString] | List[str],
+        initial_value: Optional[TLiteralString | str] = None,
+    ) -> GuiHandle[TLiteralString] | GuiHandle[str]:
+        """Add a button group to the GUI.
+
+        Args:
+            name: The name of the button group.
+            options: The options to choose from.
+            initial_value: The initial value of the button group.
+        """
+        assert len(options) > 0
+        if initial_value is None:
+            initial_value = options[0]
+        return self._add_gui_impl(
+            name,
+            initial_value,
+            leva_conf={"type": "BUTTON_GROUP", "label": name, "options": options},
+        )
+
     def add_gui_slider(
         self,
         name: str,