From 42f9cc9c2a21085f3f3aa4aa19cf9a2e1d75b9c0 Mon Sep 17 00:00:00 2001 From: Marina Barannikov Date: Tue, 4 Jun 2024 11:14:54 +0000 Subject: [PATCH] Updated transforms arguments --- lerobot/common/datasets/transforms.py | 6 +++--- lerobot/configs/default.yaml | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/lerobot/common/datasets/transforms.py b/lerobot/common/datasets/transforms.py index f31a3ea7..6e8848ea 100644 --- a/lerobot/common/datasets/transforms.py +++ b/lerobot/common/datasets/transforms.py @@ -47,11 +47,11 @@ class RandomSubsetApply(Transform): def make_transforms(cfg): image_transforms = [] - if 'jit' in cfg.list_transforms: + if 'jit' in cfg.image_transform.list: image_transforms.append(v2.ColorJitter(brightness=cfg.colorjitter_range, contrast=cfg.colorjitter_range)) - if 'sharpness' in cfg.list_transforms: + if 'sharpness' in cfg.image_transform.list: image_transforms.append(v2.RandomAdjustSharpness(cfg.sharpness_factor, p=cfg.sharpness_p)) - if 'blur' in cfg.list_transforms: + if 'blur' in cfg.image_transform.list: image_transforms.append(v2.RandomAdjustSharpness(cfg.blur_factor, p=cfg.blur_p)) return v2.Compose(RandomSubsetApply(image_transforms, n_subset=cfg.n_subset), v2.ToDtype(torch.float32, scale=True)) diff --git a/lerobot/configs/default.yaml b/lerobot/configs/default.yaml index 4f003fe9..e4cec15e 100644 --- a/lerobot/configs/default.yaml +++ b/lerobot/configs/default.yaml @@ -61,6 +61,7 @@ wandb: image_transform: enable: false colorjittor_range: (0, 1) + colorjittor_p: 1 # Range from which to sample colorjittor factor sharpness_factor: 3 # Should be more than 1, setting parameter to 1 does not change the image