mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 16:49:55 +00:00
Michel Aractingi
65b0e73ae4
* add example scripts to load behaviour 1k data in `load_behaviour1k_dataset.py`
131 lines
4.2 KiB
Python
131 lines
4.2 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Test script to verify BEHAVIOR-1K dataset loading with v3.0 wrapper.
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
|
|
from behavior_lerobot_dataset_v3 import BehaviorLeRobotDatasetV3
|
|
|
|
from lerobot.utils.utils import init_logging
|
|
|
|
init_logging()
|
|
|
|
|
|
def load_behavior1k_dataset(repo_id, root):
|
|
"""Test basic dataset loading."""
|
|
logging.info("=" * 80)
|
|
logging.info("Testing BEHAVIOR-1K dataset loading")
|
|
logging.info("=" * 80)
|
|
|
|
logging.info(f"\n1. Loading dataset with repo_id: {repo_id}")
|
|
dataset = BehaviorLeRobotDatasetV3(
|
|
repo_id=repo_id,
|
|
root=root,
|
|
modalities=["rgb"],
|
|
cameras=["head"],
|
|
chunk_streaming_using_keyframe=False,
|
|
check_timestamp_sync=False,
|
|
)
|
|
|
|
logging.info("\n2. Dataset loaded successfully!")
|
|
logging.info(f" - Number of episodes: {dataset.num_episodes}")
|
|
logging.info(f" - Number of frames: {dataset.num_frames}")
|
|
logging.info(f" - FPS: {dataset.fps}")
|
|
logging.info(f" - Features: {list(dataset.features)}")
|
|
|
|
return dataset
|
|
|
|
|
|
def load_behavior1k_dataset_with_multiple_modalities(repo_id, root):
|
|
"""Test loading multiple modalities and cameras."""
|
|
logging.info("\n" + "=" * 80)
|
|
logging.info("Testing multi-modality loading with repo_id: {repo_id}")
|
|
logging.info("=" * 80)
|
|
|
|
logging.info(f"\n1. Loading dataset with RGB + Depth with repo_id: {repo_id}")
|
|
dataset = BehaviorLeRobotDatasetV3(
|
|
repo_id=repo_id,
|
|
root=root,
|
|
modalities=["rgb", "depth"],
|
|
cameras=["head", "left_wrist", "right_wrist"],
|
|
chunk_streaming_using_keyframe=False,
|
|
check_timestamp_sync=False,
|
|
video_backend="pyav",
|
|
)
|
|
|
|
logging.info(f"\n2. Dataset loaded with modalities: {list(dataset.features)}")
|
|
logging.info(f" - Total features: {len(dataset.features)}")
|
|
|
|
rgb_keys = [k for k in dataset.features if "rgb" in k]
|
|
depth_keys = [k for k in dataset.features if "depth" in k]
|
|
logging.info(f" - RGB features: {rgb_keys}")
|
|
logging.info(f" - Depth features: {depth_keys}")
|
|
|
|
logging.info("\n3. SUCCESS! Multi-modality loading works.")
|
|
|
|
return dataset
|
|
|
|
|
|
def stream_behavior1k_dataset(repo_id, root):
|
|
"""Test chunk streaming mode."""
|
|
logging.info("\n" + "=" * 80)
|
|
logging.info("Testing chunk streaming mode")
|
|
logging.info("=" * 80)
|
|
|
|
logging.info("\n1. Loading dataset with chunk streaming...")
|
|
dataset = BehaviorLeRobotDatasetV3(
|
|
repo_id=repo_id,
|
|
root=root,
|
|
modalities=["rgb"],
|
|
cameras=["head"],
|
|
chunk_streaming_using_keyframe=True,
|
|
shuffle=True,
|
|
seed=42,
|
|
check_timestamp_sync=False,
|
|
)
|
|
|
|
logging.info("\n2. Dataset loaded in streaming mode")
|
|
logging.info(f" - Number of chunks: {len(dataset.chunks)}")
|
|
logging.info(f" - First chunk range: {dataset.chunks[0]}")
|
|
|
|
logging.info("\n3. Testing frame access in streaming mode...")
|
|
for i in range(min(3, len(dataset))):
|
|
frame = dataset[i]
|
|
logging.info(
|
|
f" - Frame {i}: episode_index={frame['episode_index'].item()}, "
|
|
f"task_index={frame['task_index'].item()}"
|
|
)
|
|
|
|
logging.info("\n4. SUCCESS! Chunk streaming works.")
|
|
|
|
return dataset
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--repo-id", type=str, default=None)
|
|
parser.add_argument("--root", type=str, default=None)
|
|
|
|
args = parser.parse_args()
|
|
|
|
load_behavior1k_dataset(args.repo_id, args.root)
|
|
load_behavior1k_dataset_with_multiple_modalities(args.repo_id, args.root)
|
|
stream_behavior1k_dataset(args.repo_id, args.root)
|