tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix tests

Browse Source
This commit is contained in:
Cadene
2024-04-16 17:29:31 +00:00
parent 4327e43f19
commit 91badebdfc
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
lerobot/scripts/visualize_dataset.py
View File

@@ -77,7 +77,7 @@ def render_dataset(dataset, out_dir, max_num_episodes):
# add current frame to list of frames to render # add current frame to list of frames to render
frames[im_key].append(item[im_key]) frames[im_key].append(item[im_key])
end_of_episode = item["index"].item() == item["episode_data_index_to"].item() end_of_episode = item["index"].item() == item["episode_data_index_to"].item() - 1
out_dir.mkdir(parents=True, exist_ok=True) out_dir.mkdir(parents=True, exist_ok=True)
for im_key in dataset.image_keys: for im_key in dataset.image_keys:

6
tests/test_datasets.py
View File

@@ -157,7 +157,7 @@ def test_load_previous_and_future_frames_within_tolerance():
"timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
"index": [0, 1, 2, 3, 4], "index": [0, 1, 2, 3, 4],
"episode_data_index_from": [0, 0, 0, 0, 0], "episode_data_index_from": [0, 0, 0, 0, 0],
"episode_data_index_to": [4, 4, 4, 4, 4], "episode_data_index_to": [5, 5, 5, 5, 5],
}) })
data_dict = data_dict.with_format("torch") data_dict = data_dict.with_format("torch")
item = data_dict[2] item = data_dict[2]
@@ -173,7 +173,7 @@ def test_load_previous_and_future_frames_outside_tolerance_inside_episode_range(
"timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
"index": [0, 1, 2, 3, 4], "index": [0, 1, 2, 3, 4],
"episode_data_index_from": [0, 0, 0, 0, 0], "episode_data_index_from": [0, 0, 0, 0, 0],
"episode_data_index_to": [4, 4, 4, 4, 4], "episode_data_index_to": [5, 5, 5, 5, 5],
}) })
data_dict = data_dict.with_format("torch") data_dict = data_dict.with_format("torch")
item = data_dict[2] item = data_dict[2]
@@ -187,7 +187,7 @@ def test_load_previous_and_future_frames_outside_tolerance_outside_episode_range
"timestamp": [0.1, 0.2, 0.3, 0.4, 0.5], "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5],
"index": [0, 1, 2, 3, 4], "index": [0, 1, 2, 3, 4],
"episode_data_index_from": [0, 0, 0, 0, 0], "episode_data_index_from": [0, 0, 0, 0, 0],
"episode_data_index_to": [4, 4, 4, 4, 4], "episode_data_index_to": [5, 5, 5, 5, 5],
}) })
data_dict = data_dict.with_format("torch") data_dict = data_dict.with_format("torch")
item = data_dict[2] item = data_dict[2]