Make it easy to log any HuggingFace dataset

README.md

@@ -1,4 +1,6 @@
-# Rerun visualization for LeRobot
+# Rerun visualization of HuggingFace datasets
+Tested with LeRobot datasets:
+
 * https://huggingface.co/lerobot
 * https://huggingface.co/datasets/lerobot/pusht
 

main.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import argparse
+
 import rerun as rr
 from datasets import load_dataset
 from PIL import Image
@@ -43,16 +45,34 @@ def log_dataset_to_rerun(dataset) -> None:
             elif isinstance(cell, float) or isinstance(cell, int):
                 rr.log(column_name, rr.Scalar(cell))
             else:
+                # TODO(emilk): check if it is a tensor and then log it using rr.Tensor
                 rr.log(column_name, rr.TextDocument(str(cell)))
 
 
 def main():
+    # Define the available datasets
+    available_datasets = [
+        "lerobot/aloha_sim_insertion_human",
+        "lerobot/aloha_sim_insertion_scripted",
+        "lerobot/aloha_sim_transfer_cube_human",
+        "lerobot/aloha_sim_transfer_cube_scripted",
+        "lerobot/pusht",
+        "lerobot/xarm_lift_medium",
+    ]
+
+    # Create the parser
+    parser = argparse.ArgumentParser(description="Log a HuggingFace dataset to Rerun.")
+    parser.add_argument("--dataset", choices=available_datasets, default="pusht", help="The name of the dataset to load")
+    parser.add_argument("--episode-id", default=1, help="Which episode to select")
+
+    # Parse the arguments
+    args = parser.parse_args()
+
     print("Loading dataset…")
-    # dataset = load_dataset("lerobot/pusht", split="train")
-    dataset = load_dataset("lerobot/aloha_sim_transfer_cube_human", split="train")
+    dataset = load_dataset(args.dataset, split="train")
 
-    print("Selecting specific episode…")
-    ds_subset = dataset.filter(lambda frame: frame["episode_id"] == 3)
+    print("Selecting episode {args.episode_id}…")
+    ds_subset = dataset.filter(lambda frame: frame["episode_id"] == args.episode_id)
 
     print("Starting Rerun…")
     rr.init("rerun_example_lerobot", spawn=True)