Click callbacks for scene node handles (#46)

* initial work on adding click callbacks * WIP outline on hover * Improve performance * Aesthetic cleanup * Fix bug, add clicks for add_mesh_trimesh * WIP example * Automate `.clickable`, nits * Example nits * Fix mypy error, add matplotlib * Docs sync * Cleanup * Dependencies tweak * Remove foobar * Fix clicking for frustums, update examples * More specific handle types * Docs * Run ruff * Docs nit --------- Co-authored-by: Chung Min Kim <[email protected]>
nerfstudio-project · Jun 15, 2023 · 35011ea · 35011ea
1 parent f84db3a
commit 35011ea
Show file tree

Hide file tree

Showing 35 changed files with 1,104 additions and 615 deletions.
diff --git a/docs/source/_static/css/custom.css b/docs/source/_static/css/custom.css
@@ -1,7 +1,7 @@
 img.sidebar-logo {
-    width: 3em;
-    margin: 1em 0 0 0;
+  width: 3em;
+  margin: 1em 0 0 0;
 }
 .sidebar-brand-text {
-    display: none;
+  display: none;
 }
diff --git a/docs/source/core.md b/docs/source/core.md
diff --git a/docs/source/examples/05_camera_commands.rst b/docs/source/examples/05_camera_commands.rst
@@ -22,35 +22,31 @@ corresponding client automatically.
         import viser.transforms as tf
 
         server = viser.ViserServer()
-
-        num_buttons = 20
+        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 button for each frame.
+            """For each client that connects, we create a set of random frames + a click handler for each frame.
 
-            When a button is clicked, we move the camera to the corresponding frame.
+            When a frame is clicked, we move the camera to the corresponding frame.
             """
 
             rng = onp.random.default_rng(0)
 
-            def make_frame(i: int) -> viser.GuiButtonHandle:
+            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 a button to move to that frame.
+                # Create a coordinate frame and label.
                 frame = client.add_frame(f"/frame_{i}", wxyz=wxyz, position=position)
                 client.add_label(f"/frame_{i}/label", text=f"Frame {i}")
-                button = client.add_gui_button(f"Frame {i}")
 
-                @button.on_click
+                # Move the camera when we click a frame.
+                @frame.on_click
                 def _(_):
-                    for b in buttons:
-                        b.disabled = True
-
                     T_world_current = tf.SE3.from_rotation_and_translation(
                         tf.SO3(client.camera.wxyz), client.camera.position
                     )
@@ -75,14 +71,8 @@ corresponding client automatically.
                     # Mouse interactions should orbit around the frame origin.
                     client.camera.look_at = frame.position
 
-                    for b in buttons:
-                        b.disabled = False
-
-                return button
-
-            buttons = []
-            for i in range(num_buttons):
-                buttons.append(make_frame(i))
+            for i in range(num_frames):
+                make_frame(i)
 
 
         while True:

diff --git a/docs/source/examples/07_record3d_visualizer.rst b/docs/source/examples/07_record3d_visualizer.rst
@@ -22,6 +22,7 @@ Parse and stream record3d captures. To get the demo data, see ``./assets/downloa
         from tqdm.auto import tqdm
 
         import viser
+        import viser.extras
         import viser.transforms as tf
 
 
@@ -91,7 +92,7 @@ Parse and stream record3d captures. To get the demo data, see ``./assets/downloa
                 position=(0, 0, 0),
                 show_axes=False,
             )
-            frame_nodes: List[viser.SceneNodeHandle] = []
+            frame_nodes: List[viser.FrameHandle] = []
             for i in tqdm(range(num_frames)):
                 frame = loader.get_frame(i)
                 position, color = frame.get_point_cloud(downsample_factor)

diff --git a/docs/source/examples/09_urdf_visualizer.rst b/docs/source/examples/09_urdf_visualizer.rst
@@ -49,9 +49,14 @@ Examples:
 
             for frame_name, mesh in urdf.scene.geometry.items():
                 assert isinstance(mesh, trimesh.Trimesh)
