From e7191fc3ad40cfa1051e58c81b9a37c7b6efa618 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Wed, 20 May 2026 16:14:14 +0200 Subject: [PATCH] feat(record): plumb DepthEncoderConfig through lerobot-record --- src/lerobot/configs/__init__.py | 4 ++++ src/lerobot/datasets/lerobot_dataset.py | 6 +++++- src/lerobot/scripts/lerobot_record.py | 2 ++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/lerobot/configs/__init__.py b/src/lerobot/configs/__init__.py index be4491811..a21932610 100644 --- a/src/lerobot/configs/__init__.py +++ b/src/lerobot/configs/__init__.py @@ -36,7 +36,9 @@ from .video import ( VALID_VIDEO_CODECS, VIDEO_ENCODER_INFO_KEYS, VideoEncoderConfig, + DepthEncoderConfig, camera_encoder_defaults, + depth_encoder_defaults, ) __all__ = [ @@ -57,8 +59,10 @@ __all__ = [ "WandBConfig", "load_recipe", "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 52728858e..cb21563a1 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -762,6 +762,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, @@ -791,6 +792,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. @@ -837,12 +840,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 0deb54b90..6fbb4a223 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -404,6 +404,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, @@ -433,6 +434,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,