Skip to content

Commit

Permalink
added visualization for min and max transforms
Browse files Browse the repository at this point in the history
  • Loading branch information
marinabar committed Jun 13, 2024
1 parent 125bd93 commit 91b5a28
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions lerobot/scripts/visualize_image_transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ def save_config_single_transforms(cfg, original_frame, output_dir):
]
print("Individual transforms examples saved to:")
for transform in transforms:
# Apply one transformation with random value in min_max range
kwargs = {
f"{transform}_weight": cfg[f"{transform}"].weight,
f"{transform}_min_max": cfg[f"{transform}"].min_max,
Expand All @@ -118,6 +119,25 @@ def save_config_single_transforms(cfg, original_frame, output_dir):
transformed_frame = tf(original_frame)
to_pil(transformed_frame).save(output_dir_single / f"{i}.png", quality=100)

# Apply min transformation
min_value, max_value = cfg[f"{transform}"].min_max
kwargs = {
f"{transform}_weight": cfg[f"{transform}"].weight,
f"{transform}_min_max": (min_value, min_value),
}
tf = get_image_transforms(**kwargs)
transformed_frame = tf(original_frame)
to_pil(transformed_frame).save(output_dir_single / "min.png", quality=100)

# Apply max transformation
kwargs = {
f"{transform}_weight": cfg[f"{transform}"].weight,
f"{transform}_min_max": (max_value, max_value),
}
tf = get_image_transforms(**kwargs)
transformed_frame = tf(original_frame)
to_pil(transformed_frame).save(output_dir_single / "max.png", quality=100)

print(f" {output_dir_single}")


Expand Down

0 comments on commit 91b5a28

Please sign in to comment.