3D gui containers (#69)

* 3d gui handle created, in progress

* Tweaks / fixes

* Add example

* Revert share in example

* Revert autobuild comment

* Cleanup

* Formatting

* Docstring update

---------

Co-authored-by: Brent Yi <[email protected]>

examples/15_gui_in_scene.py

@@ -0,0 +1,115 @@
+# mypy: disable-error-code="assignment"
+#
+# Asymmetric properties are supported in Pyright, but not yet in mypy.
+# - https://github.com/python/mypy/issues/3004
+# - https://github.com/python/mypy/pull/11643
+"""3D GUI Elements
+
+`add_3d_gui_container()` allows standard GUI elements to be incorporated directly into a
+3D scene. In this example, we click on coordinate frames to show actions that can be
+performed on them.
+"""
+
+import time
+from typing import Optional
+
+import numpy as onp
+
+import viser
+import viser.transforms as tf
+
+server = viser.ViserServer()
+num_frames = 20
+
+
+@server.on_client_connect
+def _(client: viser.ClientHandle) -> None:
+    """For each client that connects, we create a set of random frames + a click handler for each frame.
+
+    When a frame is clicked, we display a 3D gui node.
+    """
+
+    rng = onp.random.default_rng(0)
+
+    displayed_3d_container: Optional[viser.Gui3dContainerHandle] = None
+    displayed_index: Optional[int] = None
+
+    def make_frame(i: int) -> None:
+        # Sample a random orientation + position.
+        wxyz = rng.normal(size=4)
+        wxyz /= onp.linalg.norm(wxyz)
+        position = rng.uniform(-3.0, 3.0, size=(3,))
+
+        # Create a coordinate frame and label.
+        frame = client.add_frame(f"/frame_{i}", wxyz=wxyz, position=position)
+
+        # Move the camera when we click a frame.
+        @frame.on_click
+        def _(_):
+            nonlocal displayed_3d_container
+            nonlocal displayed_index
+
+            # Close previously opened GUI.
+            if displayed_3d_container is not None:
+                displayed_3d_container.remove()
+
+            # Don't re-show the same GUI element.
+            if displayed_index == i:
+                return
+
+            displayed_index = i
+
+            displayed_3d_container = client.add_3d_gui_container(f"/frame_{i}/gui")
+            with displayed_3d_container:
+                go_to = client.add_gui_button("Go to")
+                randomize_orientation = client.add_gui_button("Randomize orientation")
+                close = client.add_gui_button("Close GUI")
+
+            @go_to.on_click
+            def _(_) -> None:
+                T_world_current = tf.SE3.from_rotation_and_translation(
+                    tf.SO3(client.camera.wxyz), client.camera.position
+                )
+                T_world_target = tf.SE3.from_rotation_and_translation(
+                    tf.SO3(frame.wxyz), frame.position
+                ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))
+
+                T_current_target = T_world_current.inverse() @ T_world_target
+
+                for j in range(20):
+                    T_world_set = T_world_current @ tf.SE3.exp(
+                        T_current_target.log() * j / 19.0
+                    )
+
+                    # Important bit: we atomically set both the orientation and the position
+                    # of the camera.
+                    with client.atomic():
+                        client.camera.wxyz = T_world_set.rotation().wxyz
+                        client.camera.position = T_world_set.translation()
+                    time.sleep(1.0 / 60.0)
+
+                # Mouse interactions should orbit around the frame origin.
+                client.camera.look_at = frame.position
+
+            @randomize_orientation.on_click
+            def _(_) -> None:
+                wxyz = rng.normal(size=4)
+                wxyz /= onp.linalg.norm(wxyz)
+                frame.wxyz = wxyz
+
+            @close.on_click
+            def _(_) -> None:
+                nonlocal displayed_3d_container
+                nonlocal displayed_index
+                if displayed_3d_container is None:
+                    return
+                displayed_3d_container.remove()
+                displayed_3d_container = None
+                displayed_index = None
+
+    for i in range(num_frames):
+        make_frame(i)
+
+
+while True:
+    time.sleep(1.0)

sync_message_defs.py

@@ -19,4 +19,4 @@
     print(f"Wrote to {target_path}")
 
     # Run prettier.
-    subprocess.run(args=["prettier", "-w", str(target_path)])
+    subprocess.run(args=["npx", "prettier", "-w", str(target_path)])

