Integrate with rerun dataset converter

app.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 import urllib
 from collections import namedtuple
 from math import cos, sin
@@ -5,10 +6,13 @@ from typing import Any
 
 import gradio as gr
 import numpy as np
+from dataset_conversion import log_dataset_to_rerun
 import rerun as rr
 import rerun.blueprint as rrb
+from datasets import load_dataset
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from gradio_huggingfacehub_search import HuggingfaceHubSearch
 
 CUSTOM_PATH = "/"
 
@@ -24,81 +28,47 @@ app.add_middleware(
 )
 
 
-ColorGrid = namedtuple("ColorGrid", ["positions", "colors"])
-
-
-def build_color_grid(x_count: int = 10, y_count: int = 10, z_count: int = 10, twist: float = 0) -> ColorGrid:
-    """
-    Create a cube of points with colors.
-
-    The total point cloud will have x_count * y_count * z_count points.
-
-    Parameters
-    ----------
-    x_count, y_count, z_count:
-        Number of points in each dimension.
-    twist:
-        Angle to twist from bottom to top of the cube
-
-    """
-
-    grid = np.mgrid[
-        slice(-x_count, x_count, x_count * 1j),
-        slice(-y_count, y_count, y_count * 1j),
-        slice(-z_count, z_count, z_count * 1j),
-    ]
-
-    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)
-    for z in range(z_count):
-        xv, yv, zv = grid[:, :, :, z]
-        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])
-        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])
-        grid[:, :, :, z] = [rot_xv, rot_yv, zv]
-
-    positions = np.vstack([xyz.ravel() for xyz in grid])
-
-    colors = np.vstack([
-        xyz.ravel()
-        for xyz in np.mgrid[
-            slice(0, 255, x_count * 1j),
-            slice(0, 255, y_count * 1j),
-            slice(0, 255, z_count * 1j),
-        ]
-    ])
-
-    return ColorGrid(positions.T, colors.T.astype(np.uint8))
-
-
 def html_template(rrd: str, app_url: str = "https://app.rerun.io") -> str:
     encoded_url = urllib.parse.quote(rrd)
     return f"""<div style="width:100%; height:70vh;"><iframe style="width:100%; height:100%;" src="{app_url}?url={encoded_url}" frameborder="0" allowfullscreen=""></iframe></div>"""
 
 
-def show_cube(x: int, y: int, z: int) -> str:
-    rr.init("my data")
+def show_dataset(dataset_id: str, episode_id: int) -> str:
+    rr.init("dataset")
+
+    # TODO(jleibs): manage cache better and put in proper storage
+    filename = Path(f"tmp/{dataset_id}_{episode_id}.rrd")
+    if not filename.exists():
+        filename.parent.mkdir(parents=True, exist_ok=True)
+
+        rr.save(filename.as_posix())
 
-    cube = build_color_grid(int(x), int(y), int(z), twist=0)
-    rr.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))
+        dataset = load_dataset(dataset_id, split="train", streaming=True)
 
-    blueprint = rrb.Spatial3DView(origin="cube")
+        # This is for LeRobot datasets (https://huggingface.co/lerobot):
+        ds_subset = dataset.filter(lambda frame: "episode_index" not in frame or frame["episode_index"] == episode_id)
 
-    rr.save("cube.rrd", default_blueprint=blueprint)
+        log_dataset_to_rerun(ds_subset)
 
-    return "cube.rrd"
+    return filename.as_posix()
 
 
 with gr.Blocks() as demo:
     with gr.Row():
-        x_count = gr.Number(minimum=1, maximum=10, value=5, precision=0, label="X Count")
-        y_count = gr.Number(minimum=1, maximum=10, value=5, precision=0, label="Y Count")
-        z_count = gr.Number(minimum=1, maximum=10, value=5, precision=0, label="Z Count")
-        button = gr.Button("Show Cube")
+        search_in = HuggingfaceHubSearch(
+            "lerobot/pusht",
+            label="Search Huggingface Hub",
+            placeholder="Search for models on Huggingface",
+            search_type="dataset",
+        )
+        episode_id = gr.Number(1, label="Episode ID")
+        button = gr.Button("Show Dataset")
     with gr.Row():
         rrd = gr.File()
     with gr.Row():
         viewer = gr.HTML()
 
-    button.click(show_cube, inputs=[x_count, y_count, z_count], outputs=rrd)
+    button.click(show_dataset, inputs=[search_in, episode_id], outputs=rrd)
     rrd.change(
         html_template,
         js="""(rrd) => { console.log(rrd.url); return rrd.url}""",

requirements.txt

@@ -1,6 +1,7 @@
 datasets
 h5py
 gradio==4.27.0
+gradio_huggingfacehub_search
 pillow
 rerun-sdk>=0.15.0,<0.16.0
 tqdm