hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

fix unit tests, stats was missing, visualize_dataset was broken

Browse Source
This commit is contained in:
Cadene
2024-04-16 12:53:31 +00:00
parent 69eeced9d9
commit 4a3eac4743
9 changed files with 37 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
lerobot/scripts/visualize_dataset.py
View File

@@ -62,12 +62,12 @@ def render_dataset(dataset, out_dir, max_num_episodes):
)
dl_iter = iter(dataloader)
num_episodes = len(dataset.data_ids_per_episode)
for ep_id in range(min(max_num_episodes, num_episodes)):
for ep_id in range(min(max_num_episodes, dataset.num_episodes)):
logging.info(f"Rendering episode {ep_id}")
frames = {}
for _ in dataset.data_ids_per_episode[ep_id]:
end_of_episode = False
while not end_of_episode:
item = next(dl_iter)
for im_key in dataset.image_keys:
@@ -77,6 +77,8 @@ def render_dataset(dataset, out_dir, max_num_episodes):
# add current frame to list of frames to render
frames[im_key].append(item[im_key])
end_of_episode = item["index"].item() == item["episode_data_id_to"].item()
out_dir.mkdir(parents=True, exist_ok=True)
for im_key in dataset.image_keys:
if len(dataset.image_keys) > 1: