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

Updated formatting

Browse Source
This commit is contained in:
Marina Barannikov
2024-06-04 12:06:36 +00:00
parent 31e3c82386
commit 22bd1f0669
3 changed files with 24 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -47,11 +47,15 @@ class RandomSubsetApply(Transform):
def make_transforms(cfg):
image_transforms = []
if 'colorjitter' in cfg.list:
image_transforms.append(v2.ColorJitter(brightness=cfg.colorjitter_factor, contrast=cfg.colorjitter_factor))
if 'sharpness' in cfg.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.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)]
)