Better instructions

README.md

@@ -1,5 +1,7 @@
 # Rerun visualization of HuggingFace datasets
-Tested with LeRobot datasets:
+Visualize HuggingFace datasets using [Rerun](https://www.rerun.io/).
+
+Originally built for the LeRobot datasets:
 
 * https://huggingface.co/lerobot
 * https://huggingface.co/datasets/lerobot/pusht
@@ -12,10 +14,10 @@ Requires Python 3.10 or higher.
 
 ```sh
 pip install -r requirements.txt
-python main.py
+python main.py --dataset lerobot/aloha_sim_insertion_human
 ```
 
-Example datasets to explore (use `python main.py --dataset`):
+Example datasets to explore:
 * `lerobot/aloha_sim_insertion_human`
 * `lerobot/aloha_sim_insertion_scripted`
 * `lerobot/aloha_sim_transfer_cube_human`

main.py

@@ -75,7 +75,7 @@ def main():
  print("Loading dataset…")
  dataset = load_dataset(args.dataset, split="train", streaming=True)
 
- print(f"Selecting episode {args.episode_id}…")
+ # This is for LeRobot datasets (https://huggingface.co/lerobot):
  ds_subset = dataset.filter(lambda frame: "episode_id" not in frame or frame["episode_id"] == args.episode_id)
 
  print("Starting Rerun…")

requirements.txt

@@ -3,3 +3,4 @@ h5py
 pillow
 rerun-sdk>=0.15.0,<0.16.0
 tqdm
+webdataset