diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index c3fe246cd..b95c7f86b 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -35,7 +35,9 @@ from .video import ( VALID_VIDEO_CODECS, VIDEO_ENCODER_INFO_KEYS, VideoEncoderConfig, + DepthEncoderConfig, camera_encoder_defaults, + depth_encoder_defaults, ) __all__ = [ @@ -53,8 +55,10 @@ __all__ = [ "PreTrainedConfig", "WandBConfig", "VideoEncoderConfig", + "DepthEncoderConfig", # Defaults "camera_encoder_defaults", + "depth_encoder_defaults", # Constants "VALID_VIDEO_CODECS", "VIDEO_ENCODER_INFO_KEYS", diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 8ef6f22c6..b6ddcee83 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -761,6 +761,7 @@ class LeRobotDataset(torch.utils.data.Dataset): video_backend: str | None = None, batch_encoding_size: int = 1, camera_encoder: VideoEncoderConfig | None = None, + depth_encoder: DepthEncoderConfig | None = None, encoder_threads: int | None = None, image_writer_processes: int = 0, image_writer_threads: int = 0, @@ -790,6 +791,8 @@ class LeRobotDataset(torch.utils.data.Dataset): batch-encoding videos. camera_encoder: Video encoder settings for cameras (codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. + depth_encoder: Video encoder settings for depth cameras (codec, quality, etc.). + When ``None``, :func:`~lerobot.configs.depth.depth_encoder_defaults` is used. encoder_threads: Number of encoder threads (global). ``None`` lets the codec decide. image_writer_processes: Subprocesses for async image writing. @@ -836,12 +839,13 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_enc = None if streaming_encoding and len(obj.meta.video_keys) > 0: streaming_enc = cls._build_streaming_encoder( - obj.meta.fps, camera_encoder, encoder_queue_maxsize, encoder_threads + obj.meta.fps, camera_encoder, depth_encoder, encoder_queue_maxsize, encoder_threads ) obj.writer = DatasetWriter( meta=obj.meta, root=obj.root, camera_encoder=camera_encoder, + depth_encoder=depth_encoder, encoder_threads=encoder_threads, batch_encoding_size=batch_encoding_size, streaming_encoder=streaming_enc, diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index c8419cb14..d4a2cdc14 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -399,6 +399,7 @@ def record( root=cfg.dataset.root, batch_encoding_size=cfg.dataset.video_encoding_batch_size, camera_encoder=cfg.dataset.camera_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, @@ -428,6 +429,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=cfg.dataset.camera_encoder, + depth_encoder=cfg.dataset.depth_encoder, encoder_threads=cfg.dataset.encoder_threads, streaming_encoding=cfg.dataset.streaming_encoding, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,