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