Refactor scene pointer events + support rectangular selection (nerfst…

…udio-project#157) * Initial commit for 3D scribble, WIP * Add Viewer2DCanvas, use that to draw cursor feedback! * Remove scribble support, smoothen out box select, add example * Correctly update canvas size on resize * Refactor ScenePointerMessage, send 2D coordinates for box * Update example, walk through OpenCV/NDC coordinates * Rename SceneClickEnableMessage->ScenePointerEnableMessage * ruff, mypy * Put camera in a reasonable initial location * Add event type to on_scene_pointer * Rename example file from scene_click->scene_pointer * mypy, ruff * Add backwards compatibility * ruff, mypy * ruff, mypy * Rename click->pointer * Use resizeobserver for element size tracking * Use OpenCV image coords (normalized) * mypy * Remove irrelevant file * mypy * ruff, pyright * mypy * Nits * Try to fix mypy * Rename box to rect-select * Listen to only one scenepointerevent at a given time. * mypy, ruff * ruff * Track events for clients! * Clear pre-existing callbacks for both server/clients, include warnings. * ruff * Code nit * Nits * oops * Move pointer cleanup logic --------- Co-authored-by: Brent Yi <[email protected]>
yzslab · Mar 23, 2024 · 35331a8 · 35331a8
1 parent 752dcd7
commit 35331a8
Show file tree

Hide file tree

Showing 13 changed files with 753 additions and 278 deletions.
diff --git a/docs/source/examples/02_gui.rst b/docs/source/examples/02_gui.rst
@@ -83,6 +83,15 @@ Examples of basic GUI elements that we can create, read from, and write to.
                         initial_value=3,
                         marks=((0, "0"), (5, "5"), (7, "7"), 10),
                     )
+                    gui_upload_button = server.add_gui_upload_button(
+                        "Upload", icon=viser.Icon.UPLOAD
+                    )
+
+            @gui_upload_button.on_upload
+            def _(_) -> None:
+                """Callback for when a file is uploaded."""
+                file = gui_upload_button.value
+                print(file.name, len(file.content), "bytes")
 
             # Pre-generate a point cloud to send.
             point_positions = onp.random.uniform(low=-1.0, high=1.0, size=(5000, 3))

