lerobot/download_and_upload_dataset.py

"""
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)