mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-22 20:19:43 +00:00
feat(record): plumb DepthEncoderConfig through lerobot-record
This commit is contained in:
@@ -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",
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user