From 37174dc73105bc5fa2e65baccec03f2e00b6a68c Mon Sep 17 00:00:00 2001 From: Simon Alibert Date: Wed, 11 Dec 2024 10:18:33 +0100 Subject: [PATCH] Fix example 6 --- examples/6_add_image_transforms.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/6_add_image_transforms.py b/examples/6_add_image_transforms.py index 82b70f5c1..882710e3d 100644 --- a/examples/6_add_image_transforms.py +++ b/examples/6_add_image_transforms.py @@ -10,10 +10,10 @@ from lerobot.common.datasets.lerobot_dataset import LeRobotDataset -dataset_repo_id = "lerobot/aloha_static_tape" +dataset_repo_id = "lerobot/aloha_static_screw_driver" # Create a LeRobotDataset with no transformations -dataset = LeRobotDataset(dataset_repo_id) +dataset = LeRobotDataset(dataset_repo_id, episodes=[0]) # This is equivalent to `dataset = LeRobotDataset(dataset_repo_id, image_transforms=None)` # Get the index of the first observation in the first episode @@ -28,12 +28,13 @@ [ v2.ColorJitter(brightness=(0.5, 1.5)), v2.ColorJitter(contrast=(0.5, 1.5)), + v2.ColorJitter(hue=(-0.1, 0.1)), v2.RandomAdjustSharpness(sharpness_factor=2, p=1), ] ) # Create another LeRobotDataset with the defined transformations -transformed_dataset = LeRobotDataset(dataset_repo_id, image_transforms=transforms) +transformed_dataset = LeRobotDataset(dataset_repo_id, episodes=[0], image_transforms=transforms) # Get a frame from the transformed dataset transformed_frame = transformed_dataset[first_idx][transformed_dataset.meta.camera_keys[0]]