GUI fixes and improvements

viser/_gui_api.py

@@ -26,7 +26,7 @@
 
 import imageio.v3 as iio
 import numpy as onp
-from typing_extensions import LiteralString
+from typing_extensions import Literal, LiteralString
 
 from . import _messages
 from ._gui_handles import (
@@ -45,6 +45,8 @@
     _GuiInputHandle,
     _make_unique_id,
 )
+from ._icons import base64_from_icon
+from ._icons_enum import Icon
 from ._message_api import MessageApi, _encode_image_base64, cast_vector
 
 if TYPE_CHECKING:
@@ -159,7 +161,18 @@ def _handle_gui_updates(
 
         # Trigger callbacks.
         for cb in handle_state.update_cb:
-            cb(GuiEvent(client_id, handle))
+            from ._viser import ClientHandle, ViserServer
+
+            # Get the handle of the client that triggered this event.
+            api = self._get_api()
+            if isinstance(api, ClientHandle):
+                client = api
+            elif isinstance(api, ViserServer):
+                client = api.get_clients()[client_id]
+            else:
+                assert False
+
+            cb(GuiEvent(client_id, client, handle))
         if handle_state.sync_cb is not None:
             handle_state.sync_cb(client_id, value)
 
@@ -262,6 +275,25 @@ def add_gui_button(
         disabled: bool = False,
         visible: bool = True,
         hint: Optional[str] = None,
+        color: Optional[
+            Literal[
+                "dark",
+                "gray",
+                "red",
+                "pink",
+                "grape",
+                "violet",
+                "indigo",
+                "blue",
+                "cyan",
+                "green",
+                "lime",
+                "yellow",
+                "orange",
+                "teal",
+            ]
+        ] = None,
+        icon: Optional[Icon] = None,
     ) -> GuiButtonHandle:
         """Add a button to the GUI. The value of this input is set to `True` every time
         it is clicked; to detect clicks, we can manually set it back to `False`."""
@@ -278,6 +310,8 @@ def add_gui_button(
                     container_id=self._get_container_id(),
                     hint=hint,
                     initial_value=False,
+                    color=color,
+                    icon_base64=None if icon is None else base64_from_icon(icon),
                 ),
                 disabled=disabled,
                 visible=visible,

viser/_gui_handles.py

@@ -36,6 +36,7 @@
 
 if TYPE_CHECKING:
     from ._gui_api import GuiApi
+    from ._viser import ClientHandle
 
 
 T = TypeVar("T")
@@ -138,7 +139,13 @@ def value(self, value: Union[T, onp.ndarray]) -> None:
             # Pushing callbacks into separate threads helps prevent deadlocks when we
             # have a lock in a callback. TODO: revisit other callbacks.
             threading.Thread(
-                target=lambda: cb(GuiEvent(client_id=None, target=self))
+                target=lambda: cb(
+                    GuiEvent(
+                        client_id=None,
+                        client=None,
+                        target=self,
+                    )
+                )
             ).start()
 
     @property
@@ -223,6 +230,7 @@ class GuiEvent(Generic[TGuiHandle]):
     Passed as input to callback functions."""
 
     client_id: Optional[ClientId]
+    client: Optional[ClientHandle]
     target: TGuiHandle
 
 

viser/_message_api.py

@@ -568,6 +568,16 @@ def add_3d_gui_container(
         # Avoids circular import.
         from ._gui_api import GuiApi, _make_unique_id
 
+        # New name to make the type checker happy; ViserServer and ClientHandle inherit
+        # from both GuiApi and MessageApi. The pattern below is unideal.
+        gui_api = self
+        assert isinstance(gui_api, GuiApi)
+
+        # Remove the 3D GUI container if it already exists. This will make sure
+        # contained GUI elements are removed, preventing potential memory leaks.
+        if name in gui_api._handle_from_node_name:
+            gui_api._handle_from_node_name[name].remove()
+
         container_id = _make_unique_id()
         self._queue(
             _messages.Gui3DMessage(
@@ -576,7 +586,5 @@ def add_3d_gui_container(
                 container_id=container_id,
             )
         )
-        assert isinstance(self, MessageApi)
         node_handle = SceneNodeHandle._make(self, name, wxyz, position)
-        assert isinstance(self, GuiApi)
-        return Gui3dContainerHandle(node_handle._impl, self, container_id)
+        return Gui3dContainerHandle(node_handle._impl, gui_api, container_id)

viser/_messages.py

@@ -333,6 +333,25 @@ class GuiAddButtonMessage(_GuiAddInputBase):
     # All GUI elements currently need an `initial_value` field.
     # This makes our job on the frontend easier.
     initial_value: bool
+    color: Optional[
+        Literal[
+            "dark",
+            "gray",
+            "red",
+            "pink",
+            "grape",
+            "violet",
+            "indigo",
+            "blue",
+            "cyan",
+            "green",
+            "lime",
+            "yellow",
+            "orange",
+            "teal",
+        ]
+    ]
+    icon_base64: Optional[str]
 
 
 @dataclasses.dataclass

viser/_scene_handles.py

@@ -232,9 +232,10 @@ class Gui3dContainerHandle(SceneNodeHandle):
         default_factory=dict
     )
 
-    def __enter__(self) -> None:
+    def __enter__(self) -> Gui3dContainerHandle:
         self._container_id_restore = self._gui_api._get_container_id()
         self._gui_api._set_container_id(self._container_id)
+        return self
 
     def __exit__(self, *args) -> None:
         del args

viser/client/src/App.tsx

@@ -22,7 +22,7 @@ import { SceneNodeThreeObject, UseSceneTree } from "./SceneTree";
 import "./index.css";
 
 import ControlPanel from "./ControlPanel/ControlPanel";
-import { UseGui, useGuiState } from "./ControlPanel/GuiState";
+import { UseGui, useGuiState, useMantineTheme } from "./ControlPanel/GuiState";
 import { searchParamKey } from "./SearchParamsUtils";
 import {
   WebsocketMessageProducer,
@@ -105,23 +105,11 @@ function ViewerRoot() {
 function ViewerContents() {
   const viewer = React.useContext(ViewerContext)!;
   const control_layout = viewer.useGui((state) => state.theme.control_layout);
-  const colors = viewer.useGui((state) => state.theme.colors);
   return (
     <MantineProvider
       withGlobalStyles
       withNormalizeCSS
-      theme={{
-        colorScheme: viewer.useGui((state) => state.theme.dark_mode)
-          ? "dark"
-          : "light",
-        primaryColor: colors === null ? undefined : "custom",
-        colors:
-          colors === null
-            ? undefined
-            : {
-                custom: colors,
-              },
-      }}
+      theme={useMantineTheme()}
     >
       <Titlebar />
       <ViserModal />

viser/client/src/ControlPanel/Generated.tsx

@@ -27,6 +27,7 @@ import { ErrorBoundary } from "react-error-boundary";
 
 /** Root of generated inputs. */
 export default function GeneratedGuiContainer({
+  // We need to take viewer as input in drei's <Html /> elements, where contexts break.
   containerId,
   viewer,
 }: {
@@ -116,7 +117,7 @@ function GeneratedInput({
         <Button
           id={conf.id}
           fullWidth
-          color={"0x000"}
+          color={conf.color ?? undefined}
           onClick={() =>
             messageSender({
               type: "GuiUpdateMessage",
@@ -127,6 +128,19 @@ function GeneratedInput({
           style={{ height: "1.875rem" }}
           disabled={disabled}
           size="sm"
+          leftIcon={
+            conf.icon_base64 === null ? undefined : (
+              <Image
+                /*^In Safari, both the icon's height and width need to be set, otherwise the icon is clipped.*/
+                height={"0.9rem"}
+                width={"0.9rem"}
+                sx={(theme) => ({
+                  filter: theme.colorScheme == "dark" ? "invert(1)" : undefined,
+                })}
+                src={"data:image/svg+xml;base64," + conf.icon_base64}
+              />
+            )
+          }
         >
           {conf.label}
         </Button>

viser/client/src/ControlPanel/GuiState.tsx

@@ -2,6 +2,8 @@ import * as Messages from "../WebsocketMessages";
 import React from "react";
 import { create } from "zustand";
 import { immer } from "zustand/middleware/immer";
+import { ViewerContext } from "../App";
+import { MantineThemeOverride } from "@mantine/core";
 
 export type GuiConfig =
   | Messages.GuiAddButtonMessage
@@ -137,5 +139,22 @@ export function useGuiState(initialServer: string) {
   )[0];
 }
 
+export function useMantineTheme(): MantineThemeOverride {
+  const viewer = React.useContext(ViewerContext)!;
+  const colors = viewer.useGui((state) => state.theme.colors);
+  return {
+    colorScheme: viewer.useGui((state) => state.theme.dark_mode)
+      ? "dark"
+      : "light",
+    primaryColor: colors === null ? undefined : "custom",
+    colors:
+      colors === null
+        ? undefined
+        : {
+            custom: colors,
+          },
+  };
+}
+
 /** Type corresponding to a zustand-style useGuiState hook. */
 export type UseGui = ReturnType<typeof useGuiState>;

viser/client/src/SceneTree.tsx

@@ -81,6 +81,7 @@ function SceneNodeThreeChildren(props: {
       {children.map((child_id) => {
         return <SceneNodeThreeObject key={child_id} name={child_id} />;
       })}
+      <SceneNodeLabel name={props.name} />
     </group>,
     props.parent,
   );
@@ -124,17 +125,11 @@ export function SceneNodeThreeObject(props: { name: string }) {
   const [obj, setRef] = React.useState<THREE.Object3D | null>(null);
 
   // Create object + children.
-  const objNode = React.useMemo(
-    () => makeObject && makeObject(setRef),
-    [setRef, makeObject],
-  );
-  const children = React.useMemo(
-    () =>
-      obj === null ? null : (
-        <SceneNodeThreeChildren name={props.name} parent={obj} />
-      ),
-    [props.name, obj],
-  );
+  const objNode = makeObject && makeObject(setRef);
+  const children =
+    obj === null ? null : (
+      <SceneNodeThreeChildren name={props.name} parent={obj} />
+    );
 
   // Update attributes on a per-frame basis. Currently does redundant work,
   // although this shouldn't be a bottleneck.
@@ -186,7 +181,9 @@ export function SceneNodeThreeObject(props: { name: string }) {
     return nodeAttributes?.visibility ?? false;
   }
 
-  if (clickable) {
+  if (objNode === undefined) {
+    return <>{children}</>;
+  } else if (clickable) {
     return (
       <>
         <group
@@ -222,15 +219,15 @@ export function SceneNodeThreeObject(props: { name: string }) {
         >
           <Select enabled={hovered}>{objNode}</Select>
         </group>
-        <SceneNodeLabel name={props.name} />
         {children}
       </>
     );
   } else {
     return (
       <>
-        {objNode}
-        <SceneNodeLabel name={props.name} />
+        {/* This <group /> does nothing, but switching between clickable vs not
+        causes strange transform behavior without it. */}
+        <group>{objNode}</group>
         {children}
       </>
     );