+                T_parent_child = urdf.get_transform(
+                    frame_name, urdf.scene.graph.transforms.parents[frame_name]
+                )
                 server.add_mesh_trimesh(
-                    frame_name_with_parents(frame_name) + "/mesh",
+                    frame_name_with_parents(frame_name),
                     mesh,
+                    wxyz=tf.SO3.from_matrix(T_parent_child[:3, :3]).wxyz,
+                    position=T_parent_child[:3, 3],
                 )
 
             gui_joints: List[viser.GuiHandle[float]] = []

diff --git a/docs/source/examples/11_colmap_visualizer.rst b/docs/source/examples/11_colmap_visualizer.rst
@@ -98,6 +98,15 @@ Visualize COLMAP sparse reconstruction outputs. To get demo data, see ``./assets
                 onp.random.shuffle(img_ids)
                 img_ids = sorted(img_ids[: gui_frames.value])
 
+                def attach_callback(
+                    frustum: viser.CameraFrustumHandle, frame: viser.FrameHandle
+                ) -> None:
+                    @frustum.on_click
+                    def _(_) -> None:
+                        for client in server.get_clients().values():
+                            client.camera.wxyz = frame.wxyz
+                            client.camera.position = frame.position
+
                 for img_id in tqdm(img_ids):
                     img = images[img_id]
                     cam = cameras[img.camera_id]
@@ -110,7 +119,7 @@ Visualize COLMAP sparse reconstruction outputs. To get demo data, see ``./assets
                     T_world_camera = tf.SE3.from_rotation_and_translation(
                         tf.SO3(img.qvec), img.tvec
                     ).inverse()
