From c4870e58922e5505f9195534ff6ee095238e5e69 Mon Sep 17 00:00:00 2001
From: Marina Barannikov <marina.barannikov@huggingface.co>
Date: Fri, 31 May 2024 15:42:31 +0000
Subject: [PATCH] Added data augmentation feature to MultiLeRobotDataset

---
 lerobot/common/datasets/lerobot_dataset.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py
index 929b83e3..491ba00b 100644
--- 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):