viser/__init__.py

@@ -8,6 +8,7 @@
 from ._icons_enum import Icon as Icon
 from ._scene_handles import CameraFrustumHandle as CameraFrustumHandle
 from ._scene_handles import FrameHandle as FrameHandle
+from ._scene_handles import Gui3dContainerHandle as Gui3dContainerHandle
 from ._scene_handles import ImageHandle as ImageHandle
 from ._scene_handles import LabelHandle as LabelHandle
 from ._scene_handles import MeshHandle as MeshHandle

viser/_message_api.py

@@ -27,6 +27,7 @@
 from ._scene_handles import (
     CameraFrustumHandle,
     FrameHandle,
+    Gui3dContainerHandle,
     ImageHandle,
     LabelHandle,
     MeshHandle,
@@ -521,3 +522,28 @@ def _handle_click_updates(
             return
         for cb in handle._impl.click_cb:
             cb(handle)
+
+    def add_3d_gui_container(
+        self,
+        name: str,
+        wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),
+        position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
+    ) -> Gui3dContainerHandle:
+        """Add a 3D gui container to the scene. The returned container handle can be
+        used as a context to place GUI elements into the 3D scene."""
+
+        # Avoids circular import.
+        from ._gui_api import GuiApi, _make_unique_id
+
+        container_id = _make_unique_id()
+        self._queue(
+            _messages.Gui3DMessage(
+                order=time.time(),
+                name=name,
+                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)

viser/_messages.py

@@ -86,6 +86,15 @@ class LabelMessage(Message):
     text: str
 
 
+@dataclasses.dataclass
+class Gui3DMessage(Message):
+    """Add a 3D gui element to the scene."""
+
+    order: float
+    name: str
+    container_id: str
+
+
 @dataclasses.dataclass
 class PointCloudMessage(Message):
     """Point cloud message.

viser/_scene_handles.py

@@ -13,6 +13,7 @@
 from . import _messages
 
 if TYPE_CHECKING:
+    from ._gui_api import GuiApi
     from ._message_api import ClientId, MessageApi
 
 
@@ -201,3 +202,33 @@ def on_update(
         """Attach a callback for when the gizmo is moved."""
         self._impl_aux.update_cb.append(func)
         return func
+
+
+@dataclasses.dataclass
+class Gui3dContainerHandle(SceneNodeHandle):
+    """Use as a context to place GUI elements into a 3D GUI container."""
+
+    _gui_api: GuiApi
+    _container_id: str
+    _container_id_restore: Optional[str] = None
+
+    def __enter__(self) -> None:
+        self._container_id_restore = self._gui_api._get_container_id()
+        self._gui_api._set_container_id(self._container_id)
+
+    def __exit__(self, *args) -> None:
+        del args
+        assert self._container_id_restore is not None
+        self._gui_api._set_container_id(self._container_id_restore)
+        self._container_id_restore = None
+
+    def remove(self) -> None:
+        """Permanently remove this GUI container from the visualizer."""
+
+        # Call scene node remove.
+        super().remove()
+
+        # Clean up contained GUI elements.
+        self._gui_api._get_api()._queue(
+            _messages.GuiRemoveContainerChildrenMessage(self._container_id)
+        )

viser/client/src/App.tsx

@@ -29,7 +29,7 @@ import WebsocketInterface from "./WebsocketInterface";
 import { Titlebar } from "./Titlebar";
 import { useSceneTreeState } from "./SceneTreeState";
 
-type ViewerContextContents = {
+export type ViewerContextContents = {
   useSceneTree: UseSceneTree;
   useGui: UseGui;
   websocketRef: React.MutableRefObject<WebSocket | null>;

viser/client/src/ControlPanel/Generated.tsx

@@ -2,7 +2,7 @@ import {
   GuiAddFolderMessage,
   GuiAddTabGroupMessage,
 } from "../WebsocketMessages";
-import { ViewerContext } from "../App";
+import { ViewerContext, ViewerContextContents } from "../App";
 import { makeThrottledMessageSender } from "../WebsocketFunctions";
 import { GuiConfig } from "./GuiState";
 import { Image, Tabs, TabsValue } from "@mantine/core";
@@ -26,12 +26,15 @@ import React from "react";
 /** Root of generated inputs. */
 export default function GeneratedGuiContainer({
   containerId,
+  viewer,
 }: {
   containerId: string;
+  viewer?: ViewerContextContents;
 }) {
-  const viewer = React.useContext(ViewerContext)!;
+  if (viewer === undefined) viewer = React.useContext(ViewerContext)!;
+
   const guiIdSet = viewer.useGui(
-    (state) => state.guiIdSetFromContainerId[containerId],
+    (state) => state.guiIdSetFromContainerId[containerId]
   );
   const guiConfigFromId = viewer.useGui((state) => state.guiConfigFromId);
 
@@ -44,7 +47,12 @@ export default function GeneratedGuiContainer({
           .sort((a, b) => a.order - b.order)
           .map((conf, index) => {
             return (
-              <GeneratedInput conf={conf} key={conf.id} first={index == 0} />
+              <GeneratedInput
+                conf={conf}
+                key={conf.id}
+                first={index == 0}
+                viewer={viewer}
+              />
             );
           })}
       </>
@@ -54,15 +62,24 @@ export default function GeneratedGuiContainer({
 }
 
 /** A single generated GUI element. */
-function GeneratedInput({ conf, first }: { conf: GuiConfig; first: boolean }) {
+function GeneratedInput({
+  conf,
+  first,
+  viewer,
+}: {
+  conf: GuiConfig;
+  first: boolean;
+  viewer?: ViewerContextContents;
+}) {
   // Handle nested containers.
   if (conf.type == "GuiAddFolderMessage")
     return <GeneratedFolder conf={conf} first={first} />;
   if (conf.type == "GuiAddTabGroupMessage")
     return <GeneratedTabGroup conf={conf} />;
 
   // Handle GUI input types.
-  const viewer = React.useContext(ViewerContext)!;
+  if (viewer === undefined) viewer = React.useContext(ViewerContext)!;
+
   const messageSender = makeThrottledMessageSender(viewer.websocketRef, 50);
   function updateValue(value: any) {
     setGuiValue(conf.id, value);
@@ -406,7 +423,7 @@ function VectorInput(
         precision: number;
         onChange: (value: number[]) => void;
         disabled: boolean;
-      },
+      }
 ) {
   return (
     <Flex justify="space-between" style={{ columnGap: "0.3rem" }}>

viser/client/src/WebsocketInterface.tsx

@@ -15,7 +15,8 @@ import { Html, PivotControls } from "@react-three/drei";
 import { isTexture, makeThrottledMessageSender } from "./WebsocketFunctions";
 import { isGuiConfig } from "./ControlPanel/GuiState";
 import { useFrame } from "@react-three/fiber";
-
+import GeneratedGuiContainer from "./ControlPanel/Generated";
+import { Box, Paper } from "@mantine/core";
 /** Float **/
 function threeColorBufferFromUint8Buffer(colors: ArrayBuffer) {
   return new THREE.Float32BufferAttribute(
@@ -452,6 +453,32 @@ function useMessageHandler() {
         );
         return;
       }
+      case "Gui3DMessage": {
+        addSceneNodeMakeParents(
+          new SceneNode<THREE.Group>(message.name, (ref) => {
+            // We wrap with <group /> because Html doesn't implement THREE.Object3D.
+            return (
+              <group ref={ref}>
+                <Html>
+                  <Paper
+                    sx={{
+                      width: "20em",
+                      fontSize: "0.8em",
+                    }}
+                    withBorder
+                  >
+                    <GeneratedGuiContainer
+                      containerId={message.container_id}
+                      viewer={viewer}
+                    />
+                  </Paper>
+                </Html>
+              </group>
+            );
+          })
+        );
+        return;
+      }
       // Add an image.
       case "ImageMessage": {
         // It's important that we load the texture outside of the node

viser/client/src/WebsocketMessages.tsx

@@ -34,6 +34,12 @@ export interface LabelMessage {
   name: string;
   text: string;
 }
+export interface Gui3DMessage {
+  type: "Gui3DMessage";
+  order: number;
+  name: string;
+  container_id: string;
+}
 export interface PointCloudMessage {
   type: "PointCloudMessage";
   name: string;
@@ -329,6 +335,7 @@ export type Message =
   | CameraFrustumMessage
   | FrameMessage
   | LabelMessage
+  | Gui3DMessage
   | PointCloudMessage
   | MeshMessage
   | TransformControlsMessage