Added data augmentation feature to MultiLeRobotDataset

2024-05-31 15:42:31 +00:00 · 2024-05-31 15:42:31 +00:00 · c4870e5892
parent 20a3715469
commit c4870e5892
1 changed files with 5 additions and 0 deletions
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@ -381,6 +381,11 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
        for data_key in self.disabled_data_keys:
            if data_key in item:
                del item[data_key]
+
+        if self.transform is not None:
+            for cam in self.camera_keys:
+                item[cam] = self.transform(item[cam])
+
        return item

    def __repr__(self):