mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
feat(record): plumb DepthEncoderConfig through lerobot-record
Surface DepthEncoderConfig and depth_encoder_defaults from
lerobot.datasets, and wire dataset.depth_encoder_config through
LeRobotDataset.create() and LeRobotDataset.resume() so depth-capable
recordings (e.g. RealSense use_depth=True) can be tuned from the CLI:
--dataset.depth_encoder_config.depth_min=0.1
--dataset.depth_encoder_config.depth_max=4.0
--dataset.depth_encoder_config.vcodec=ffv1
The default factory keeps depth-stream defaults (12-bit HEVC, log
quantization), so non-depth recordings are unaffected.
Made-with: Cursor
This commit is contained in:
@@ -49,9 +49,11 @@ from .sampler import EpisodeAwareSampler
|
|||||||
from .streaming_dataset import StreamingLeRobotDataset
|
from .streaming_dataset import StreamingLeRobotDataset
|
||||||
from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
|
from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
|
||||||
from .video_utils import (
|
from .video_utils import (
|
||||||
|
DepthEncoderConfig,
|
||||||
VideoEncoderConfig,
|
VideoEncoderConfig,
|
||||||
VideoEncodingManager,
|
VideoEncodingManager,
|
||||||
camera_encoder_defaults,
|
camera_encoder_defaults,
|
||||||
|
depth_encoder_defaults,
|
||||||
)
|
)
|
||||||
|
|
||||||
# NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.)
|
# NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.)
|
||||||
@@ -67,9 +69,11 @@ __all__ = [
|
|||||||
"LeRobotDatasetMetadata",
|
"LeRobotDatasetMetadata",
|
||||||
"MultiLeRobotDataset",
|
"MultiLeRobotDataset",
|
||||||
"StreamingLeRobotDataset",
|
"StreamingLeRobotDataset",
|
||||||
|
"DepthEncoderConfig",
|
||||||
"VideoEncoderConfig",
|
"VideoEncoderConfig",
|
||||||
"VideoEncodingManager",
|
"VideoEncodingManager",
|
||||||
"camera_encoder_defaults",
|
"camera_encoder_defaults",
|
||||||
|
"depth_encoder_defaults",
|
||||||
"add_features",
|
"add_features",
|
||||||
"aggregate_datasets",
|
"aggregate_datasets",
|
||||||
"aggregate_pipeline_dataset_features",
|
"aggregate_pipeline_dataset_features",
|
||||||
|
|||||||
@@ -104,10 +104,12 @@ from lerobot.common.control_utils import (
|
|||||||
from lerobot.configs import parser
|
from lerobot.configs import parser
|
||||||
from lerobot.configs.dataset import DatasetRecordConfig
|
from lerobot.configs.dataset import DatasetRecordConfig
|
||||||
from lerobot.datasets import (
|
from lerobot.datasets import (
|
||||||
|
DepthEncoderConfig,
|
||||||
LeRobotDataset,
|
LeRobotDataset,
|
||||||
VideoEncodingManager,
|
VideoEncodingManager,
|
||||||
aggregate_pipeline_dataset_features,
|
aggregate_pipeline_dataset_features,
|
||||||
create_initial_features,
|
create_initial_features,
|
||||||
|
depth_encoder_defaults,
|
||||||
safe_stop_image_writer,
|
safe_stop_image_writer,
|
||||||
)
|
)
|
||||||
from lerobot.processor import (
|
from lerobot.processor import (
|
||||||
@@ -399,6 +401,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_config=cfg.dataset.camera_encoder_config,
|
camera_encoder_config=cfg.dataset.camera_encoder_config,
|
||||||
|
depth_encoder_config=cfg.dataset.depth_encoder_config,
|
||||||
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 +431,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_config=cfg.dataset.camera_encoder_config,
|
camera_encoder_config=cfg.dataset.camera_encoder_config,
|
||||||
|
depth_encoder_config=cfg.dataset.depth_encoder_config,
|
||||||
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