Add dataset visualization with rerun.io (#131)

Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>

lerobot/common/datasets/_video_benchmark/run_video_benchmark.py

@@ -1,5 +1,4 @@
 import json
-import os
 import random
 import shutil
 import subprocess
@@ -41,10 +40,8 @@ def run_video_benchmark(
     repo_id = cfg["repo_id"]
 
     # TODO(rcadene): rewrite with hardcoding of original images and episodes
-    dataset = LeRobotDataset(
-        repo_id,
-        root=Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None,
-    )
+    dataset = LeRobotDataset(repo_id)
+
     # Get fps
     fps = dataset.fps
 

lerobot/common/datasets/factory.py

@@ -1,14 +1,10 @@
 import logging
-import os
-from pathlib import Path
 
 import torch
 from omegaconf import OmegaConf
 
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 
-DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None
-
 
 def make_dataset(
     cfg,
@@ -31,7 +27,6 @@ def make_dataset(
     dataset = LeRobotDataset(
         cfg.dataset_repo_id,
         split=split,
-        root=DATA_DIR,
         delta_timestamps=delta_timestamps,
     )
 

lerobot/common/datasets/lerobot_dataset.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 import datasets
@@ -13,7 +14,8 @@ from lerobot.common.datasets.utils import (
 )
 from lerobot.common.datasets.video_utils import VideoFrame, load_from_videos
 
-CODEBASE_VERSION = "v1.2"
+DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None
+CODEBASE_VERSION = "v1.3"
 
 
 class LeRobotDataset(torch.utils.data.Dataset):
@@ -21,7 +23,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         self,
         repo_id: str,
         version: str | None = CODEBASE_VERSION,
-        root: Path | None = None,
+        root: Path | None = DATA_DIR,
         split: str = "train",
         transform: callable = None,
         delta_timestamps: dict[list[float]] | None = None,