diff --git a/docs/source/examples/20_scene_click.rst b/docs/source/examples/20_scene_click.rst
diff --git a/docs/source/examples/20_scene_pointer.rst b/docs/source/examples/20_scene_pointer.rst
@@ -0,0 +1,168 @@
+.. Comment: this file is automatically generated by `update_example_docs.py`.
+   It should not be modified manually.
+
+Scene pointer events.
+==========================================
+
+
+This example shows how to use scene pointer events to specify rays, and how they can be
+used to interact with the scene (e.g., ray-mesh intersections).
+
+To get the demo data, see ``./assets/download_dragon_mesh.sh``.
+
+
+
+.. code-block:: python
+        :linenos:
+
+
+        import time
+        from pathlib import Path
+        from typing import List, cast
+
+        import numpy as onp
+        import trimesh
+        import trimesh.creation
+        import trimesh.ray
+        import viser
+        import viser.transforms as tf
+
+        server = viser.ViserServer()
+        server.configure_theme(brand_color=(130, 0, 150))
+        server.set_up_direction("+y")
+
+        mesh = cast(
+            trimesh.Trimesh, trimesh.load_mesh(str(Path(__file__).parent / "assets/dragon.obj"))
+        )
+        mesh.apply_scale(0.05)
+
+        mesh_handle = server.add_mesh_trimesh(
+            name="/mesh",
+            mesh=mesh,
+            position=(0.0, 0.0, 0.0),
+        )
+
+        hit_pos_handles: List[viser.GlbHandle] = []
+
+
+        # Buttons + callbacks will operate on a per-client basis, but will modify the global scene! :)
+        @server.on_client_connect
+        def _(client: viser.ClientHandle) -> None:
+            # Set up the camera -- this gives a nice view of the full mesh.
+            client.camera.position = onp.array([0.0, 0.0, -10.0])
+            client.camera.wxyz = onp.array([0.0, 0.0, 0.0, 1.0])
+
+            # Tests "click" scenepointerevent.
+            click_button_handle = client.add_gui_button("Add sphere", icon=viser.Icon.POINTER)
+
+            @click_button_handle.on_click
+            def _(_):
+                click_button_handle.disabled = True
+
+                @client.on_scene_pointer(event_type="click")
+                def _(event: viser.ScenePointerEvent) -> None:
+                    # Check for intersection with the mesh, using trimesh's ray-mesh intersection.
+                    # Note that mesh is in the mesh frame, so we need to transform the ray.
+                    R_world_mesh = tf.SO3(mesh_handle.wxyz)
+                    R_mesh_world = R_world_mesh.inverse()
+                    origin = (R_mesh_world @ onp.array(event.ray_origin)).reshape(1, 3)
+                    direction = (R_mesh_world @ onp.array(event.ray_direction)).reshape(1, 3)
+                    intersector = trimesh.ray.ray_triangle.RayMeshIntersector(mesh)
+                    hit_pos, _, _ = intersector.intersects_location(origin, direction)
+
+                    if len(hit_pos) == 0:
+                        return
+
+                    # Get the first hit position (based on distance from the ray origin).
+                    hit_pos = min(hit_pos, key=lambda x: onp.linalg.norm(x - origin))
+
+                    # Create a sphere at the hit location.
+                    hit_pos_mesh = trimesh.creation.icosphere(radius=0.1)
+                    hit_pos_mesh.vertices += R_world_mesh @ hit_pos
+                    hit_pos_mesh.visual.vertex_colors = (0.5, 0.0, 0.7, 1.0)  # type: ignore
+                    hit_pos_handle = server.add_mesh_trimesh(
+                        name=f"/hit_pos_{len(hit_pos_handles)}", mesh=hit_pos_mesh
+                    )
+                    hit_pos_handles.append(hit_pos_handle)
+
+                @client.on_scene_pointer_done
+                def _():
+                    click_button_handle.disabled = False
+                    client.remove_scene_pointer_callback()
+
+            # Tests "rect-select" scenepointerevent.
+            paint_button_handle = client.add_gui_button("Paint mesh", icon=viser.Icon.PAINT)
+
+            @paint_button_handle.on_click
+            def _(_):
+                paint_button_handle.disabled = True
+
+                @client.on_scene_pointer(event_type="rect-select")
+                def _(message: viser.ScenePointerEvent) -> None:
+                    global mesh_handle
+                    camera = message.client.camera
+
+                    # Put the mesh in the camera frame.
+                    R_world_mesh = tf.SO3(mesh_handle.wxyz)
+                    R_mesh_world = R_world_mesh.inverse()
+                    R_camera_world = tf.SE3.from_rotation_and_translation(
+                        tf.SO3(camera.wxyz), camera.position
+                    ).inverse()
+                    vertices = mesh.vertices
+                    vertices = (R_mesh_world.as_matrix() @ vertices.T).T
+                    vertices = (
+                        R_camera_world.as_matrix()
+                        @ onp.hstack([vertices, onp.ones((vertices.shape[0], 1))]).T
+                    ).T[:, :3]
+
+                    # Get the camera intrinsics, and project the vertices onto the image plane.
+                    fov, aspect = camera.fov, camera.aspect
+                    vertices_proj = vertices[:, :2] / vertices[:, 2].reshape(-1, 1)
+                    vertices_proj /= onp.tan(fov / 2)
+                    vertices_proj[:, 0] /= aspect
+
+                    # Move the origin to the upper-left corner, and scale to [0, 1].
+                    # ... make sure to match the OpenCV's image coordinates!
+                    vertices_proj = (1 + vertices_proj) / 2
+
+                    # Select the vertices that lie inside the 2D selected box, once projected.
+                    mask = (
+                        (vertices_proj > onp.array(message.screen_pos[0]))
+                        & (vertices_proj < onp.array(message.screen_pos[1]))
+                    ).all(axis=1)[..., None]
+
+                    # Update the mesh color based on whether the vertices are inside the box
+                    mesh.visual.vertex_colors = onp.where(  # type: ignore
+                        mask, (0.5, 0.0, 0.7, 1.0), (0.9, 0.9, 0.9, 1.0)
+                    )
+                    mesh_handle = server.add_mesh_trimesh(
+                        name="/mesh",
+                        mesh=mesh,
+                        position=(0.0, 0.0, 0.0),
+                    )
+
+                @client.on_scene_pointer_done
+                def _():
+                    paint_button_handle.disabled = False
+                    client.remove_scene_pointer_callback()
+
+            # Button to clear spheres.
+            clear_button_handle = client.add_gui_button("Clear scene", icon=viser.Icon.X)
+
+            @clear_button_handle.on_click
+            def _(_):
+                """Reset the mesh color and remove all click-generated spheres."""
+                global mesh_handle
+                for handle in hit_pos_handles:
+                    handle.remove()
+                hit_pos_handles.clear()
+                mesh.visual.vertex_colors = (0.9, 0.9, 0.9, 1.0)  # type: ignore
+                mesh_handle = server.add_mesh_trimesh(
+                    name="/mesh",
+                    mesh=mesh,
+                    position=(0.0, 0.0, 0.0),
+                )
+
+
+        while True:
+            time.sleep(10.0)
diff --git a/examples/20_scene_click.py b/examples/20_scene_click.py