Spaces:
Sleeping
Sleeping
Rename vars for clarity
Browse files
app.py
CHANGED
@@ -33,11 +33,11 @@ def html_template(rrd: str, app_url: str = "https://app.rerun.io") -> str:
|
|
33 |
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>"""
|
34 |
|
35 |
|
36 |
-
def show_dataset(dataset_id: str,
|
37 |
rr.init("dataset")
|
38 |
|
39 |
# TODO(jleibs): manage cache better and put in proper storage
|
40 |
-
filename = Path(f"tmp/{dataset_id}_{
|
41 |
if not filename.exists():
|
42 |
filename.parent.mkdir(parents=True, exist_ok=True)
|
43 |
|
@@ -46,7 +46,9 @@ def show_dataset(dataset_id: str, episode_id: int) -> str:
|
|
46 |
dataset = load_dataset(dataset_id, split="train", streaming=True)
|
47 |
|
48 |
# This is for LeRobot datasets (https://huggingface.co/lerobot):
|
49 |
-
ds_subset = dataset.filter(
|
|
|
|
|
50 |
|
51 |
log_dataset_to_rerun(ds_subset)
|
52 |
|
@@ -61,14 +63,14 @@ with gr.Blocks() as demo:
|
|
61 |
placeholder="Search for models on Huggingface",
|
62 |
search_type="dataset",
|
63 |
)
|
64 |
-
|
65 |
button = gr.Button("Show Dataset")
|
66 |
with gr.Row():
|
67 |
rrd = gr.File()
|
68 |
with gr.Row():
|
69 |
viewer = gr.HTML()
|
70 |
|
71 |
-
button.click(show_dataset, inputs=[search_in,
|
72 |
rrd.change(
|
73 |
html_template,
|
74 |
js="""(rrd) => { console.log(rrd.url); return rrd.url}""",
|
|
|
33 |
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>"""
|
34 |
|
35 |
|
36 |
+
def show_dataset(dataset_id: str, episode_index: int) -> str:
|
37 |
rr.init("dataset")
|
38 |
|
39 |
# TODO(jleibs): manage cache better and put in proper storage
|
40 |
+
filename = Path(f"tmp/{dataset_id}_{episode_index}.rrd")
|
41 |
if not filename.exists():
|
42 |
filename.parent.mkdir(parents=True, exist_ok=True)
|
43 |
|
|
|
46 |
dataset = load_dataset(dataset_id, split="train", streaming=True)
|
47 |
|
48 |
# This is for LeRobot datasets (https://huggingface.co/lerobot):
|
49 |
+
ds_subset = dataset.filter(
|
50 |
+
lambda frame: "episode_index" not in frame or frame["episode_index"] == episode_index
|
51 |
+
)
|
52 |
|
53 |
log_dataset_to_rerun(ds_subset)
|
54 |
|
|
|
63 |
placeholder="Search for models on Huggingface",
|
64 |
search_type="dataset",
|
65 |
)
|
66 |
+
episode_index = gr.Number(1, label="Episode Index")
|
67 |
button = gr.Button("Show Dataset")
|
68 |
with gr.Row():
|
69 |
rrd = gr.File()
|
70 |
with gr.Row():
|
71 |
viewer = gr.HTML()
|
72 |
|
73 |
+
button.click(show_dataset, inputs=[search_in, episode_index], outputs=rrd)
|
74 |
rrd.change(
|
75 |
html_template,
|
76 |
js="""(rrd) => { console.log(rrd.url); return rrd.url}""",
|
main.py
CHANGED
@@ -20,14 +20,16 @@ def main() -> None:
|
|
20 |
|
21 |
parser = argparse.ArgumentParser(description="Log a HuggingFace dataset to Rerun.")
|
22 |
parser.add_argument("--dataset", default="lerobot/pusht", help="The name of the dataset to load")
|
23 |
-
parser.add_argument("--episode-
|
24 |
args = parser.parse_args()
|
25 |
|
26 |
print("Loading dataset…")
|
27 |
dataset = load_dataset(args.dataset, split="train", streaming=True)
|
28 |
|
29 |
# This is for LeRobot datasets (https://huggingface.co/lerobot):
|
30 |
-
ds_subset = dataset.filter(
|
|
|
|
|
31 |
|
32 |
print("Starting Rerun…")
|
33 |
rr.init(f"rerun_example_huggingface {args.dataset}", spawn=True)
|
|
|
20 |
|
21 |
parser = argparse.ArgumentParser(description="Log a HuggingFace dataset to Rerun.")
|
22 |
parser.add_argument("--dataset", default="lerobot/pusht", help="The name of the dataset to load")
|
23 |
+
parser.add_argument("--episode-index", default=1, help="Which episode to select")
|
24 |
args = parser.parse_args()
|
25 |
|
26 |
print("Loading dataset…")
|
27 |
dataset = load_dataset(args.dataset, split="train", streaming=True)
|
28 |
|
29 |
# This is for LeRobot datasets (https://huggingface.co/lerobot):
|
30 |
+
ds_subset = dataset.filter(
|
31 |
+
lambda frame: "episode_index" not in frame or frame["episode_index"] == args.episode_index
|
32 |
+
)
|
33 |
|
34 |
print("Starting Rerun…")
|
35 |
rr.init(f"rerun_example_huggingface {args.dataset}", spawn=True)
|