76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
|
from lerobot.common.datasets.transforms import (
|
|
ImageTransformConfig,
|
|
ImageTransforms,
|
|
ImageTransformsConfig,
|
|
make_transform_from_config,
|
|
)
|
|
from lerobot.common.utils.utils import seeded_context
|
|
|
|
ARTIFACT_DIR = Path("tests/data/save_image_transforms_to_safetensors")
|
|
DATASET_REPO_ID = "lerobot/aloha_mobile_shrimp"
|
|
|
|
|
|
def save_default_config_transform(original_frame: torch.Tensor, output_dir: Path):
|
|
cfg = ImageTransformsConfig(enable=True)
|
|
default_tf = ImageTransforms(cfg)
|
|
|
|
with seeded_context(1337):
|
|
img_tf = default_tf(original_frame)
|
|
|
|
save_file({"default": img_tf}, output_dir / "default_transforms.safetensors")
|
|
|
|
|
|
def save_single_transforms(original_frame: torch.Tensor, output_dir: Path):
|
|
transforms = {
|
|
("ColorJitter", "brightness", [(0.5, 0.5), (2.0, 2.0)]),
|
|
("ColorJitter", "contrast", [(0.5, 0.5), (2.0, 2.0)]),
|
|
("ColorJitter", "saturation", [(0.5, 0.5), (2.0, 2.0)]),
|
|
("ColorJitter", "hue", [(-0.25, -0.25), (0.25, 0.25)]),
|
|
("SharpnessJitter", "sharpness", [(0.5, 0.5), (2.0, 2.0)]),
|
|
}
|
|
|
|
frames = {"original_frame": original_frame}
|
|
for tf_type, tf_name, min_max_values in transforms.items():
|
|
for min_max in min_max_values:
|
|
tf_cfg = ImageTransformConfig(type=tf_type, kwargs={tf_name: min_max})
|
|
tf = make_transform_from_config(tf_cfg)
|
|
key = f"{tf_name}_{min_max[0]}_{min_max[1]}"
|
|
frames[key] = tf(original_frame)
|
|
|
|
save_file(frames, output_dir / "single_transforms.safetensors")
|
|
|
|
|
|
def main():
|
|
dataset = LeRobotDataset(DATASET_REPO_ID, episodes=[0], image_transforms=None)
|
|
output_dir = Path(ARTIFACT_DIR)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
original_frame = dataset[0][dataset.meta.camera_keys[0]]
|
|
|
|
save_single_transforms(original_frame, output_dir)
|
|
save_default_config_transform(original_frame, output_dir)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|