-                    server.add_frame(
+                    frame = server.add_frame(
                         f"/colmap/frame_{img_id}",
                         wxyz=T_world_camera.rotation().wxyz,
                         position=T_world_camera.translation(),
@@ -126,13 +135,14 @@ Visualize COLMAP sparse reconstruction outputs. To get demo data, see ``./assets
                     fy = cam.params[1]
                     image = iio.imread(image_filename)
                     image = image[::downsample_factor, ::downsample_factor]
-                    server.add_camera_frustum(
+                    frustum = server.add_camera_frustum(
                         f"/colmap/frame_{img_id}/frustum",
                         fov=2 * onp.arctan2(H / 2, fy),
                         aspect=W / H,
                         scale=0.15,
                         image=image,
                     )
+                    attach_callback(frustum, frame)
 
             need_update = True
 

diff --git a/docs/source/examples/12_click_meshes.rst b/docs/source/examples/12_click_meshes.rst
@@ -0,0 +1,96 @@
+.. Comment: this file is automatically generated by `update_example_docs.py`.
+   It should not be modified manually.
+
+Click callback demonstration.
+==========================================
+
+
+Click on meshes to select them. The index of the last clicked mesh is displayed in the GUI.
+
+
+
+.. code-block:: python
+        :linenos:
+
+
+        import time
+
+        import matplotlib
+        import numpy as onp
+        import trimesh.creation
+
+        import viser
+
+
+        def main() -> None:
+            grid_shape = (4, 5)
+            server = viser.ViserServer()
+
+            with server.gui_folder("Last clicked"):
+                x_value = server.add_gui_number(
+                    name="x",
+                    initial_value=0,
+                    disabled=True,
+                    hint="x coordinate of the last clicked mesh",
+                )
+                y_value = server.add_gui_number(
+                    name="y",
+                    initial_value=0,
+                    disabled=True,
+                    hint="y coordinate of the last clicked mesh",
+                )
+
+            def add_swappable_mesh(i: int, j: int) -> None:
+                """Simple callback that swaps between:
+                 - a gray box
+                 - a colored box
+                 - a colored sphere
+
+                Color is chosen based on the position (i, j) of the mesh in the grid.
+                """
+
+                colormap = matplotlib.colormaps["tab20"]
+
+                def create_mesh(counter: int) -> None:
+                    if counter == 0:
+                        mesh = trimesh.creation.box((0.5, 0.5, 0.5))
+                    elif counter == 1:
+                        mesh = trimesh.creation.box((0.5, 0.5, 0.5))
+                    else:
+                        mesh = trimesh.creation.icosphere(subdivisions=2, radius=0.4)
+
+                    colors = colormap(
+                        (i * grid_shape[1] + j + onp.random.rand(mesh.vertices.shape[0]))
+                        / (grid_shape[0] * grid_shape[1])
+                    )
+                    if counter != 0:
+                        assert mesh.visual is not None
+                        mesh.visual.vertex_colors = colors
+
+                    handle = server.add_mesh_trimesh(
+                        name=f"/sphere_{i}_{j}",
+                        mesh=mesh,
+                        position=(i, j, 0.0),
+                    )
+
+                    @handle.on_click
+                    def _(_) -> None:
+                        x_value.value = i
+                        y_value.value = j
+
+                        # The new mesh will replace the old one because the names (/sphere_{i}_{j}) are
+                        # the same.
+                        create_mesh((counter + 1) % 3)
+
+                create_mesh(0)
+
+            for i in range(grid_shape[0]):
+                for j in range(grid_shape[1]):
+                    add_swappable_mesh(i, j)
+
+            while True:
+                time.sleep(10.0)
+
+
+        if __name__ == "__main__":
+            main()
diff --git a/docs/source/index.md b/docs/source/index.md
@@ -57,12 +57,23 @@ URL (default: `http://localhost:8080`).
    ./development.md
 
 .. toctree::
-   :caption: API Reference
+   :caption: API (Core)
+   :hidden:
+   :maxdepth: 1
+   :titlesonly:
+
+   ./server.md
+   ./client_handles.md
+   ./gui_handles.md
+   ./scene_handles.md
+
+
+.. toctree::
+   :caption: API (Additional)
    :hidden:
    :maxdepth: 1
    :titlesonly:
 
-   ./core.md
    ./infrastructure.md
    ./transforms.md
 

diff --git a/docs/source/infrastructure.md b/docs/source/infrastructure.md
@@ -1,4 +1,4 @@
-# Infrastructure
+# Communication
 
 <!-- prettier-ignore-start -->
 

diff --git a/examples/05_camera_commands.py b/examples/05_camera_commands.py
@@ -17,35 +17,31 @@
 import viser.transforms as tf
 
 server = viser.ViserServer()
-
-num_buttons = 20
+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 button for each frame.
+    """For each client that connects, we create a set of random frames + a click handler for each frame.
 
-    When a button is clicked, we move the camera to the corresponding frame.
+    When a frame is clicked, we move the camera to the corresponding frame.
     """
 
     rng = onp.random.default_rng(0)
 
-    def make_frame(i: int) -> viser.GuiButtonHandle:
+    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 a button to move to that frame.
+        # Create a coordinate frame and label.
         frame = client.add_frame(f"/frame_{i}", wxyz=wxyz, position=position)
         client.add_label(f"/frame_{i}/label", text=f"Frame {i}")
-        button = client.add_gui_button(f"Frame {i}")
 
-        @button.on_click
+        # Move the camera when we click a frame.
+        @frame.on_click
         def _(_):
-            for b in buttons:
-                b.disabled = True
-
             T_world_current = tf.SE3.from_rotation_and_translation(
                 tf.SO3(client.camera.wxyz), client.camera.position
             )
@@ -70,14 +66,8 @@ def _(_):
             # Mouse interactions should orbit around the frame origin.
             client.camera.look_at = frame.position
 
-            for b in buttons:
-                b.disabled = False
-
-        return button
-
-    buttons = []
-    for i in range(num_buttons):
-        buttons.append(make_frame(i))
+    for i in range(num_frames):
+        make_frame(i)
 
 
 while True: