rerun/cube_example

1import numpy as np2from math import cos, sin, pi3from collections import namedtuple4import time5 6import gradio as gr7from gradio_rerun import Rerun8 9import rerun as rr10import rerun.blueprint as rrb11 12 13ColorGrid = namedtuple("ColorGrid", ["positions", "colors"])14 15 16def build_color_grid(x_count=10, y_count=10, z_count=10, twist=0):17    """18    Create a cube of points with colors.19 20    The total point cloud will have x_count * y_count * z_count points.21 22    Parameters23    ----------24    x_count, y_count, z_count:25        Number of points in each dimension.26    twist:27        Angle to twist from bottom to top of the cube28 29    """30 31    grid = np.mgrid[32        slice(-x_count, x_count, x_count * 1j),33        slice(-y_count, y_count, y_count * 1j),34        slice(-z_count, z_count, z_count * 1j),35    ]36 37    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)38    for z in range(z_count):39        xv, yv, zv = grid[:, :, :, z]40        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])41        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])42        grid[:, :, :, z] = [rot_xv, rot_yv, zv]43 44    positions = np.vstack([xyz.ravel() for xyz in grid])45 46    colors = np.vstack(47        [48            xyz.ravel()49            for xyz in np.mgrid[50                slice(0, 255, x_count * 1j),51                slice(0, 255, y_count * 1j),52                slice(0, 255, z_count * 1j),53            ]54        ]55    )56 57    return ColorGrid(positions.T, colors.T.astype(np.uint8))58 59 60@rr.thread_local_stream("rerun_example_cube")61def show_cube(x, y, z, twist):62    stream = rr.binary_stream()63 64    blueprint = rrb.Blueprint(65        rrb.Spatial3DView(origin="cube"),66        collapse_panels=True,67    )68    rr.send_blueprint(blueprint)69 70    yield stream.read()71 72    for step in range(twist):73        rr.set_time_sequence("steps", step)74        cube = build_color_grid(int(x), int(y), int(z), twist=pi * step / 180)75        rr.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))76        yield stream.read()77        time.sleep(0.01)78 79 80with gr.Blocks() as demo:81    with gr.Row():82        x_count = gr.Number(83            minimum=1, maximum=20, value=10, precision=0, label="X Count"84        )85        y_count = gr.Number(86            minimum=1, maximum=20, value=10, precision=0, label="Y Count"87        )88        z_count = gr.Number(89            minimum=1, maximum=20, value=10, precision=0, label="Z Count"90        )91        twist_steps = gr.Number(92            minimum=0, maximum=180, value=90, precision=0, label="Twist (degrees)"93        )94        button = gr.Button("Show Cube")95    with gr.Row():96        viewer = Rerun(streaming=True)97 98    button.click(99        show_cube, inputs=[x_count, y_count, z_count, twist_steps], outputs=viewer100    )101 102demo.queue(default_concurrency_limit=10)103demo.launch()104