Implement gui.add_image() + plotting example (#343)

* Implement `gui.add_image()` + plotting example

* ruff

* UnionType import fix

* Add Image.tsx

* Add title to plots, image prop for scene nodes

* Props

docs/source/examples/24_plots_as_images.rst

@@ -0,0 +1,283 @@
+.. Comment: this file is automatically generated by `update_example_docs.py`.
+   It should not be modified manually.
+
+Plots as Images
+==========================================
+
+
+Examples of sending plots as images to Viser's GUI panel. This can be faster
+than using Plotly.
+
+
+
+.. code-block:: python
+        :linenos:
+
+
+        import colorsys
+        import time
+
+        import cv2
+        import numpy as np
+        import tyro
+
+        import viser
+        import viser.transforms as vtf
+
+
+        def get_line_plot(
+            xs: np.ndarray,
+            ys: np.ndarray,
+            height: int,
+            width: int,
+            *,
+            x_bounds: tuple[float, float] | None = None,
+            y_bounds: tuple[float, float] | None = None,
+            title: str | None = None,
+            line_thickness: int = 2,
+            grid_x_lines: int = 8,
+            grid_y_lines: int = 5,
+            font_scale: float = 0.4,
+            background_color: tuple[int, int, int] = (0, 0, 0),
+            plot_area_color: tuple[int, int, int] = (0, 0, 0),
+            grid_color: tuple[int, int, int] = (60, 60, 60),
+            axes_color: tuple[int, int, int] = (100, 100, 100),
+            line_color: tuple[int, int, int] = (255, 255, 255),
+            text_color: tuple[int, int, int] = (200, 200, 200),
+        ) -> np.ndarray:
+            """Create a line plot using OpenCV with axes, labels, and grid.
+
+            This is much faster than using libraries like Matplotlib or Plotly, but is
+            less flexible.
+            """
+
+            if x_bounds is None:
+                x_bounds = (np.min(xs), np.max(xs.round(decimals=4)))
+            if y_bounds is None:
+                y_bounds = (np.min(ys), np.max(ys))
+
+            # Calculate text sizes for padding.
+            font = cv2.FONT_HERSHEY_DUPLEX
+            sample_y_label = f"{max(abs(y_bounds[0]), abs(y_bounds[1])):.1f}"
+            y_text_size = cv2.getTextSize(sample_y_label, font, font_scale, 1)[0]
+
+            sample_x_label = f"{max(abs(x_bounds[0]), abs(x_bounds[1])):.1f}"
+            x_text_size = cv2.getTextSize(sample_x_label, font, font_scale, 1)[0]
+
+            # Define padding based on font scale.
+            extra_padding = 8
+            left_pad = int(y_text_size[0] * 1.5) + extra_padding  # Space for y-axis labels
+            right_pad = int(10 * font_scale) + extra_padding
+
+            # Calculate top padding, accounting for title if present
+            top_pad = int(10 * font_scale) + extra_padding
+            title_font_scale = font_scale * 1.5  # Make title slightly larger
+            if title is not None:
+                title_size = cv2.getTextSize(title, font, title_font_scale, 1)[0]
+                top_pad += title_size[1] + int(10 * font_scale)
+
+            bottom_pad = int(x_text_size[1] * 2.0) + extra_padding  # Space for x-axis labels
+
+            # Create larger image to accommodate padding.
+            total_height = height
+            total_width = width
+            plot_width = width - left_pad - right_pad
+            plot_height = height - top_pad - bottom_pad
+            assert plot_width > 0 and plot_height > 0
+
+            # Create image with specified background color
+            img = np.ones((total_height, total_width, 3), dtype=np.uint8)
+            img[:] = background_color
+
+            # Create plot area with specified color
+            plot_area = np.ones((plot_height, plot_width, 3), dtype=np.uint8)
+            plot_area[:] = plot_area_color
+            img[top_pad : top_pad + plot_height, left_pad : left_pad + plot_width] = plot_area
+
+            def scale_to_pixels(values, bounds, pixels):
+                """Scale values from bounds range to pixel coordinates."""
+                min_val, max_val = bounds
+                normalized = (values - min_val) / (max_val - min_val)
+                return (normalized * (pixels - 1)).astype(np.int32)
+
+            # Vertical grid lines.
+            for i in range(grid_x_lines):
+                x_pos = left_pad + int(plot_width * i / (grid_x_lines - 1))
+                cv2.line(img, (x_pos, top_pad), (x_pos, top_pad + plot_height), grid_color, 1)
+
+            # Horizontal grid lines.
+            for i in range(grid_y_lines):
+                y_pos = top_pad + int(plot_height * i / (grid_y_lines - 1))
+                cv2.line(img, (left_pad, y_pos), (left_pad + plot_width, y_pos), grid_color, 1)
+
+            # Draw axes.
+            cv2.line(
+                img,
+                (left_pad, top_pad + plot_height),
+                (left_pad + plot_width, top_pad + plot_height),
+                axes_color,
+                1,
+            )  # x-axis
+            cv2.line(
+                img, (left_pad, top_pad), (left_pad, top_pad + plot_height), axes_color, 1
+            )  # y-axis
+
+            # Scale and plot the data.
+            x_scaled = scale_to_pixels(xs, x_bounds, plot_width) + left_pad
+            y_scaled = top_pad + plot_height - 1 - scale_to_pixels(ys, y_bounds, plot_height)
+            pts = np.column_stack((x_scaled, y_scaled)).reshape((-1, 1, 2))
+
+            # Draw the main plot line.
+            cv2.polylines(
+                img, [pts], False, line_color, thickness=line_thickness, lineType=cv2.LINE_AA
+            )
+
+            # Draw title if specified
+            if title is not None:
+                title_size = cv2.getTextSize(title, font, title_font_scale, 1)[0]
+                title_x = left_pad + (plot_width - title_size[0]) // 2
+                title_y = int(top_pad / 2) + title_size[1] // 2 - 1
+                cv2.putText(
+                    img,
+                    title,
+                    (title_x, title_y),
+                    font,
+                    title_font_scale,
+                    text_color,
+                    1,
+                    cv2.LINE_AA,
+                )
+
+            # X-axis labels.
+            for i in range(grid_x_lines):
+                x_val = x_bounds[0] + (x_bounds[1] - x_bounds[0]) * i / (grid_x_lines - 1)
+                x_pos = left_pad + int(plot_width * i / (grid_x_lines - 1))
+                label = f"{x_val:.1f}"
+                if label == "-0.0":
+                    label = "0.0"
+                text_size = cv2.getTextSize(label, font, font_scale, 1)[0]
+                cv2.putText(
+                    img,
+                    label,
+                    (x_pos - text_size[0] // 2, top_pad + plot_height + text_size[1] + 10),
+                    font,
+                    font_scale,
+                    text_color,
+                    1,
+                    cv2.LINE_AA,
+                )
+
+            # Y-axis labels.
+            for i in range(grid_y_lines):
+                y_val = y_bounds[0] + (y_bounds[1] - y_bounds[0]) * (grid_y_lines - 1 - i) / (
+                    grid_y_lines - 1
+                )
+                y_pos = top_pad + int(plot_height * i / (grid_y_lines - 1))
+                label = f"{y_val:.1f}"
+                if label == "-0.0":
+                    label = "0.0"
+                text_size = cv2.getTextSize(label, font, font_scale, 1)[0]
+                cv2.putText(
+                    img,
+                    label,
+                    (left_pad - text_size[0] - 5, y_pos + 5),
+                    font,
+                    font_scale,
+                    text_color,
+                    1,
+                    cv2.LINE_AA,
+                )
+
+            return img
+
+
+        def create_sine_plot(title: str, counter: int) -> np.ndarray:
+            """Create a sine wave plot with the given counter offset."""
+            xs = np.linspace(0, 2 * np.pi, 20)
+            rgb = colorsys.hsv_to_rgb(counter / 4000 % 1, 1, 1)
+            return get_line_plot(
+                xs=xs,
+                ys=np.sin(xs + counter / 20),
+                title=title,
+                line_color=(int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255)),
+                height=150,
+                width=350,
+            )
+
+
+        def main(num_plots: int = 8) -> None:
+            server = viser.ViserServer()
+
+            # Create GUI elements for display runtimes.
+            with server.gui.add_folder("Runtime"):
+                draw_time = server.gui.add_text("Draw / plot (ms)", "0.00", disabled=True)
+                send_gui_time = server.gui.add_text(
+                    "Gui update / plot (ms)", "0.00", disabled=True
+                )
+                send_scene_time = server.gui.add_text(
+                    "Scene update / plot (ms)", "0.00", disabled=True
+                )
+
+            # Add 2D plots to the GUI.
+            with server.gui.add_folder("Plots"):
+                plots_cb = server.gui.add_checkbox("Update plots", True)
+                gui_image_handles = [
+                    server.gui.add_image(
+                        create_sine_plot(f"Plot {i}", counter=0),
+                        label=f"Image {i}",
+                        format="jpeg",
+                    )
+                    for i in range(num_plots)
+                ]
+
+            # Add 2D plots to the scene. We flip them with a parent coordinate frame.
+            server.scene.add_frame(
+                "/images", wxyz=vtf.SO3.from_y_radians(np.pi).wxyz, show_axes=False
+            )
+            scene_image_handles = [
+                server.scene.add_image(
+                    f"/images/plot{i}",
+                    image=gui_image_handles[i].image,
+                    render_width=3.5,
+                    render_height=1.5,
+                    format="jpeg",
+                    position=(
+                        (i % 2 - 0.5) * 3.5,
+                        (i // 2 - (num_plots - 1) / 4) * 1.5,
+                        0,
+                    ),
+                )
+                for i in range(num_plots)
+            ]
+
+            counter = 0
+
+            while True:
+                if plots_cb.value:
+                    # Create and time the plot generation.
+                    start = time.time()
+                    images = [
+                        create_sine_plot(f"Plot {i}", counter=counter * (i + 1))
+                        for i in range(num_plots)
+                    ]
+                    draw_time.value = f"{0.98 * float(draw_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
+
+                    # Update all plot images.
+                    start = time.time()
+                    for i, handle in enumerate(gui_image_handles):
+                        handle.image = images[i]
+                    send_gui_time.value = f"{0.98 * float(send_gui_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
+
+                    # Update all scene images.
+                    start = time.time()
+                    for i, handle in enumerate(scene_image_handles):
+                        handle.image = gui_image_handles[i].image
+                    send_scene_time.value = f"{0.98 * float(send_scene_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
+
+                # Sleep a bit before continuing.
+                time.sleep(0.02)
+                counter += 1
+
+
+        if __name__ == "__main__":
+            tyro.cli(main)