""" This file contains all obsolete download scripts. They are centralized here to not have to load useless dependencies when using datasets. """ import io import json import pickle import shutil from pathlib import Path import einops import h5py import numpy as np import torch import tqdm from datasets import Dataset, Features, Image, Sequence, Value from huggingface_hub import HfApi from PIL import Image as PILImage from safetensors.torch import save_file from lerobot.common.datasets.utils import compute_stats, flatten_dict, hf_transform_to_torch def download_and_upload(root, revision, dataset_id): # TODO(rcadene, adilzouitine): add community_id/user_id (e.g. "lerobot", "cadene") or repo_id (e.g. "lerobot/pusht") if "pusht" in dataset_id: download_and_upload_pusht(root, revision, dataset_id) elif "xarm" in dataset_id: download_and_upload_xarm(root, revision, dataset_id) elif "aloha" in dataset_id: download_and_upload_aloha(root, revision, dataset_id) elif "umi" in dataset_id: download_and_upload_umi(root, revision, dataset_id) else: raise ValueError(dataset_id) def concatenate_episodes(ep_dicts): data_dict = {} keys = ep_dicts[0].keys() for key in keys: if torch.is_tensor(ep_dicts[0][key][0]): data_dict[key] = torch.cat([ep_dict[key] for ep_dict in ep_dicts]) else: if key not in data_dict: data_dict[key] = [] for ep_dict in ep_dicts: for x in ep_dict[key]: data_dict[key].append(x) total_frames = data_dict["frame_index"].shape[0] data_dict["index"] = torch.arange(0, total_frames, 1) return data_dict def download_and_extract_zip(url: str, destination_folder: Path) -> bool: import zipfile import requests print(f"downloading from {url}") response = requests.get(url, stream=True) if response.status_code == 200: total_size = int(response.headers.get("content-length", 0)) progress_bar = tqdm.tqdm(total=total_size, unit="B", unit_scale=True) zip_file = io.BytesIO() for chunk in response.iter_content(chunk_size=1024): if chunk: zip_file.write(chunk) progress_bar.update(len(chunk)) progress_bar.close() zip_file.seek(0) with zipfile.ZipFile(zip_file, "r") as zip_ref: zip_ref.extractall(destination_folder) return True else: return False def push_to_hub(hf_dataset, episode_data_index, info, stats, root, revision, dataset_id): # push to main to indicate latest version hf_dataset.push_to_hub(f"lerobot/{dataset_id}", token=True) # push to version branch hf_dataset.push_to_hub(f"lerobot/{dataset_id}", token=True, revision=revision) # create and store meta_data meta_data_dir = root / dataset_id / "meta_data" meta_data_dir.mkdir(parents=True, exist_ok=True) api = HfApi() # info info_path = meta_data_dir / "info.json" with open(str(info_path), "w") as f: json.dump(info, f, indent=4) api.upload_file( path_or_fileobj=info_path, path_in_repo=str(info_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", ) api.upload_file( path_or_fileobj=info_path, path_in_repo=str(info_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", revision=revision, ) # stats stats_path = meta_data_dir / "stats.safetensors" save_file(flatten_dict(stats), stats_path) api.upload_file( path_or_fileobj=stats_path, path_in_repo=str(stats_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", ) api.upload_file( path_or_fileobj=stats_path, path_in_repo=str(stats_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", revision=revision, ) # episode_data_index episode_data_index = {key: torch.tensor(episode_data_index[key]) for key in episode_data_index} ep_data_idx_path = meta_data_dir / "episode_data_index.safetensors" save_file(episode_data_index, ep_data_idx_path) api.upload_file( path_or_fileobj=ep_data_idx_path, path_in_repo=str(ep_data_idx_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", ) api.upload_file( path_or_fileobj=ep_data_idx_path, path_in_repo=str(ep_data_idx_path).replace(f"{root}/{dataset_id}", ""), repo_id=f"lerobot/{dataset_id}", repo_type="dataset", revision=revision, ) # copy in tests folder, the first episode and the meta_data directory num_items_first_ep = episode_data_index["to"][0] - episode_data_index["from"][0] hf_dataset.select(range(num_items_first_ep)).with_format("torch").save_to_disk( f"tests/data/lerobot/{dataset_id}/train" ) if Path(f"tests/data/lerobot/{dataset_id}/meta_data").exists(): shutil.rmtree(f"tests/data/lerobot/{dataset_id}/meta_data") shutil.copytree(meta_data_dir, f"tests/data/lerobot/{dataset_id}/meta_data") def download_and_upload_pusht(root, revision, dataset_id="pusht", fps=10): try: import pymunk from gym_pusht.envs.pusht import PushTEnv, pymunk_to_shapely from lerobot.common.datasets._diffusion_policy_replay_buffer import ( ReplayBuffer as DiffusionPolicyReplayBuffer, ) except ModuleNotFoundError as e: print("`gym_pusht` is not installed. Please install it with `pip install 'lerobot[gym_pusht]'`") raise e # as define in env success_threshold = 0.95 # 95% coverage, pusht_url = "https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip" pusht_zarr = Path("pusht/pusht_cchi_v7_replay.zarr") root = Path(root) raw_dir = root / f"{dataset_id}_raw" zarr_path = (raw_dir / pusht_zarr).resolve() if not zarr_path.is_dir(): raw_dir.mkdir(parents=True, exist_ok=True) download_and_extract_zip(pusht_url, raw_dir) # load dataset_dict = DiffusionPolicyReplayBuffer.copy_from_path(zarr_path) # , keys=['img', 'state', 'action']) episode_ids = torch.from_numpy(dataset_dict.get_episode_idxs()) num_episodes = dataset_dict.meta["episode_ends"].shape[0] assert len( {dataset_dict[key].shape[0] for key in dataset_dict.keys()} # noqa: SIM118 ), "Some data type dont have the same number of total frames." # TODO: verify that goal pose is expected to be fixed goal_pos_angle = np.array([256, 256, np.pi / 4]) # x, y, theta (in radians) goal_body = PushTEnv.get_goal_pose_body(goal_pos_angle) imgs = torch.from_numpy(dataset_dict["img"]) # b h w c states = torch.from_numpy(dataset_dict["state"]) actions = torch.from_numpy(dataset_dict["action"]) ep_dicts = [] episode_data_index = {"from": [], "to": []} id_from = 0 for episode_id in tqdm.tqdm(range(num_episodes)): id_to = dataset_dict.meta["episode_ends"][episode_id] num_frames = id_to - id_from assert (episode_ids[id_from:id_to] == episode_id).all() image = imgs[id_from:id_to] assert image.min() >= 0.0 assert image.max() <= 255.0 image = image.type(torch.uint8) state = states[id_from:id_to] agent_pos = state[:, :2] block_pos = state[:, 2:4] block_angle = state[:, 4] reward = torch.zeros(num_frames) success = torch.zeros(num_frames, dtype=torch.bool) done = torch.zeros(num_frames, dtype=torch.bool) for i in range(num_frames): space = pymunk.Space() space.gravity = 0, 0 space.damping = 0 # Add walls. walls = [ PushTEnv.add_segment(space, (5, 506), (5, 5), 2), PushTEnv.add_segment(space, (5, 5), (506, 5), 2), PushTEnv.add_segment(space, (506, 5), (506, 506), 2), PushTEnv.add_segment(space, (5, 506), (506, 506), 2), ] space.add(*walls) block_body = PushTEnv.add_tee(space, block_pos[i].tolist(), block_angle[i].item()) goal_geom = pymunk_to_shapely(goal_body, block_body.shapes) block_geom = pymunk_to_shapely(block_body, block_body.shapes) intersection_area = goal_geom.intersection(block_geom).area goal_area = goal_geom.area coverage = intersection_area / goal_area reward[i] = np.clip(coverage / success_threshold, 0, 1) success[i] = coverage > success_threshold # last step of demonstration is considered done done[-1] = True ep_dict = { "observation.image": [PILImage.fromarray(x.numpy()) for x in image], "observation.state": agent_pos, "action": actions[id_from:id_to], "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int), "frame_index": torch.arange(0, num_frames, 1), "timestamp": torch.arange(0, num_frames, 1) / fps, # "next.observation.image": image[1:], # "next.observation.state": agent_pos[1:], # TODO(rcadene): verify that reward and done are aligned with image and agent_pos "next.reward": torch.cat([reward[1:], reward[[-1]]]), "next.done": torch.cat([done[1:], done[[-1]]]), "next.success": torch.cat([success[1:], success[[-1]]]), } ep_dicts.append(ep_dict) episode_data_index["from"].append(id_from) episode_data_index["to"].append(id_from + num_frames) id_from += num_frames data_dict = concatenate_episodes(ep_dicts) features = { "observation.image": Image(), "observation.state": Sequence( length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None) ), "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)), "episode_index": Value(dtype="int64", id=None), "frame_index": Value(dtype="int64", id=None), "timestamp": Value(dtype="float32", id=None), "next.reward": Value(dtype="float32", id=None), "next.done": Value(dtype="bool", id=None), "next.success": Value(dtype="bool", id=None), "index": Value(dtype="int64", id=None), } features = Features(features) hf_dataset = Dataset.from_dict(data_dict, features=features) hf_dataset.set_transform(hf_transform_to_torch) info = { "fps": fps, } stats = compute_stats(hf_dataset) push_to_hub(hf_dataset, episode_data_index, info, stats, root, revision, dataset_id) def download_and_upload_xarm(root, revision, dataset_id, fps=15): root = Path(root) raw_dir = root / "xarm_datasets_raw" if not raw_dir.exists(): import zipfile import gdown raw_dir.mkdir(parents=True, exist_ok=True) # from https://github.com/fyhMer/fowm/blob/main/scripts/download_datasets.py url = "https://drive.google.com/uc?id=1nhxpykGtPDhmQKm-_B8zBSywVRdgeVya" zip_path = raw_dir / "data.zip" gdown.download(url, str(zip_path), quiet=False) print("Extracting...") with zipfile.ZipFile(str(zip_path), "r") as zip_f: for member in zip_f.namelist(): if member.startswith("data/xarm") and member.endswith(".pkl"): print(member) zip_f.extract(member=member) zip_path.unlink() dataset_path = root / f"{dataset_id}" / "buffer.pkl" print(f"Using offline dataset '{dataset_path}'") with open(dataset_path, "rb") as f: dataset_dict = pickle.load(f) ep_dicts = [] episode_data_index = {"from": [], "to": []} id_from = 0 id_to = 0 episode_id = 0 total_frames = dataset_dict["actions"].shape[0] for i in tqdm.tqdm(range(total_frames)): id_to += 1 if not dataset_dict["dones"][i]: continue num_frames = id_to - id_from image = torch.tensor(dataset_dict["observations"]["rgb"][id_from:id_to]) image = einops.rearrange(image, "b c h w -> b h w c") state = torch.tensor(dataset_dict["observations"]["state"][id_from:id_to]) action = torch.tensor(dataset_dict["actions"][id_from:id_to]) # TODO(rcadene): we have a missing last frame which is the observation when the env is done # it is critical to have this frame for tdmpc to predict a "done observation/state" # next_image = torch.tensor(dataset_dict["next_observations"]["rgb"][id_from:id_to]) # next_state = torch.tensor(dataset_dict["next_observations"]["state"][id_from:id_to]) next_reward = torch.tensor(dataset_dict["rewards"][id_from:id_to]) next_done = torch.tensor(dataset_dict["dones"][id_from:id_to]) ep_dict = { "observation.image": [PILImage.fromarray(x.numpy()) for x in image], "observation.state": state, "action": action, "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int), "frame_index": torch.arange(0, num_frames, 1), "timestamp": torch.arange(0, num_frames, 1) / fps, # "next.observation.image": next_image, # "next.observation.state": next_state, "next.reward": next_reward, "next.done": next_done, } ep_dicts.append(ep_dict) episode_data_index["from"].append(id_from) episode_data_index["to"].append(id_from + num_frames) id_from = id_to episode_id += 1 data_dict = concatenate_episodes(ep_dicts) features = { "observation.image": Image(), "observation.state": Sequence( length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None) ), "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)), "episode_index": Value(dtype="int64", id=None), "frame_index": Value(dtype="int64", id=None), "timestamp": Value(dtype="float32", id=None), "next.reward": Value(dtype="float32", id=None), "next.done": Value(dtype="bool", id=None), #'next.success': Value(dtype='bool', id=None), "index": Value(dtype="int64", id=None), } features = Features(features) hf_dataset = Dataset.from_dict(data_dict, features=features) hf_dataset.set_transform(hf_transform_to_torch) info = { "fps": fps, } stats = compute_stats(hf_dataset) push_to_hub(hf_dataset, episode_data_index, info, stats, root, revision, dataset_id) def download_and_upload_aloha(root, revision, dataset_id, fps=50): folder_urls = { "aloha_sim_insertion_human": "https://drive.google.com/drive/folders/1RgyD0JgTX30H4IM5XZn8I3zSV_mr8pyF", "aloha_sim_insertion_scripted": "https://drive.google.com/drive/folders/1TsojQQSXtHEoGnqgJ3gmpPQR2DPLtS2N", "aloha_sim_transfer_cube_human": "https://drive.google.com/drive/folders/1sc-E4QYW7A0o23m1u2VWNGVq5smAsfCo", "aloha_sim_transfer_cube_scripted": "https://drive.google.com/drive/folders/1aRyoOhQwxhyt1J8XgEig4s6kzaw__LXj", } ep48_urls = { "aloha_sim_insertion_human": "https://drive.google.com/file/d/18Cudl6nikDtgRolea7je8iF_gGKzynOP/view?usp=drive_link", "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/1wfMSZ24oOh5KR_0aaP3Cnu_c4ZCveduB/view?usp=drive_link", "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/18smMymtr8tIxaNUQ61gW6dG50pt3MvGq/view?usp=drive_link", "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1pnGIOd-E4-rhz2P3VxpknMKRZCoKt6eI/view?usp=drive_link", } ep49_urls = { "aloha_sim_insertion_human": "https://drive.google.com/file/d/1C1kZYyROzs-PrLc0SkDgUgMi4-L3lauE/view?usp=drive_link", "aloha_sim_insertion_scripted": "https://drive.google.com/file/d/17EuCUWS6uCCr6yyNzpXdcdE-_TTNCKtf/view?usp=drive_link", "aloha_sim_transfer_cube_human": "https://drive.google.com/file/d/1Nk7l53d9sJoGDBKAOnNrExX5nLacATc6/view?usp=drive_link", "aloha_sim_transfer_cube_scripted": "https://drive.google.com/file/d/1GKReZHrXU73NMiC5zKCq_UtqPVtYq8eo/view?usp=drive_link", } num_episodes = { "aloha_sim_insertion_human": 50, "aloha_sim_insertion_scripted": 50, "aloha_sim_transfer_cube_human": 50, "aloha_sim_transfer_cube_scripted": 50, } episode_len = { "aloha_sim_insertion_human": 500, "aloha_sim_insertion_scripted": 400, "aloha_sim_transfer_cube_human": 400, "aloha_sim_transfer_cube_scripted": 400, } cameras = { "aloha_sim_insertion_human": ["top"], "aloha_sim_insertion_scripted": ["top"], "aloha_sim_transfer_cube_human": ["top"], "aloha_sim_transfer_cube_scripted": ["top"], } root = Path(root) raw_dir = root / f"{dataset_id}_raw" if not raw_dir.is_dir(): import gdown assert dataset_id in folder_urls assert dataset_id in ep48_urls assert dataset_id in ep49_urls raw_dir.mkdir(parents=True, exist_ok=True) gdown.download_folder(folder_urls[dataset_id], output=str(raw_dir)) # because of the 50 files limit per directory, two files episode 48 and 49 were missing gdown.download(ep48_urls[dataset_id], output=str(raw_dir / "episode_48.hdf5"), fuzzy=True) gdown.download(ep49_urls[dataset_id], output=str(raw_dir / "episode_49.hdf5"), fuzzy=True) ep_dicts = [] episode_data_index = {"from": [], "to": []} id_from = 0 for ep_id in tqdm.tqdm(range(num_episodes[dataset_id])): ep_path = raw_dir / f"episode_{ep_id}.hdf5" with h5py.File(ep_path, "r") as ep: num_frames = ep["/action"].shape[0] assert episode_len[dataset_id] == num_frames # last step of demonstration is considered done done = torch.zeros(num_frames, dtype=torch.bool) done[-1] = True state = torch.from_numpy(ep["/observations/qpos"][:]) action = torch.from_numpy(ep["/action"][:]) ep_dict = {} for cam in cameras[dataset_id]: image = torch.from_numpy(ep[f"/observations/images/{cam}"][:]) # b h w c # image = einops.rearrange(image, "b h w c -> b c h w").contiguous() ep_dict[f"observation.images.{cam}"] = [PILImage.fromarray(x.numpy()) for x in image] # ep_dict[f"next.observation.images.{cam}"] = image ep_dict.update( { "observation.state": state, "action": action, "episode_index": torch.tensor([ep_id] * num_frames), "frame_index": torch.arange(0, num_frames, 1), "timestamp": torch.arange(0, num_frames, 1) / fps, # "next.observation.state": state, # TODO(rcadene): compute reward and success # "next.reward": reward, "next.done": done, # "next.success": success, } ) assert isinstance(ep_id, int) ep_dicts.append(ep_dict) episode_data_index["from"].append(id_from) episode_data_index["to"].append(id_from + num_frames) id_from += num_frames data_dict = concatenate_episodes(ep_dicts) features = { "observation.images.top": Image(), "observation.state": Sequence( length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None) ), "action": Sequence(length=data_dict["action"].shape[1], feature=Value(dtype="float32", id=None)), "episode_index": Value(dtype="int64", id=None), "frame_index": Value(dtype="int64", id=None), "timestamp": Value(dtype="float32", id=None), # "next.reward": Value(dtype="float32", id=None), "next.done": Value(dtype="bool", id=None), # "next.success": Value(dtype="bool", id=None), "index": Value(dtype="int64", id=None), } features = Features(features) hf_dataset = Dataset.from_dict(data_dict, features=features) hf_dataset.set_transform(hf_transform_to_torch) info = { "fps": fps, } stats = compute_stats(hf_dataset) push_to_hub(hf_dataset, episode_data_index, info, stats, root, revision, dataset_id) def download_and_upload_umi(root, revision, dataset_id, fps=10): # fps is equal to 10 source:https://arxiv.org/pdf/2402.10329.pdf#table.caption.16 import os import re import shutil from glob import glob import numpy as np import torch import tqdm import zarr from datasets import Dataset, Features, Image, Sequence, Value from lerobot.common.datasets._umi_imagecodecs_numcodecs import register_codecs # NOTE: This is critical otherwise ValueError: codec not available: 'imagecodecs_jpegxl' # will be raised register_codecs() url_cup_in_the_wild = "https://real.stanford.edu/umi/data/zarr_datasets/cup_in_the_wild.zarr.zip" cup_in_the_wild_zarr = Path("umi/cup_in_the_wild/cup_in_the_wild.zarr") root = Path(root) raw_dir = root / f"{dataset_id}_raw" zarr_path = (raw_dir / cup_in_the_wild_zarr).resolve() if not zarr_path.is_dir(): raw_dir.mkdir(parents=True, exist_ok=True) download_and_extract_zip(url_cup_in_the_wild, zarr_path) zarr_data = zarr.open(zarr_path, mode="r") # We process the image data separately because it is too large to fit in memory end_pose = torch.from_numpy(zarr_data["data/robot0_demo_end_pose"][:]) start_pos = torch.from_numpy(zarr_data["data/robot0_demo_start_pose"][:]) eff_pos = torch.from_numpy(zarr_data["data/robot0_eef_pos"][:]) eff_rot_axis_angle = torch.from_numpy(zarr_data["data/robot0_eef_rot_axis_angle"][:]) gripper_width = torch.from_numpy(zarr_data["data/robot0_gripper_width"][:]) states_pos = torch.cat([eff_pos, eff_rot_axis_angle], dim=1) states = torch.cat([states_pos, gripper_width], dim=1) def get_episode_idxs(episode_ends: np.ndarray) -> np.ndarray: # Optimized and simplified version of this function: https://github.com/real-stanford/universal_manipulation_interface/blob/298776ce251f33b6b3185a98d6e7d1f9ad49168b/diffusion_policy/common/replay_buffer.py#L374 from numba import jit @jit(nopython=True) def _get_episode_idxs(episode_ends): result = np.zeros((episode_ends[-1],), dtype=np.int64) start_idx = 0 for episode_number, end_idx in enumerate(episode_ends): result[start_idx:end_idx] = episode_number start_idx = end_idx return result return _get_episode_idxs(episode_ends) episode_ends = zarr_data["meta/episode_ends"][:] num_episodes: int = episode_ends.shape[0] episode_ids = torch.from_numpy(get_episode_idxs(episode_ends)) # We convert it in torch tensor later because the jit function does not support torch tensors episode_ends = torch.from_numpy(episode_ends) ep_dicts = [] episode_data_index = {"from": [], "to": []} id_from = 0 for episode_id in tqdm.tqdm(range(num_episodes)): id_to = episode_ends[episode_id] num_frames = id_to - id_from assert ( episode_ids[id_from:id_to] == episode_id ).all(), f"episode_ids[{id_from}:{id_to}] != {episode_id}" state = states[id_from:id_to] ep_dict = { # observation.image will be filled later "observation.state": state, "episode_index": torch.tensor([episode_id] * num_frames, dtype=torch.int), "frame_index": torch.arange(0, num_frames, 1), "timestamp": torch.arange(0, num_frames, 1) / fps, "episode_data_index_from": torch.tensor([id_from] * num_frames), "episode_data_index_to": torch.tensor([id_from + num_frames] * num_frames), "end_pose": end_pose[id_from:id_to], "start_pos": start_pos[id_from:id_to], "gripper_width": gripper_width[id_from:id_to], } ep_dicts.append(ep_dict) episode_data_index["from"].append(id_from) episode_data_index["to"].append(id_from + num_frames) id_from += num_frames data_dict = concatenate_episodes(ep_dicts) total_frames = id_from data_dict["index"] = torch.arange(0, total_frames, 1) print("Saving images to disk in temporary folder...") # datasets.Image() can take a list of paths to images, so we save the images to a temporary folder # to avoid loading them all in memory _umi_save_images_concurrently(zarr_data, "tmp_umi_images", max_workers=12) print("Saving images to disk in temporary folder... Done") # Sort files by number eg. 1.png, 2.png, 3.png, 9.png, 10.png instead of 1.png, 10.png, 2.png, 3.png, 9.png # to correctly match the images with the data images_path = sorted(glob("tmp_umi_images/*"), key=lambda x: int(re.search(r"(\d+)\.png$", x).group(1))) data_dict["observation.image"] = images_path features = { "observation.image": Image(), "observation.state": Sequence( length=data_dict["observation.state"].shape[1], feature=Value(dtype="float32", id=None) ), "episode_index": Value(dtype="int64", id=None), "frame_index": Value(dtype="int64", id=None), "timestamp": Value(dtype="float32", id=None), "index": Value(dtype="int64", id=None), "episode_data_index_from": Value(dtype="int64", id=None), "episode_data_index_to": Value(dtype="int64", id=None), # `start_pos` and `end_pos` respectively represent the positions of the end-effector # at the beginning and the end of the episode. # `gripper_width` indicates the distance between the grippers, and this value is included # in the state vector, which comprises the concatenation of the end-effector position # and gripper width. "end_pose": Sequence(length=data_dict["end_pose"].shape[1], feature=Value(dtype="float32", id=None)), "start_pos": Sequence( length=data_dict["start_pos"].shape[1], feature=Value(dtype="float32", id=None) ), "gripper_width": Sequence( length=data_dict["gripper_width"].shape[1], feature=Value(dtype="float32", id=None) ), } features = Features(features) hf_dataset = Dataset.from_dict(data_dict, features=features) hf_dataset.set_transform(hf_transform_to_torch) info = { "fps": fps, } stats = compute_stats(hf_dataset) push_to_hub( hf_dataset=hf_dataset, episode_data_index=episode_data_index, info=info, stats=stats, root=root, revision=revision, dataset_id=dataset_id, ) # Cleanup if os.path.exists("tmp_umi_images"): print("Removing temporary images folder") shutil.rmtree("tmp_umi_images") print("Cleanup done") def _umi_clear_folder(folder_path: str): import os """ Clears all the content of the specified folder. Creates the folder if it does not exist. Args: folder_path (str): Path to the folder to clear. Examples: >>> import os >>> os.makedirs('example_folder', exist_ok=True) >>> with open('example_folder/temp_file.txt', 'w') as f: ... f.write('example') >>> clear_folder('example_folder') >>> os.listdir('example_folder') [] """ if os.path.exists(folder_path): for filename in os.listdir(folder_path): file_path = os.path.join(folder_path, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as e: print(f"Failed to delete {file_path}. Reason: {e}") else: os.makedirs(folder_path) def _umi_save_image(img_array: np.array, i: int, folder_path: str): import os """ Saves a single image to the specified folder. Args: img_array (ndarray): The numpy array of the image. i (int): Index of the image, used for naming. folder_path (str): Path to the folder where the image will be saved. """ img = PILImage.fromarray(img_array) img_format = "PNG" if img_array.dtype == np.uint8 else "JPEG" img.save(os.path.join(folder_path, f"{i}.{img_format.lower()}"), quality=100) def _umi_save_images_concurrently(zarr_data: dict, folder_path: str, max_workers: int = 4): from concurrent.futures import ThreadPoolExecutor """ Saves images from the zarr_data to the specified folder using multithreading. Args: zarr_data (dict): A dictionary containing image data in an array format. folder_path (str): Path to the folder where images will be saved. max_workers (int): The maximum number of threads to use for saving images. """ num_images = len(zarr_data["data/camera0_rgb"]) _umi_clear_folder(folder_path) # Clear or create folder first with ThreadPoolExecutor(max_workers=max_workers) as executor: [ executor.submit(_umi_save_image, zarr_data["data/camera0_rgb"][i], i, folder_path) for i in range(num_images) ] if __name__ == "__main__": root = "data" revision = "v1.1" dataset_ids = [ "pusht", "xarm_lift_medium", "xarm_lift_medium_replay", "xarm_push_medium", "xarm_push_medium_replay", "aloha_sim_insertion_human", "aloha_sim_insertion_scripted", "aloha_sim_transfer_cube_human", "aloha_sim_transfer_cube_scripted", "umi_cup_in_the_wild", ] for dataset_id in dataset_ids: download_and_upload(root, revision, dataset_id)