hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Added visualisations for image augmentation

Browse Source
This commit is contained in:
Marina Barannikov
2024-06-04 11:57:45 +00:00
parent 42f9cc9c2a
commit 5eea2542d9
3 changed files with 74 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
lerobot/common/datasets/transforms.py
View File

@@ -47,11 +47,11 @@ class RandomSubsetApply(Transform):
def make_transforms(cfg):
image_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.image_transform.list:
if 'colorjitter' in cfg.list:
image_transforms.append(v2.ColorJitter(brightness=cfg.colorjitter_factor, contrast=cfg.colorjitter_factor))
if 'sharpness' in cfg.list:
image_transforms.append(v2.RandomAdjustSharpness(cfg.sharpness_factor, p=cfg.sharpness_p))
if 'blur' in cfg.image_transform.list:
if 'blur' in cfg.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))
return v2.Compose([RandomSubsetApply(image_transforms, n_subset=cfg.n_subset), v2.ToDtype(torch.float32, scale=True)])