wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add sanity check

This commit is contained in:
Simon Alibert 2024-10-28 12:26:07 +01:00
parent 51e87f6f97
commit 4c22de20a6
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -749,6 +749,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
if not keep_image_files and self.image_writer is not None: if not keep_image_files and self.image_writer is not None:
shutil.rmtree(self.image_writer.write_dir) shutil.rmtree(self.image_writer.write_dir)
video_files = list(self.root.rglob("*.mp4"))
assert len(video_files) == self.num_episodes * len(self.video_keys)
parquet_files = list(self.root.rglob("*.parquet"))
assert len(parquet_files) == self.num_episodes
if run_compute_stats: if run_compute_stats:
self.stop_image_writer() self.stop_image_writer()
self.stats = compute_stats(self) self.stats = compute_stats(self)
@ -762,9 +768,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
# TODO(aliberts) # TODO(aliberts)
# - [X] add video info in info.json # - [X] add video info in info.json
# Sanity checks: # Sanity checks:
# - [X] number of files
# - [ ] shapes # - [ ] shapes
# - [ ] ep_lenghts # - [ ] ep_lenghts
# - [ ] number of files
@classmethod @classmethod
def create( def create(