feat(record): plumb DepthEncoderConfig through lerobot-record

This commit is contained in:
CarolinePascal
2026-05-20 16:14:14 +02:00
parent d6f97ae17f
commit 14df709201
3 changed files with 11 additions and 1 deletions
+4
View File
@@ -35,7 +35,9 @@ from .video import (
VALID_VIDEO_CODECS, VALID_VIDEO_CODECS,
VIDEO_ENCODER_INFO_KEYS, VIDEO_ENCODER_INFO_KEYS,
VideoEncoderConfig, VideoEncoderConfig,
DepthEncoderConfig,
camera_encoder_defaults, camera_encoder_defaults,
depth_encoder_defaults,
) )
__all__ = [ __all__ = [
@@ -53,8 +55,10 @@ __all__ = [
"PreTrainedConfig", "PreTrainedConfig",
"WandBConfig", "WandBConfig",
"VideoEncoderConfig", "VideoEncoderConfig",
"DepthEncoderConfig",
# Defaults # Defaults
"camera_encoder_defaults", "camera_encoder_defaults",
"depth_encoder_defaults",
# Constants # Constants
"VALID_VIDEO_CODECS", "VALID_VIDEO_CODECS",
"VIDEO_ENCODER_INFO_KEYS", "VIDEO_ENCODER_INFO_KEYS",
+5 -1
View File
@@ -761,6 +761,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
video_backend: str | None = None, video_backend: str | None = None,
batch_encoding_size: int = 1, batch_encoding_size: int = 1,
camera_encoder: VideoEncoderConfig | None = None, camera_encoder: VideoEncoderConfig | None = None,
depth_encoder: DepthEncoderConfig | None = None,
encoder_threads: int | None = None, encoder_threads: int | None = None,
image_writer_processes: int = 0, image_writer_processes: int = 0,
image_writer_threads: int = 0, image_writer_threads: int = 0,
@@ -790,6 +791,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
batch-encoding videos. batch-encoding videos.
camera_encoder: Video encoder settings for cameras (codec, quality, etc.). camera_encoder: Video encoder settings for cameras (codec, quality, etc.).
When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults` is used. 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`` encoder_threads: Number of encoder threads (global). ``None``
lets the codec decide. lets the codec decide.
image_writer_processes: Subprocesses for async image writing. image_writer_processes: Subprocesses for async image writing.
@@ -836,12 +839,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
streaming_enc = None streaming_enc = None
if streaming_encoding and len(obj.meta.video_keys) > 0: if streaming_encoding and len(obj.meta.video_keys) > 0:
streaming_enc = cls._build_streaming_encoder( 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( obj.writer = DatasetWriter(
meta=obj.meta, meta=obj.meta,
root=obj.root, root=obj.root,
camera_encoder=camera_encoder, camera_encoder=camera_encoder,
depth_encoder=depth_encoder,
encoder_threads=encoder_threads, encoder_threads=encoder_threads,
batch_encoding_size=batch_encoding_size, batch_encoding_size=batch_encoding_size,
streaming_encoder=streaming_enc, streaming_encoder=streaming_enc,
+2
View File
@@ -399,6 +399,7 @@ def record(
root=cfg.dataset.root, root=cfg.dataset.root,
batch_encoding_size=cfg.dataset.video_encoding_batch_size, batch_encoding_size=cfg.dataset.video_encoding_batch_size,
camera_encoder=cfg.dataset.camera_encoder, camera_encoder=cfg.dataset.camera_encoder,
depth_encoder=cfg.dataset.depth_encoder,
encoder_threads=cfg.dataset.encoder_threads, encoder_threads=cfg.dataset.encoder_threads,
streaming_encoding=cfg.dataset.streaming_encoding, streaming_encoding=cfg.dataset.streaming_encoding,
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, 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), image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera * len(robot.cameras),
batch_encoding_size=cfg.dataset.video_encoding_batch_size, batch_encoding_size=cfg.dataset.video_encoding_batch_size,
camera_encoder=cfg.dataset.camera_encoder, camera_encoder=cfg.dataset.camera_encoder,
depth_encoder=cfg.dataset.depth_encoder,
encoder_threads=cfg.dataset.encoder_threads, encoder_threads=cfg.dataset.encoder_threads,
streaming_encoding=cfg.dataset.streaming_encoding, streaming_encoding=cfg.dataset.streaming_encoding,
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize, encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,