fixed record function with names

huggingface · Nov 18, 2024 · 7603ad6 · 7603ad6
1 parent baab583
commit 7603ad6
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 148 deletions.
diff --git a/lerobot/common/datasets/populate_dataset.py b/lerobot/common/datasets/populate_dataset.py
@@ -48,7 +48,9 @@ def wrapper(*args, **kwargs):
 
 
 def save_image(img_tensor, key, frame_index, episode_index, videos_dir: str):
-    img = Image.fromarray(img_tensor.numpy())
+    if torch.is_tensor(img_tensor):
+        img_tensor = img_tensor.numpy()
+    img = Image.fromarray(img_tensor)
     path = Path(videos_dir) / f"{key}_episode_{episode_index:06d}" / f"frame_{frame_index:06d}.png"
     path.parent.mkdir(parents=True, exist_ok=True)
     img.save(str(path), quality=100)