Websocket refactor + theming (#82)

* Refactor websocket interface, theming improvements

* Clean up control panel, icons

* CSS, example cleanup

* Fix mypy errors

* Match 3D gui style

* Add E741 to ruff ignore

examples/08_smplx_visualizer.py

@@ -36,6 +36,7 @@ def main(
     ext: Literal["npz", "pkl"] = "npz",
 ) -> None:
     server = viser.ViserServer()
+    server.configure_theme(control_layout="collapsible", dark_mode=True)
     model = smplx.create(
         model_path=str(model_path),
         model_type=model_type,

examples/13_theming.py

@@ -1,6 +1,9 @@
+# mypy: disable-error-code="arg-type"
+#
+# Waiting on PEP 675 support in mypy. https://github.com/python/mypy/issues/12554
 """Theming
 
-Viser is adding support for theming. Work-in-progress.
+Viser includes support for light theming.
 """
 
 import time
@@ -33,15 +36,34 @@
     image_alt="NerfStudio Logo",
     href="https://docs.nerf.studio/",
 )
-
-# image = None
-
 titlebar_theme = TitlebarConfig(buttons=buttons, image=image)
 
-server.configure_theme(
-    dark_mode=True, titlebar_content=titlebar_theme, control_layout="fixed"
+server.add_gui_markdown(
+    "Viser includes support for light theming via the `.configure_theme()` method."
+)
+
+# GUI elements for controllable values.
+titlebar = server.add_gui_checkbox("Titlebar", initial_value=True)
+dark_mode = server.add_gui_checkbox("Dark mode", initial_value=True)
+control_layout = server.add_gui_dropdown(
+    "Control layout", ("floating", "fixed", "collapsible")
 )
-server.world_axes.visible = True
+brand_color = server.add_gui_rgb("Brand color", (230, 180, 30))
+synchronize = server.add_gui_button("Apply theme")
+
+
+@synchronize.on_click
+def synchronize_theme(_) -> None:
+    server.configure_theme(
+        dark_mode=dark_mode.value,
+        titlebar_content=titlebar_theme if titlebar.value else None,
+        control_layout=control_layout.value,
+        brand_color=brand_color.value,
+    )
+    server.world_axes.visible = True
+
+
+synchronize_theme(synchronize)
 
 while True:
     time.sleep(10.0)

examples/16_modal.py

@@ -14,7 +14,9 @@ def main():
     def _(client: viser.ClientHandle) -> None:
         with client.add_gui_modal("Modal example"):
             client.add_gui_markdown(
-                markdown="**The slider below determines how many modals will appear...**"
+                markdown=(
+                    "**The slider below determines how many modals will appear...**"
+                )
             )
 
             gui_slider = client.add_gui_slider(

pyproject.toml

@@ -91,6 +91,7 @@ select = [
     "PLW",  # Pylint warnings.
 ]
 ignore = [
+    "E741", # Ambiguous variable name. (l, O, or I)
     "E501",  # Line too long.
     "F722",  # Forward annotation false positive from jaxtyping. Should be caught by pyright.
     "F821",  # Forward annotation false positive from jaxtyping. Should be caught by pyright.

viser/_message_api.py

@@ -9,6 +9,7 @@
 
 import abc
 import base64
+import colorsys
 import io
 import threading
 import time
@@ -45,6 +46,16 @@
 P = ParamSpec("P")
 
 
+def _hex_from_hls(h: float, l: float, s: float) -> str:
+    """Converts HLS values in [0.0, 1.0] to a hex-formatted string, eg 0xffffff."""
+    return "#" + "".join(
+        [
+            int(min(255, max(0, channel * 255.0)) + 0.5).to_bytes(1, "little").hex()
+            for channel in colorsys.hls_to_rgb(h, l, s)
+        ]
+    )
+
+
 def _colors_to_uint8(colors: onp.ndarray) -> onpt.NDArray[onp.uint8]:
     """Convert intensity values to uint8. We assume the range [0,1] for floats, and
     [0,255] for integers."""
@@ -143,13 +154,51 @@ def configure_theme(
         titlebar_content: Optional[theme.TitlebarConfig] = None,
         control_layout: Literal["floating", "collapsible", "fixed"] = "floating",
         dark_mode: bool = False,
+        brand_color: Optional[Tuple[int, int, int]] = None,
     ) -> None:
         """Configure the viser front-end's visual appearance."""
+
+        colors_cast: Optional[
+            Tuple[str, str, str, str, str, str, str, str, str, str]
+        ] = None
+
+        if brand_color is not None:
+            assert len(brand_color) in (3, 10)
+            if len(brand_color) == 3:
+                assert all(
+                    map(lambda val: isinstance(val, int), brand_color)
+                ), "All channels should be integers."
+
+                # RGB => HLS.
+                h, l, s = colorsys.rgb_to_hls(
+                    brand_color[0] / 255.0,
+                    brand_color[1] / 255.0,
+                    brand_color[2] / 255.0,
+                )
+
+                # Automatically generate a 10-color palette.
+                min_l = max(l - 0.08, 0.0)
+                max_l = min(0.8 + 0.5, 0.9)
+                l = max(min_l, min(max_l, l))
+
+                primary_index = 8
+                ls = tuple(
+                    onp.interp(
+                        x=onp.arange(10), xp=(0, primary_index, 9), fp=(max_l, l, min_l)
+                    )
+                )
+                colors_cast = tuple(_hex_from_hls(h, ls[i], s) for i in range(10))  # type: ignore
+
+        assert colors_cast is None or all(
+            [isinstance(val, str) and val.startswith("#") for val in colors_cast]
+        ), "All string colors should be in hexadecimal + prefixed with #, eg #ffffff."
+
         self._queue(
             _messages.ThemeConfigurationMessage(
                 titlebar_content=titlebar_content,
                 control_layout=control_layout,
                 dark_mode=dark_mode,
+                colors=colors_cast,
             ),
         )
 

viser/_messages.py

@@ -451,4 +451,5 @@ class ThemeConfigurationMessage(Message):
 
     titlebar_content: Optional[theme.TitlebarConfig]
     control_layout: Literal["floating", "collapsible", "fixed"]
+    colors: Optional[Tuple[str, str, str, str, str, str, str, str, str, str]]
     dark_mode: bool

viser/client/src/App.tsx

@@ -24,20 +24,28 @@ import "./index.css";
 import ControlPanel from "./ControlPanel/ControlPanel";
 import { UseGui, useGuiState } from "./ControlPanel/GuiState";
 import { searchParamKey } from "./SearchParamsUtils";
-import WebsocketInterface from "./WebsocketInterface";
+import {
+  WebsocketMessageProducer,
+  FrameSynchronizedMessageHandler,
+} from "./WebsocketInterface";
 
 import { Titlebar } from "./Titlebar";
 import { ViserModal } from "./Modal";
 import { useSceneTreeState } from "./SceneTreeState";
+import { Message } from "./WebsocketMessages";
 
 export type ViewerContextContents = {
+  // Zustand hooks.
   useSceneTree: UseSceneTree;
   useGui: UseGui;
+  // Useful references.
   websocketRef: React.MutableRefObject<WebSocket | null>;
   canvasRef: React.MutableRefObject<HTMLCanvasElement | null>;
   sceneRef: React.MutableRefObject<THREE.Scene | null>;
   cameraRef: React.MutableRefObject<THREE.PerspectiveCamera | null>;
   cameraControlRef: React.MutableRefObject<CameraControls | null>;
+  // Scene node attributes.
+  // This is intentionally placed outside of the Zustand state to reduce overhead.
   nodeAttributesFromName: React.MutableRefObject<{
     [name: string]:
       | undefined
@@ -47,15 +55,16 @@ export type ViewerContextContents = {
           visibility?: boolean;
         };
   }>;
+  messageQueueRef: React.MutableRefObject<Message[]>;
 };
 export const ViewerContext = React.createContext<null | ViewerContextContents>(
   null,
 );
 
 THREE.ColorManagement.enabled = true;
 
-function SingleViewer() {
-  // Default server logic.
+function ViewerRoot() {
+  // What websocket server should we connect to?
   function getDefaultServerFromUrl() {
     // https://localhost:8080/ => ws://localhost:8080
     // https://localhost:8080/?server=some_url => ws://localhost:8080
@@ -80,18 +89,23 @@ function SingleViewer() {
     sceneRef: React.useRef(null),
     cameraRef: React.useRef(null),
     cameraControlRef: React.useRef(null),
-    // Scene node attributes that aren't placed in the zustand state, for performance reasons.
+    // Scene node attributes that aren't placed in the zustand state for performance reasons.
     nodeAttributesFromName: React.useRef({}),
+    messageQueueRef: React.useRef([]),
   };
 
-  // Memoize the websocket interface so it isn't remounted when the theme or
-  // viewer context changes.
-  const memoizedWebsocketInterface = React.useMemo(
-    () => <WebsocketInterface />,
-    [],
+  return (
+    <ViewerContext.Provider value={viewer}>
+      <WebsocketMessageProducer />
+      <ViewerContents />
+    </ViewerContext.Provider>
   );
+}
 
+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
@@ -100,37 +114,43 @@ function SingleViewer() {
         colorScheme: viewer.useGui((state) => state.theme.dark_mode)
           ? "dark"
           : "light",
+        primaryColor: colors === null ? undefined : "custom",
+        colors:
+          colors === null
+            ? undefined
+            : {
+                custom: colors,
+              },
       }}
     >
-      <ViewerContext.Provider value={viewer}>
-        <Titlebar />
-        <ViserModal />
-        <Box
-          sx={{
-            width: "100%",
-            height: "1px",
-            position: "relative",
-            flex: "1 0 auto",
-          }}
-        >
-          <MediaQuery smallerThan={"xs"} styles={{ right: 0, bottom: "3.5em" }}>
-            <Box
-              sx={(theme) => ({
-                top: 0,
-                bottom: 0,
-                left: 0,
-                right: control_layout === "fixed" ? "20em" : 0,
-                position: "absolute",
-                backgroundColor:
-                  theme.colorScheme === "light" ? "#fff" : theme.colors.dark[9],
-              })}
-            >
-              <ViewerCanvas>{memoizedWebsocketInterface}</ViewerCanvas>
-            </Box>
-          </MediaQuery>
-          <ControlPanel control_layout={control_layout} />
-        </Box>
-      </ViewerContext.Provider>
+      <Titlebar />
+      <ViserModal />
+      <Box
+        sx={{
+          width: "100%",
+          height: "1px",
+          position: "relative",
+          flexGrow: 1,
+          display: "flex",
+          flexDirection: "row",
+        }}
+      >
+        <MediaQuery smallerThan={"xs"} styles={{ right: 0, bottom: "3.5em" }}>
+          <Box
+            sx={(theme) => ({
+              backgroundColor:
+                theme.colorScheme === "light" ? "#fff" : theme.colors.dark[9],
+              flexGrow: 1,
+              width: "10em",
+            })}
+          >
+            <ViewerCanvas>
+              <FrameSynchronizedMessageHandler />
+            </ViewerCanvas>
+          </Box>
+        </MediaQuery>
+        <ControlPanel control_layout={control_layout} />
+      </Box>
     </MantineProvider>
   );
 }
@@ -195,7 +215,7 @@ export function Root() {
         flexDirection: "column",
       }}
     >
-      <SingleViewer />
+      <ViewerRoot />
     </Box>
   );
 }