mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
Merge branch 'main' into chore/bump_transformers_v5
This commit is contained in:
@@ -170,13 +170,13 @@ Once you can drive the robot well, you can start recording data to train AI mode
|
|||||||
We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens):
|
We use Hugging Face to store your data online. First, log in with your token from [Hugging Face settings](https://huggingface.co/settings/tokens):
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Store your Hugging Face username:
|
Store your Hugging Face username:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HF_USER=$(huggingface-cli whoami | head -n 1)
|
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}')
|
||||||
echo $HF_USER
|
echo $HF_USER
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -155,10 +155,10 @@ Upload your repository to Hugging Face:
|
|||||||
pip install huggingface_hub
|
pip install huggingface_hub
|
||||||
|
|
||||||
# Login to Hugging Face
|
# Login to Hugging Face
|
||||||
huggingface-cli login
|
hf auth login
|
||||||
|
|
||||||
# Create a new repository
|
# Create a new repository
|
||||||
huggingface-cli repo create my-custom-env --type space --org my-org
|
hf repo create my-org/my-custom-env
|
||||||
|
|
||||||
# Initialize git and push
|
# Initialize git and push
|
||||||
git init
|
git init
|
||||||
|
|||||||
@@ -159,7 +159,7 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
|
|||||||
Add your token to the CLI by running this command:
|
Add your token to the CLI by running this command:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Then store your Hugging Face repository name in a variable:
|
Then store your Hugging Face repository name in a variable:
|
||||||
@@ -327,7 +327,7 @@ You can look for other LeRobot datasets on the hub by searching for `LeRobot` [t
|
|||||||
You can also push your local dataset to the Hub manually, running:
|
You can also push your local dataset to the Hub manually, running:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset
|
hf upload ${HF_USER}/record-test ~/.cache/huggingface/lerobot/{repo-id} --repo-type dataset
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Record function
|
#### Record function
|
||||||
@@ -491,7 +491,7 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
|
|||||||
Once training is done, upload the latest checkpoint with:
|
Once training is done, upload the latest checkpoint with:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli upload ${HF_USER}/act_so101_test \
|
hf upload ${HF_USER}/act_so101_test \
|
||||||
outputs/train/act_so101_test/checkpoints/last/pretrained_model
|
outputs/train/act_so101_test/checkpoints/last/pretrained_model
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -499,7 +499,7 @@ You can also upload intermediate checkpoints with:
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
CKPT=010000
|
CKPT=010000
|
||||||
huggingface-cli upload ${HF_USER}/act_so101_test${CKPT} \
|
hf upload ${HF_USER}/act_so101_test${CKPT} \
|
||||||
outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
|
outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -279,13 +279,13 @@ We use the Hugging Face hub features for uploading your dataset. If you haven't
|
|||||||
Add your token to the CLI by running this command:
|
Add your token to the CLI by running this command:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
huggingface-cli login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
hf auth login --token ${HUGGINGFACE_TOKEN} --add-to-git-credential
|
||||||
```
|
```
|
||||||
|
|
||||||
Then store your Hugging Face repository name in a variable:
|
Then store your Hugging Face repository name in a variable:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HF_USER=$(huggingface-cli whoami | head -n 1)
|
HF_USER=$(hf auth whoami | awk -F': *' 'NR==1 {print $2}')
|
||||||
echo $HF_USER
|
echo $HF_USER
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -89,8 +89,8 @@ def delete_episodes(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
episode_indices: List of episode indices to delete.
|
episode_indices: List of episode indices to delete.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
"""
|
"""
|
||||||
if not episode_indices:
|
if not episode_indices:
|
||||||
raise ValueError("No episodes to delete")
|
raise ValueError("No episodes to delete")
|
||||||
@@ -152,7 +152,7 @@ def split_dataset(
|
|||||||
dataset: The source LeRobotDataset to split.
|
dataset: The source LeRobotDataset to split.
|
||||||
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
||||||
split names to fractions (must sum to <= 1.0).
|
split names to fractions (must sum to <= 1.0).
|
||||||
output_dir: Base directory for output datasets. If None, uses default location.
|
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
Split by specific episodes
|
Split by specific episodes
|
||||||
@@ -243,8 +243,8 @@ def merge_datasets(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
datasets: List of LeRobotDatasets to merge.
|
datasets: List of LeRobotDatasets to merge.
|
||||||
output_repo_id: Repository ID for the merged dataset.
|
output_repo_id: Merged dataset identifier.
|
||||||
output_dir: Directory to save the merged dataset. If None, uses default location.
|
output_dir: Root directory where the merged dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/output_repo_id.
|
||||||
"""
|
"""
|
||||||
if not datasets:
|
if not datasets:
|
||||||
raise ValueError("No datasets to merge")
|
raise ValueError("No datasets to merge")
|
||||||
@@ -288,8 +288,8 @@ def modify_features(
|
|||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
|
add_features: Optional dict mapping feature names to (feature_values, feature_info) tuples.
|
||||||
remove_features: Optional feature name(s) to remove. Can be a single string or list.
|
remove_features: Optional feature name(s) to remove. Can be a single string or list.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with features modified.
|
New dataset with features modified.
|
||||||
@@ -390,8 +390,8 @@ def add_features(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
|
features: Dictionary mapping feature names to (feature_values, feature_info) tuples.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with all features added.
|
New dataset with all features added.
|
||||||
@@ -427,8 +427,8 @@ def remove_feature(
|
|||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobotDataset.
|
dataset: The source LeRobotDataset.
|
||||||
feature_names: Name(s) of features to remove. Can be a single string or list.
|
feature_names: Name(s) of features to remove. Can be a single string or list.
|
||||||
output_dir: Directory to save the new dataset. If None, uses default location.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset. If None, appends "_modified" to original.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
New dataset with features removed.
|
New dataset with features removed.
|
||||||
@@ -1529,7 +1529,7 @@ def modify_tasks(
|
|||||||
|
|
||||||
def convert_image_to_video_dataset(
|
def convert_image_to_video_dataset(
|
||||||
dataset: LeRobotDataset,
|
dataset: LeRobotDataset,
|
||||||
output_dir: Path,
|
output_dir: Path | None = None,
|
||||||
repo_id: str | None = None,
|
repo_id: str | None = None,
|
||||||
vcodec: str = "libsvtav1",
|
vcodec: str = "libsvtav1",
|
||||||
pix_fmt: str = "yuv420p",
|
pix_fmt: str = "yuv420p",
|
||||||
@@ -1548,8 +1548,8 @@ def convert_image_to_video_dataset(
|
|||||||
|
|
||||||
Args:
|
Args:
|
||||||
dataset: The source LeRobot dataset with images
|
dataset: The source LeRobot dataset with images
|
||||||
output_dir: Directory to save the new video dataset
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Repository ID for the new dataset (default: original_id + "_video")
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
vcodec: Video codec (default: libsvtav1)
|
vcodec: Video codec (default: libsvtav1)
|
||||||
pix_fmt: Pixel format (default: yuv420p)
|
pix_fmt: Pixel format (default: yuv420p)
|
||||||
g: Group of pictures size (default: 2)
|
g: Group of pictures size (default: 2)
|
||||||
@@ -1600,6 +1600,7 @@ def convert_image_to_video_dataset(
|
|||||||
# Video info will be updated after episodes are encoded
|
# Video info will be updated after episodes are encoded
|
||||||
|
|
||||||
# Create new metadata for video dataset
|
# Create new metadata for video dataset
|
||||||
|
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
|
||||||
new_meta = LeRobotDatasetMetadata.create(
|
new_meta = LeRobotDatasetMetadata.create(
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
fps=dataset.meta.fps,
|
fps=dataset.meta.fps,
|
||||||
|
|||||||
@@ -36,8 +36,11 @@ Convert a local dataset (works in place):
|
|||||||
```bash
|
```bash
|
||||||
python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
|
python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
|
||||||
--repo-id=lerobot/pusht \
|
--repo-id=lerobot/pusht \
|
||||||
--root=/path/to/local/dataset/directory
|
--root=/path/to/local/dataset/directory \
|
||||||
--push-to-hub=false
|
--push-to-hub=false
|
||||||
|
|
||||||
|
N.B. Path semantics (v2): --root is the exact dataset folder containing
|
||||||
|
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
|
||||||
```
|
```
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -201,7 +204,6 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
|
|
||||||
image_keys = get_image_keys(root)
|
image_keys = get_image_keys(root)
|
||||||
|
|
||||||
ep_idx = 0
|
|
||||||
chunk_idx = 0
|
chunk_idx = 0
|
||||||
file_idx = 0
|
file_idx = 0
|
||||||
size_in_mb = 0
|
size_in_mb = 0
|
||||||
@@ -211,9 +213,24 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
|
|
||||||
logging.info(f"Converting data files from {len(ep_paths)} episodes")
|
logging.info(f"Converting data files from {len(ep_paths)} episodes")
|
||||||
|
|
||||||
for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
|
for ep_idx, ep_path in enumerate(tqdm.tqdm(ep_paths, desc="convert data files")):
|
||||||
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
|
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
|
||||||
ep_num_frames = get_parquet_num_frames(ep_path)
|
ep_num_frames = get_parquet_num_frames(ep_path)
|
||||||
|
|
||||||
|
# Check if we need to start a new file BEFORE creating metadata
|
||||||
|
if size_in_mb + ep_size_in_mb >= data_file_size_in_mb and len(paths_to_cat) > 0:
|
||||||
|
# Write the accumulated data files
|
||||||
|
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
|
||||||
|
|
||||||
|
# Move to next file
|
||||||
|
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
|
||||||
|
|
||||||
|
# Reset for the next file
|
||||||
|
size_in_mb = 0
|
||||||
|
num_frames += ep_num_frames # Still need to accumulate total frames
|
||||||
|
paths_to_cat = []
|
||||||
|
|
||||||
|
# Now create metadata with correct chunk/file indices
|
||||||
ep_metadata = {
|
ep_metadata = {
|
||||||
"episode_index": ep_idx,
|
"episode_index": ep_idx,
|
||||||
"data/chunk_index": chunk_idx,
|
"data/chunk_index": chunk_idx,
|
||||||
@@ -224,20 +241,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
|
|||||||
size_in_mb += ep_size_in_mb
|
size_in_mb += ep_size_in_mb
|
||||||
num_frames += ep_num_frames
|
num_frames += ep_num_frames
|
||||||
episodes_metadata.append(ep_metadata)
|
episodes_metadata.append(ep_metadata)
|
||||||
ep_idx += 1
|
paths_to_cat.append(ep_path)
|
||||||
|
|
||||||
if size_in_mb < data_file_size_in_mb:
|
|
||||||
paths_to_cat.append(ep_path)
|
|
||||||
continue
|
|
||||||
|
|
||||||
if paths_to_cat:
|
|
||||||
concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys)
|
|
||||||
|
|
||||||
# Reset for the next file
|
|
||||||
size_in_mb = ep_size_in_mb
|
|
||||||
paths_to_cat = [ep_path]
|
|
||||||
|
|
||||||
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
|
|
||||||
|
|
||||||
# Write remaining data if any
|
# Write remaining data if any
|
||||||
if paths_to_cat:
|
if paths_to_cat:
|
||||||
@@ -469,7 +473,7 @@ def convert_dataset(
|
|||||||
|
|
||||||
# Set root based on whether local dataset path is provided
|
# Set root based on whether local dataset path is provided
|
||||||
use_local_dataset = False
|
use_local_dataset = False
|
||||||
root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id
|
root = HF_LEROBOT_HOME / repo_id if root is None else Path(root)
|
||||||
if root.exists():
|
if root.exists():
|
||||||
validate_local_dataset_version(root)
|
validate_local_dataset_version(root)
|
||||||
use_local_dataset = True
|
use_local_dataset = True
|
||||||
@@ -553,7 +557,7 @@ if __name__ == "__main__":
|
|||||||
"--root",
|
"--root",
|
||||||
type=str,
|
type=str,
|
||||||
default=None,
|
default=None,
|
||||||
help="Local directory to use for downloading/writing the dataset.",
|
help="Local directory to use for downloading/writing the dataset. Defaults to $HF_LEROBOT_HOME/repo_id.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--push-to-hub",
|
"--push-to-hub",
|
||||||
|
|||||||
@@ -132,10 +132,13 @@ def visualize_dataset(
|
|||||||
|
|
||||||
logging.info("Logging to Rerun")
|
logging.info("Logging to Rerun")
|
||||||
|
|
||||||
|
first_index = None
|
||||||
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
|
for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
|
||||||
|
if first_index is None:
|
||||||
|
first_index = batch["index"][0].item()
|
||||||
# iterate over the batch
|
# iterate over the batch
|
||||||
for i in range(len(batch["index"])):
|
for i in range(len(batch["index"])):
|
||||||
rr.set_time("frame_index", sequence=batch["frame_index"][i].item())
|
rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index)
|
||||||
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
|
rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
|
||||||
|
|
||||||
# display each camera image
|
# display each camera image
|
||||||
|
|||||||
@@ -21,6 +21,9 @@ This script allows you to delete episodes, split datasets, merge datasets,
|
|||||||
remove features, modify tasks, and convert image datasets to video format.
|
remove features, modify tasks, and convert image datasets to video format.
|
||||||
When new_repo_id is specified, creates a new dataset.
|
When new_repo_id is specified, creates a new dataset.
|
||||||
|
|
||||||
|
Path semantics (v2): --root and --new_root are exact dataset folders containing
|
||||||
|
meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}.
|
||||||
|
|
||||||
Usage Examples:
|
Usage Examples:
|
||||||
|
|
||||||
Delete episodes 0, 2, and 5 from a dataset:
|
Delete episodes 0, 2, and 5 from a dataset:
|
||||||
@@ -29,16 +32,31 @@ Delete episodes 0, 2, and 5 from a dataset:
|
|||||||
--operation.type delete_episodes \
|
--operation.type delete_episodes \
|
||||||
--operation.episode_indices "[0, 2, 5]"
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
Delete episodes and save to a new dataset:
|
Delete episodes from a local dataset at a specific path:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--new_repo_id lerobot/pusht_filtered \
|
--root /path/to/pusht \
|
||||||
--operation.type delete_episodes \
|
--operation.type delete_episodes \
|
||||||
--operation.episode_indices "[0, 2, 5]"
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
Split dataset by fractions:
|
Delete episodes and save to a new dataset at a specific path and with a new repo_id:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
|
--new_repo_id lerobot/pusht_filtered \
|
||||||
|
--new_root /path/to/pusht_filtered \
|
||||||
|
--operation.type delete_episodes \
|
||||||
|
--operation.episode_indices "[0, 2, 5]"
|
||||||
|
|
||||||
|
Split dataset by fractions (pusht_train, pusht_val):
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--repo_id lerobot/pusht \
|
||||||
|
--operation.type split \
|
||||||
|
--operation.splits '{"train": 0.8, "val": 0.2}'
|
||||||
|
|
||||||
|
Split dataset by fractions and save split datasets to a specific folder (base_folder/train, base_folder/val):
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--repo_id lerobot/pusht \
|
||||||
|
--new_root /path/to/base_folder \
|
||||||
--operation.type split \
|
--operation.type split \
|
||||||
--operation.splits '{"train": 0.8, "val": 0.2}'
|
--operation.splits '{"train": 0.8, "val": 0.2}'
|
||||||
|
|
||||||
@@ -56,15 +74,29 @@ Split into more than two splits:
|
|||||||
|
|
||||||
Merge multiple datasets:
|
Merge multiple datasets:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht_merged \
|
--new_repo_id lerobot/pusht_merged \
|
||||||
--operation.type merge \
|
--operation.type merge \
|
||||||
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
||||||
|
|
||||||
|
Merge multiple datasets to a specific output path:
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--new_repo_id lerobot/pusht_merged \
|
||||||
|
--new_root /path/to/pusht_merged \
|
||||||
|
--operation.type merge \
|
||||||
|
--operation.repo_ids "['lerobot/pusht_train', 'lerobot/pusht_val']"
|
||||||
|
|
||||||
|
Merge multiple datasets from a list of local dataset paths:
|
||||||
|
lerobot-edit-dataset \
|
||||||
|
--new_repo_id lerobot/pusht_merged \
|
||||||
|
--operation.type merge \
|
||||||
|
--operation.repo_ids "['pusht_train', 'pusht_val']" \
|
||||||
|
--operation.roots "['/path/to/pusht_train', '/path/to/pusht_val']"
|
||||||
|
|
||||||
Remove camera feature:
|
Remove camera feature:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--operation.type remove_feature \
|
--operation.type remove_feature \
|
||||||
--operation.feature_names "['observation.images.top']"
|
--operation.feature_names "['observation.image']"
|
||||||
|
|
||||||
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
|
Modify tasks - set a single task for all episodes (WARNING: modifies in-place):
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -88,8 +120,8 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
|
|||||||
Convert image dataset to video format and save locally:
|
Convert image dataset to video format and save locally:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
--repo_id lerobot/pusht_image \
|
--repo_id lerobot/pusht_image \
|
||||||
--operation.type convert_image_to_video \
|
--new_root /path/to/output/pusht_video \
|
||||||
--operation.output_dir /path/to/output/pusht_video
|
--operation.type convert_image_to_video
|
||||||
|
|
||||||
Convert image dataset to video format and save with new repo_id:
|
Convert image dataset to video format and save with new repo_id:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -167,6 +199,7 @@ class SplitConfig(OperationConfig):
|
|||||||
@dataclass
|
@dataclass
|
||||||
class MergeConfig(OperationConfig):
|
class MergeConfig(OperationConfig):
|
||||||
repo_ids: list[str] | None = None
|
repo_ids: list[str] | None = None
|
||||||
|
roots: list[str] | None = None
|
||||||
|
|
||||||
|
|
||||||
@OperationConfig.register_subclass("remove_feature")
|
@OperationConfig.register_subclass("remove_feature")
|
||||||
@@ -200,36 +233,46 @@ class ConvertImageToVideoConfig(OperationConfig):
|
|||||||
@OperationConfig.register_subclass("info")
|
@OperationConfig.register_subclass("info")
|
||||||
@dataclass
|
@dataclass
|
||||||
class InfoConfig(OperationConfig):
|
class InfoConfig(OperationConfig):
|
||||||
type: str = "info"
|
|
||||||
show_features: bool = False
|
show_features: bool = False
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class EditDatasetConfig:
|
class EditDatasetConfig:
|
||||||
repo_id: str
|
# Operation configuration.
|
||||||
operation: OperationConfig
|
operation: OperationConfig
|
||||||
|
# Input dataset identifier. Always required unless for Merge operation.
|
||||||
|
repo_id: str | None = None
|
||||||
|
# Root directory where the input dataset is stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
root: str | None = None
|
root: str | None = None
|
||||||
|
# Edited dataset identifier. When both new_repo_id (resp. new_root) and repo_id (resp. root) are identical, modifications are applied in-place and a backup of the original dataset is created. Required for Merge operation.
|
||||||
new_repo_id: str | None = None
|
new_repo_id: str | None = None
|
||||||
|
# Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/new_repo_id. For Split operation, this is the base directory for the split datasets.
|
||||||
|
new_root: str | None = None
|
||||||
|
# Upload dataset to Hugging Face hub.
|
||||||
push_to_hub: bool = False
|
push_to_hub: bool = False
|
||||||
|
|
||||||
|
|
||||||
def get_output_path(repo_id: str, new_repo_id: str | None, root: Path | None) -> tuple[str, Path]:
|
def get_output_path(
|
||||||
if new_repo_id:
|
repo_id: str,
|
||||||
output_repo_id = new_repo_id
|
new_repo_id: str | None,
|
||||||
output_dir = root / new_repo_id if root else HF_LEROBOT_HOME / new_repo_id
|
root: Path | str | None,
|
||||||
else:
|
new_root: Path | str | None,
|
||||||
output_repo_id = repo_id
|
) -> tuple[str, Path]:
|
||||||
dataset_path = root / repo_id if root else HF_LEROBOT_HOME / repo_id
|
input_path = Path(root) if root else HF_LEROBOT_HOME / repo_id
|
||||||
old_path = Path(str(dataset_path) + "_old")
|
|
||||||
|
|
||||||
if dataset_path.exists():
|
output_repo_id = new_repo_id if new_repo_id else repo_id
|
||||||
if old_path.exists():
|
output_path = Path(new_root) if new_root else HF_LEROBOT_HOME / output_repo_id
|
||||||
shutil.rmtree(old_path)
|
|
||||||
shutil.move(str(dataset_path), str(old_path))
|
|
||||||
|
|
||||||
output_dir = dataset_path
|
# In case of in-place modification, create a backup of the original dataset (if it exists)
|
||||||
|
if output_path == input_path:
|
||||||
|
backup_path = input_path.with_name(input_path.name + "_old")
|
||||||
|
|
||||||
return output_repo_id, output_dir
|
if input_path.exists():
|
||||||
|
if backup_path.exists():
|
||||||
|
shutil.rmtree(backup_path)
|
||||||
|
shutil.move(input_path, backup_path)
|
||||||
|
|
||||||
|
return output_repo_id, output_path
|
||||||
|
|
||||||
|
|
||||||
def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
||||||
@@ -241,11 +284,15 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
|||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
output_repo_id, output_dir = get_output_path(
|
output_repo_id, output_dir = get_output_path(
|
||||||
cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
|
cfg.repo_id,
|
||||||
|
new_repo_id=cfg.new_repo_id,
|
||||||
|
root=cfg.root,
|
||||||
|
new_root=cfg.new_root,
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id is None:
|
# In case of in-place modification, make the dataset point to the backup directory
|
||||||
dataset.root = Path(str(dataset.root) + "_old")
|
if output_dir == dataset.root:
|
||||||
|
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||||
|
|
||||||
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
|
logging.info(f"Deleting episodes {cfg.operation.episode_indices} from {cfg.repo_id}")
|
||||||
new_dataset = delete_episodes(
|
new_dataset = delete_episodes(
|
||||||
@@ -272,19 +319,27 @@ def handle_split(cfg: EditDatasetConfig) -> None:
|
|||||||
"splits dict must be specified with split names as keys and fractions/episode lists as values"
|
"splits dict must be specified with split names as keys and fractions/episode lists as values"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if cfg.new_repo_id is not None:
|
||||||
|
logging.warning(
|
||||||
|
"split uses the original dataset identifier --repo_id to generate split names. The --new_repo_id parameter is ignored."
|
||||||
|
)
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
|
|
||||||
logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
|
logging.info(f"Splitting dataset {cfg.repo_id} with splits: {cfg.operation.splits}")
|
||||||
split_datasets = split_dataset(dataset, splits=cfg.operation.splits)
|
split_datasets = split_dataset(
|
||||||
|
dataset,
|
||||||
|
splits=cfg.operation.splits,
|
||||||
|
output_dir=cfg.new_root,
|
||||||
|
)
|
||||||
|
|
||||||
for split_name, split_ds in split_datasets.items():
|
for split_name, split_ds in split_datasets.items():
|
||||||
split_repo_id = f"{cfg.repo_id}_{split_name}"
|
|
||||||
logging.info(
|
logging.info(
|
||||||
f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
|
f"{split_name}: {split_ds.meta.total_episodes} episodes, {split_ds.meta.total_frames} frames"
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
logging.info(f"Pushing {split_name} split to hub as {split_repo_id}")
|
logging.info(f"Pushing {split_name} split to hub as {split_ds.repo_id}")
|
||||||
LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
|
LeRobotDataset(split_ds.repo_id, root=split_ds.root).push_to_hub()
|
||||||
|
|
||||||
|
|
||||||
@@ -295,18 +350,29 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
|
|||||||
if not cfg.operation.repo_ids:
|
if not cfg.operation.repo_ids:
|
||||||
raise ValueError("repo_ids must be specified for merge operation")
|
raise ValueError("repo_ids must be specified for merge operation")
|
||||||
|
|
||||||
if not cfg.repo_id:
|
if cfg.repo_id is not None or cfg.root is not None:
|
||||||
raise ValueError("repo_id must be specified as the output repository for merged dataset")
|
logging.warning(
|
||||||
|
"merge uses --new_repo_id and --new_root for the merged dataset. The --repo_id and --root parameters are ignored."
|
||||||
|
)
|
||||||
|
|
||||||
logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
|
if cfg.operation.roots:
|
||||||
datasets = [LeRobotDataset(repo_id, root=cfg.root) for repo_id in cfg.operation.repo_ids]
|
if len(cfg.operation.roots) != len(cfg.operation.repo_ids):
|
||||||
|
raise ValueError("repo_ids and roots must have the same length for merge operation")
|
||||||
|
logging.info(f"Loading {len(cfg.operation.roots)} datasets to merge")
|
||||||
|
datasets = [
|
||||||
|
LeRobotDataset(repo_id=repo_id, root=root)
|
||||||
|
for repo_id, root in zip(cfg.operation.repo_ids, cfg.operation.roots, strict=True)
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
logging.info(f"Loading {len(cfg.operation.repo_ids)} datasets to merge")
|
||||||
|
datasets = [LeRobotDataset(repo_id) for repo_id in cfg.operation.repo_ids]
|
||||||
|
|
||||||
output_dir = Path(cfg.root) / cfg.repo_id if cfg.root else HF_LEROBOT_HOME / cfg.repo_id
|
output_dir = Path(cfg.new_root) if cfg.new_root else HF_LEROBOT_HOME / cfg.new_repo_id
|
||||||
|
|
||||||
logging.info(f"Merging datasets into {cfg.repo_id}")
|
logging.info(f"Merging datasets into {cfg.new_repo_id}")
|
||||||
merged_dataset = merge_datasets(
|
merged_dataset = merge_datasets(
|
||||||
datasets,
|
datasets,
|
||||||
output_repo_id=cfg.repo_id,
|
output_repo_id=cfg.new_repo_id,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -316,7 +382,7 @@ def handle_merge(cfg: EditDatasetConfig) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
logging.info(f"Pushing to hub as {cfg.repo_id}")
|
logging.info(f"Pushing to hub as {cfg.new_repo_id}")
|
||||||
LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
|
LeRobotDataset(merged_dataset.repo_id, root=output_dir).push_to_hub()
|
||||||
|
|
||||||
|
|
||||||
@@ -329,11 +395,15 @@ def handle_remove_feature(cfg: EditDatasetConfig) -> None:
|
|||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
output_repo_id, output_dir = get_output_path(
|
output_repo_id, output_dir = get_output_path(
|
||||||
cfg.repo_id, cfg.new_repo_id, Path(cfg.root) if cfg.root else None
|
cfg.repo_id,
|
||||||
|
new_repo_id=cfg.new_repo_id,
|
||||||
|
root=cfg.root,
|
||||||
|
new_root=cfg.new_root,
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id is None:
|
# In case of in-place modification, make the dataset point to the backup directory
|
||||||
dataset.root = Path(str(dataset.root) + "_old")
|
if output_dir == dataset.root:
|
||||||
|
dataset.root = dataset.root.with_name(dataset.root.name + "_old")
|
||||||
|
|
||||||
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
|
logging.info(f"Removing features {cfg.operation.feature_names} from {cfg.repo_id}")
|
||||||
new_dataset = remove_feature(
|
new_dataset = remove_feature(
|
||||||
@@ -361,9 +431,10 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
|
|||||||
if new_task is None and episode_tasks_raw is None:
|
if new_task is None and episode_tasks_raw is None:
|
||||||
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
|
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
|
||||||
|
|
||||||
# Warn about in-place modification behavior
|
if cfg.new_repo_id is not None or cfg.new_root is not None:
|
||||||
if cfg.new_repo_id is not None:
|
logging.warning(
|
||||||
logging.warning("modify_tasks modifies datasets in-place. The --new_repo_id parameter is ignored.")
|
"modify_tasks modifies datasets in-place. The --new_repo_id and --new_root parameters are ignored."
|
||||||
|
)
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
|
logging.warning(f"Modifying dataset in-place at {dataset.root}. Original data will be overwritten.")
|
||||||
@@ -399,32 +470,30 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None:
|
|||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root)
|
||||||
|
|
||||||
# Determine output directory and repo_id
|
# Determine output directory and repo_id
|
||||||
# Priority: 1) new_repo_id, 2) operation.output_dir, 3) auto-generated name
|
# Priority: 1) new_root, 2) new_repo_id, 3) operation.output_dir, 4) auto-generated name
|
||||||
output_dir_config = getattr(cfg.operation, "output_dir", None)
|
output_dir_config = getattr(cfg.operation, "output_dir", None)
|
||||||
|
if output_dir_config:
|
||||||
|
logging.warning(
|
||||||
|
"--operation.output_dir is deprecated and will be removed in future versions. "
|
||||||
|
"Please use --new_root instead."
|
||||||
|
)
|
||||||
|
|
||||||
if cfg.new_repo_id:
|
if cfg.new_root:
|
||||||
# Use new_repo_id for both local storage and hub push
|
output_dir = Path(cfg.new_root)
|
||||||
|
output_repo_id = cfg.new_repo_id or f"{cfg.repo_id}_video"
|
||||||
|
logging.info(f"Saving to new_root: {output_dir} as {output_repo_id}")
|
||||||
|
elif cfg.new_repo_id:
|
||||||
output_repo_id = cfg.new_repo_id
|
output_repo_id = cfg.new_repo_id
|
||||||
# Place new dataset as a sibling to the original dataset
|
output_dir = HF_LEROBOT_HOME / cfg.new_repo_id
|
||||||
# Get the parent of the actual dataset root (not cfg.root which might be the lerobot cache dir)
|
|
||||||
# Extract just the dataset name (after last slash) for the local directory
|
|
||||||
local_dir_name = cfg.new_repo_id.split("/")[-1]
|
|
||||||
output_dir = dataset.root.parent / local_dir_name
|
|
||||||
logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
|
logging.info(f"Saving to new dataset: {cfg.new_repo_id} at {output_dir}")
|
||||||
elif output_dir_config:
|
elif output_dir_config:
|
||||||
# Use custom output directory for local-only storage
|
|
||||||
output_dir = Path(output_dir_config)
|
output_dir = Path(output_dir_config)
|
||||||
# Extract repo name from output_dir for the dataset
|
|
||||||
output_repo_id = output_dir.name
|
output_repo_id = output_dir.name
|
||||||
logging.info(f"Saving to local directory: {output_dir}")
|
logging.info(f"Saving to local directory: {output_dir} as {output_repo_id}")
|
||||||
else:
|
else:
|
||||||
# Auto-generate name: append "_video" to original repo_id
|
|
||||||
output_repo_id = f"{cfg.repo_id}_video"
|
output_repo_id = f"{cfg.repo_id}_video"
|
||||||
# Place new dataset as a sibling to the original dataset
|
output_dir = HF_LEROBOT_HOME / output_repo_id
|
||||||
# Extract just the dataset name (after last slash) for the local directory
|
logging.info(f"Saving to auto-generated location: {output_dir} as {output_repo_id}")
|
||||||
local_dir_name = output_repo_id.split("/")[-1]
|
|
||||||
output_dir = dataset.root.parent / local_dir_name
|
|
||||||
logging.info(f"Saving to auto-generated location: {output_dir}")
|
|
||||||
|
|
||||||
logging.info(f"Converting dataset {cfg.repo_id} to video format")
|
logging.info(f"Converting dataset {cfg.repo_id} to video format")
|
||||||
|
|
||||||
@@ -499,8 +568,20 @@ def handle_info(cfg: EditDatasetConfig):
|
|||||||
sys.stdout.write(f"{feature_dump_str}\n")
|
sys.stdout.write(f"{feature_dump_str}\n")
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_config(cfg: EditDatasetConfig) -> None:
|
||||||
|
if isinstance(cfg.operation, MergeConfig):
|
||||||
|
if not cfg.new_repo_id:
|
||||||
|
raise ValueError("--new_repo_id is required for merge operation (the merged dataset identifier)")
|
||||||
|
else:
|
||||||
|
if not cfg.repo_id:
|
||||||
|
raise ValueError(
|
||||||
|
f"--repo_id is required for {cfg.operation.type} operation (the input dataset identifier)"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@parser.wrap()
|
@parser.wrap()
|
||||||
def edit_dataset(cfg: EditDatasetConfig) -> None:
|
def edit_dataset(cfg: EditDatasetConfig) -> None:
|
||||||
|
_validate_config(cfg)
|
||||||
operation_type = cfg.operation.type
|
operation_type = cfg.operation.type
|
||||||
|
|
||||||
if operation_type == "delete_episodes":
|
if operation_type == "delete_episodes":
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import (
|
|||||||
OperationConfig,
|
OperationConfig,
|
||||||
RemoveFeatureConfig,
|
RemoveFeatureConfig,
|
||||||
SplitConfig,
|
SplitConfig,
|
||||||
|
_validate_config,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -51,11 +52,23 @@ class TestOperationTypeParsing:
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
||||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
cfg = parse_cfg(
|
||||||
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
||||||
|
)
|
||||||
assert isinstance(cfg.operation, expected_cls), (
|
assert isinstance(cfg.operation, expected_cls), (
|
||||||
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_merge_requires_new_repo_id(self):
|
||||||
|
cfg = parse_cfg(["--operation.type", "merge"])
|
||||||
|
with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
|
||||||
|
_validate_config(cfg)
|
||||||
|
|
||||||
|
def test_non_merge_requires_repo_id(self):
|
||||||
|
cfg = parse_cfg(["--operation.type", "delete_episodes"])
|
||||||
|
with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
|
||||||
|
_validate_config(cfg)
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"type_name, expected_cls",
|
"type_name, expected_cls",
|
||||||
[
|
[
|
||||||
@@ -69,6 +82,8 @@ class TestOperationTypeParsing:
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
||||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
cfg = parse_cfg(
|
||||||
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
||||||
|
)
|
||||||
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
||||||
assert resolved_name == type_name
|
assert resolved_name == type_name
|
||||||
|
|||||||
Reference in New Issue
Block a user