ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-03 20:31:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(tests) bug in clear_episode_buffer

Browse Source
This commit is contained in:
Michel Aractingi
2025-07-20 01:39:19 +02:00
parent 8ffc00dbcd
commit 23375cce3a
1 changed files with 2 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/lerobot/datasets/lerobot_dataset.py
View File

@@ -876,8 +876,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
None.
"""
if not episode_data:
episode_buffer = self.episode_buffer
episode_buffer = episode_data if episode_data is not None else self.episode_buffer
validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
@@ -1036,12 +1035,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
return metadata
def clear_episode_buffer(self) -> None:
episode_index = self.episode_buffer["episode_index"]
if self.image_writer is not None:
for cam_key in self.meta.camera_keys:
img_dir = self._get_image_file_path(
episode_index=episode_index, image_key=cam_key, frame_index=0
).parent
img_dir = self.root / "images" / cam_key
if img_dir.is_dir():
shutil.rmtree(img_dir)