diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index d916ec7bb..96a65cb18 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -49,9 +49,11 @@ from .sampler import EpisodeAwareSampler from .streaming_dataset import StreamingLeRobotDataset from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card from .video_utils import ( + DepthEncoderConfig, VideoEncoderConfig, VideoEncodingManager, camera_encoder_defaults, + depth_encoder_defaults, ) # NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.) @@ -67,9 +69,11 @@ __all__ = [ "LeRobotDatasetMetadata", "MultiLeRobotDataset", "StreamingLeRobotDataset", + "DepthEncoderConfig", "VideoEncoderConfig", "VideoEncodingManager", "camera_encoder_defaults", + "depth_encoder_defaults", "add_features", "aggregate_datasets", "aggregate_pipeline_dataset_features", diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index b571bd5c3..beaaeb576 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -104,10 +104,12 @@ from lerobot.common.control_utils import ( from lerobot.configs import parser from lerobot.configs.dataset import DatasetRecordConfig from lerobot.datasets import ( + DepthEncoderConfig, LeRobotDataset, VideoEncodingManager, aggregate_pipeline_dataset_features, create_initial_features, + depth_encoder_defaults, safe_stop_image_writer, ) from lerobot.processor import ( @@ -399,6 +401,7 @@ def record( root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, camera_encoder_config=cfg.dataset.camera_encoder_config, + depth_encoder_config=cfg.dataset.depth_encoder_config, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -428,6 +431,7 @@ def record( image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras), batch_encoding_size=cfg.dataset.video_encoding_batch_size, camera_encoder_config=cfg.dataset.camera_encoder_config, + depth_encoder_config=cfg.dataset.depth_encoder_config, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,