From 3646f6dd4f416f0e2fdaca5d75321c06ce122552 Mon Sep 17 00:00:00 2001 From: Marina Barannikov Date: Wed, 12 Jun 2024 18:16:49 +0200 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com> --- Makefile | 8 ++++---- lerobot/common/datasets/factory.py | 25 +++++++++++++------------ 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/Makefile b/Makefile index cddf9ffd6..fafda8254 100644 --- a/Makefile +++ b/Makefile @@ -46,7 +46,7 @@ test-act-ete-train: policy.n_action_steps=20 \ policy.chunk_size=20 \ training.batch_size=2 \ - training.image_transforms.enable=true \ + training.image_transforms.enable=true \ hydra.run.dir=tests/outputs/act/ test-act-ete-eval: @@ -74,7 +74,7 @@ test-act-ete-train-amp: policy.chunk_size=20 \ training.batch_size=2 \ hydra.run.dir=tests/outputs/act_amp/ \ - training.image_transforms.enable=true \ + training.image_transforms.enable=true \ use_amp=true test-act-ete-eval-amp: @@ -102,7 +102,7 @@ test-diffusion-ete-train: training.save_checkpoint=true \ training.save_freq=2 \ training.batch_size=2 \ - training.image_transforms.enable=true \ + training.image_transforms.enable=true \ hydra.run.dir=tests/outputs/diffusion/ test-diffusion-ete-eval: @@ -130,7 +130,7 @@ test-tdmpc-ete-train: training.save_checkpoint=true \ training.save_freq=2 \ training.batch_size=2 \ - training.image_transforms.enable=true \ + training.image_transforms.enable=true \ hydra.run.dir=tests/outputs/tdmpc/ test-tdmpc-ete-eval: diff --git a/lerobot/common/datasets/factory.py b/lerobot/common/datasets/factory.py index fc23377d0..754bc91b2 100644 --- a/lerobot/common/datasets/factory.py +++ b/lerobot/common/datasets/factory.py @@ -74,19 +74,20 @@ def make_dataset(cfg, split: str = "train") -> LeRobotDataset | MultiLeRobotData image_transforms = None if cfg.training.image_transforms.enable: + cfg_tf = cfg.training.image_transforms image_transforms = get_image_transforms( - brightness_weight=cfg.training.image_transforms.brightness.weight, - brightness_min_max=cfg.training.image_transforms.brightness.min_max, - contrast_weight=cfg.training.image_transforms.contrast.weight, - contrast_min_max=cfg.training.image_transforms.contrast.min_max, - saturation_weight=cfg.training.image_transforms.saturation.weight, - saturation_min_max=cfg.training.image_transforms.saturation.min_max, - hue_weight=cfg.training.image_transforms.hue.weight, - hue_min_max=cfg.training.image_transforms.hue.min_max, - sharpness_weight=cfg.training.image_transforms.sharpness.weight, - sharpness_min_max=cfg.training.image_transforms.sharpness.min_max, - max_num_transforms=cfg.training.image_transforms.max_num_transforms, - random_order=cfg.training.image_transforms.random_order, + brightness_weight=cfg_tf.brightness.weight, + brightness_min_max=cfg_tf.brightness.min_max, + contrast_weight=cfg_tf.contrast.weight, + contrast_min_max=cfg_tf.contrast.min_max, + saturation_weight=cfg_tf.saturation.weight, + saturation_min_max=cfg_tf.saturation.min_max, + hue_weight=cfg_tf.hue.weight, + hue_min_max=cfg_tf.hue.min_max, + sharpness_weight=cfg_tf.sharpness.weight, + sharpness_min_max=cfg_tf.sharpness.min_max, + max_num_transforms=cfg_tf.max_num_transforms, + random_order=cfg_tf.random_order, ) if isinstance(cfg.dataset_repo_id, str):