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

Added data augmentation feature to MultiLeRobotDataset

This commit is contained in:
Marina Barannikov 2024-05-31 15:42:31 +00:00
parent 20a3715469
commit c4870e5892
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -381,6 +381,11 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
for data_key in self.disabled_data_keys: for data_key in self.disabled_data_keys:
if data_key in item: if data_key in item:
del item[data_key] del item[data_key]
if self.transform is not None:
for cam in self.camera_keys:
item[cam] = self.transform(item[cam])
return item return item
def __repr__(self): def __repr__(self):