Plots as Images

Examples of sending plots as images to Viser’s GUI panel. This can be faster than using Plotly.

  1import colorsys
  2import time
  3
  4import cv2
  5import numpy as np
  6import tyro
  7
  8import viser
  9import viser.transforms as vtf
 10
 11
 12def get_line_plot(
 13    xs: np.ndarray,
 14    ys: np.ndarray,
 15    height: int,
 16    width: int,
 17    *,
 18    x_bounds: tuple[float, float] | None = None,
 19    y_bounds: tuple[float, float] | None = None,
 20    title: str | None = None,
 21    line_thickness: int = 2,
 22    grid_x_lines: int = 8,
 23    grid_y_lines: int = 5,
 24    font_scale: float = 0.4,
 25    background_color: tuple[int, int, int] = (0, 0, 0),
 26    plot_area_color: tuple[int, int, int] = (0, 0, 0),
 27    grid_color: tuple[int, int, int] = (60, 60, 60),
 28    axes_color: tuple[int, int, int] = (100, 100, 100),
 29    line_color: tuple[int, int, int] = (255, 255, 255),
 30    text_color: tuple[int, int, int] = (200, 200, 200),
 31) -> np.ndarray:
 32    """Create a line plot using OpenCV with axes, labels, and grid.
 33
 34    This is much faster than using libraries like Matplotlib or Plotly, but is
 35    less flexible.
 36    """
 37
 38    if x_bounds is None:
 39        x_bounds = (np.min(xs), np.max(xs.round(decimals=4)))
 40    if y_bounds is None:
 41        y_bounds = (np.min(ys), np.max(ys))
 42
 43    # Calculate text sizes for padding.
 44    font = cv2.FONT_HERSHEY_DUPLEX
 45    sample_y_label = f"{max(abs(y_bounds[0]), abs(y_bounds[1])):.1f}"
 46    y_text_size = cv2.getTextSize(sample_y_label, font, font_scale, 1)[0]
 47
 48    sample_x_label = f"{max(abs(x_bounds[0]), abs(x_bounds[1])):.1f}"
 49    x_text_size = cv2.getTextSize(sample_x_label, font, font_scale, 1)[0]
 50
 51    # Define padding based on font scale.
 52    extra_padding = 8
 53    left_pad = int(y_text_size[0] * 1.5) + extra_padding  # Space for y-axis labels
 54    right_pad = int(10 * font_scale) + extra_padding
 55
 56    # Calculate top padding, accounting for title if present
 57    top_pad = int(10 * font_scale) + extra_padding
 58    title_font_scale = font_scale * 1.5  # Make title slightly larger
 59    if title is not None:
 60        title_size = cv2.getTextSize(title, font, title_font_scale, 1)[0]
 61        top_pad += title_size[1] + int(10 * font_scale)
 62
 63    bottom_pad = int(x_text_size[1] * 2.0) + extra_padding  # Space for x-axis labels
 64
 65    # Create larger image to accommodate padding.
 66    total_height = height
 67    total_width = width
 68    plot_width = width - left_pad - right_pad
 69    plot_height = height - top_pad - bottom_pad
 70    assert plot_width > 0 and plot_height > 0
 71
 72    # Create image with specified background color
 73    img = np.ones((total_height, total_width, 3), dtype=np.uint8)
 74    img[:] = background_color
 75
 76    # Create plot area with specified color
 77    plot_area = np.ones((plot_height, plot_width, 3), dtype=np.uint8)
 78    plot_area[:] = plot_area_color
 79    img[top_pad : top_pad + plot_height, left_pad : left_pad + plot_width] = plot_area
 80
 81    def scale_to_pixels(values, bounds, pixels):
 82        """Scale values from bounds range to pixel coordinates."""
 83        min_val, max_val = bounds
 84        normalized = (values - min_val) / (max_val - min_val)
 85        return (normalized * (pixels - 1)).astype(np.int32)
 86
 87    # Vertical grid lines.
 88    for i in range(grid_x_lines):
 89        x_pos = left_pad + int(plot_width * i / (grid_x_lines - 1))
 90        cv2.line(img, (x_pos, top_pad), (x_pos, top_pad + plot_height), grid_color, 1)
 91
 92    # Horizontal grid lines.
 93    for i in range(grid_y_lines):
 94        y_pos = top_pad + int(plot_height * i / (grid_y_lines - 1))
 95        cv2.line(img, (left_pad, y_pos), (left_pad + plot_width, y_pos), grid_color, 1)
 96
 97    # Draw axes.
 98    cv2.line(
 99        img,
100        (left_pad, top_pad + plot_height),
101        (left_pad + plot_width, top_pad + plot_height),
102        axes_color,
103        1,
104    )  # x-axis
105    cv2.line(
106        img, (left_pad, top_pad), (left_pad, top_pad + plot_height), axes_color, 1
107    )  # y-axis
108
109    # Scale and plot the data.
110    x_scaled = scale_to_pixels(xs, x_bounds, plot_width) + left_pad
111    y_scaled = top_pad + plot_height - 1 - scale_to_pixels(ys, y_bounds, plot_height)
112    pts = np.column_stack((x_scaled, y_scaled)).reshape((-1, 1, 2))
113
114    # Draw the main plot line.
115    cv2.polylines(
116        img, [pts], False, line_color, thickness=line_thickness, lineType=cv2.LINE_AA
117    )
118
119    # Draw title if specified
120    if title is not None:
121        title_size = cv2.getTextSize(title, font, title_font_scale, 1)[0]
122        title_x = left_pad + (plot_width - title_size[0]) // 2
123        title_y = int(top_pad / 2) + title_size[1] // 2 - 1
124        cv2.putText(
125            img,
126            title,
127            (title_x, title_y),
128            font,
129            title_font_scale,
130            text_color,
131            1,
132            cv2.LINE_AA,
133        )
134
135    # X-axis labels.
136    for i in range(grid_x_lines):
137        x_val = x_bounds[0] + (x_bounds[1] - x_bounds[0]) * i / (grid_x_lines - 1)
138        x_pos = left_pad + int(plot_width * i / (grid_x_lines - 1))
139        label = f"{x_val:.1f}"
140        if label == "-0.0":
141            label = "0.0"
142        text_size = cv2.getTextSize(label, font, font_scale, 1)[0]
143        cv2.putText(
144            img,
145            label,
146            (x_pos - text_size[0] // 2, top_pad + plot_height + text_size[1] + 10),
147            font,
148            font_scale,
149            text_color,
150            1,
151            cv2.LINE_AA,
152        )
153
154    # Y-axis labels.
155    for i in range(grid_y_lines):
156        y_val = y_bounds[0] + (y_bounds[1] - y_bounds[0]) * (grid_y_lines - 1 - i) / (
157            grid_y_lines - 1
158        )
159        y_pos = top_pad + int(plot_height * i / (grid_y_lines - 1))
160        label = f"{y_val:.1f}"
161        if label == "-0.0":
162            label = "0.0"
163        text_size = cv2.getTextSize(label, font, font_scale, 1)[0]
164        cv2.putText(
165            img,
166            label,
167            (left_pad - text_size[0] - 5, y_pos + 5),
168            font,
169            font_scale,
170            text_color,
171            1,
172            cv2.LINE_AA,
173        )
174
175    return img
176
177
178def create_sine_plot(title: str, counter: int) -> np.ndarray:
179    """Create a sine wave plot with the given counter offset."""
180    xs = np.linspace(0, 2 * np.pi, 20)
181    rgb = colorsys.hsv_to_rgb(counter / 4000 % 1, 1, 1)
182    return get_line_plot(
183        xs=xs,
184        ys=np.sin(xs + counter / 20),
185        title=title,
186        line_color=(int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255)),
187        height=150,
188        width=350,
189    )
190
191
192def main(num_plots: int = 8) -> None:
193    server = viser.ViserServer()
194
195    # Create GUI elements for display runtimes.
196    with server.gui.add_folder("Runtime"):
197        draw_time = server.gui.add_text("Draw / plot (ms)", "0.00", disabled=True)
198        send_gui_time = server.gui.add_text(
199            "Gui update / plot (ms)", "0.00", disabled=True
200        )
201        send_scene_time = server.gui.add_text(
202            "Scene update / plot (ms)", "0.00", disabled=True
203        )
204
205    # Add 2D plots to the GUI.
206    with server.gui.add_folder("Plots"):
207        plots_cb = server.gui.add_checkbox("Update plots", True)
208        gui_image_handles = [
209            server.gui.add_image(
210                create_sine_plot(f"Plot {i}", counter=0),
211                label=f"Image {i}",
212                format="jpeg",
213            )
214            for i in range(num_plots)
215        ]
216
217    # Add 2D plots to the scene. We flip them with a parent coordinate frame.
218    server.scene.add_frame(
219        "/images", wxyz=vtf.SO3.from_y_radians(np.pi).wxyz, show_axes=False
220    )
221    scene_image_handles = [
222        server.scene.add_image(
223            f"/images/plot{i}",
224            image=gui_image_handles[i].image,
225            render_width=3.5,
226            render_height=1.5,
227            format="jpeg",
228            position=(
229                (i % 2 - 0.5) * 3.5,
230                (i // 2 - (num_plots - 1) / 4) * 1.5,
231                0,
232            ),
233        )
234        for i in range(num_plots)
235    ]
236
237    counter = 0
238
239    while True:
240        if plots_cb.value:
241            # Create and time the plot generation.
242            start = time.time()
243            images = [
244                create_sine_plot(f"Plot {i}", counter=counter * (i + 1))
245                for i in range(num_plots)
246            ]
247            draw_time.value = f"{0.98 * float(draw_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
248
249            # Update all plot images.
250            start = time.time()
251            for i, handle in enumerate(gui_image_handles):
252                handle.image = images[i]
253            send_gui_time.value = f"{0.98 * float(send_gui_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
254
255            # Update all scene images.
256            start = time.time()
257            for i, handle in enumerate(scene_image_handles):
258                handle.image = gui_image_handles[i].image
259            send_scene_time.value = f"{0.98 * float(send_scene_time.value) + 0.02 * (time.time() - start) / num_plots * 1000:.2f}"
260
261        # Sleep a bit before continuing.
262        time.sleep(0.02)
263        counter += 1
264
265
266if __name__ == "__main__":
267    tyro.cli(main)