mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
chore(rename): renaming camera_encoder_config to camera_encoder
This commit is contained in:
+1
-1
@@ -90,6 +90,6 @@ lerobot-record \
|
|||||||
--dataset.single_task="Your task description" \
|
--dataset.single_task="Your task description" \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--policy.path=${HF_USER}/act_policy
|
--policy.path=${HF_USER}/act_policy
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -194,7 +194,7 @@ lerobot-record \
|
|||||||
--dataset.single_task="Navigate around obstacles" \
|
--dataset.single_task="Navigate around obstacles" \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--display_data=true
|
--display_data=true
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -123,7 +123,7 @@ lerobot-record \
|
|||||||
--dataset.single_task="Grab and handover the red cube to the other arm" \
|
--dataset.single_task="Grab and handover the red cube to the other arm" \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--policy.path=<user>/groot-bimanual \ # your trained model
|
--policy.path=<user>/groot-bimanual \ # your trained model
|
||||||
--dataset.episode_time_s=30 \
|
--dataset.episode_time_s=30 \
|
||||||
--dataset.reset_time_s=10
|
--dataset.reset_time_s=10
|
||||||
|
|||||||
@@ -232,7 +232,7 @@ lerobot-record \
|
|||||||
--dataset.private=true \
|
--dataset.private=true \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--display_data=true
|
--display_data=true
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -278,6 +278,6 @@ lerobot-record \
|
|||||||
--dataset.num_episodes=10 \
|
--dataset.num_episodes=10 \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model
|
--policy.path=outputs/train/hopejr_hand/checkpoints/last/pretrained_model
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -193,7 +193,7 @@ lerobot-record \
|
|||||||
--dataset.num_episodes=5 \
|
--dataset.num_episodes=5 \
|
||||||
--dataset.single_task="Grab the black cube" \
|
--dataset.single_task="Grab the black cube" \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--dataset.encoder_threads=2
|
--dataset.encoder_threads=2
|
||||||
```
|
```
|
||||||
</hfoption>
|
</hfoption>
|
||||||
|
|||||||
@@ -43,7 +43,7 @@ lerobot-record \
|
|||||||
--dataset.num_episodes=5 \
|
--dataset.num_episodes=5 \
|
||||||
--dataset.single_task="Grab the black cube" \
|
--dataset.single_task="Grab the black cube" \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--dataset.encoder_threads=2
|
--dataset.encoder_threads=2
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -161,7 +161,7 @@ lerobot-record \
|
|||||||
--dataset.private=true \
|
--dataset.private=true \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--display_data=true
|
--display_data=true
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -203,7 +203,7 @@ lerobot-record \
|
|||||||
--dataset.private=true \
|
--dataset.private=true \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
--display_data=true
|
--display_data=true
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -108,7 +108,7 @@ lerobot-record \
|
|||||||
--dataset.num_episodes=10 \
|
--dataset.num_episodes=10 \
|
||||||
--dataset.streaming_encoding=true \
|
--dataset.streaming_encoding=true \
|
||||||
--dataset.encoder_threads=2 \
|
--dataset.encoder_threads=2 \
|
||||||
# --dataset.camera_encoder_config.vcodec=auto \
|
# --dataset.camera_encoder.vcodec=auto \
|
||||||
# <- Teleop optional if you want to teleoperate in between episodes \
|
# <- Teleop optional if you want to teleoperate in between episodes \
|
||||||
# --teleop.type=so100_leader \
|
# --teleop.type=so100_leader \
|
||||||
# --teleop.port=/dev/ttyACM0 \
|
# --teleop.port=/dev/ttyACM0 \
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ This makes `save_episode()` near-instant (the video is already encoded by the ti
|
|||||||
| Parameter | CLI Flag | Type | Default | Description |
|
| Parameter | CLI Flag | Type | Default | Description |
|
||||||
| ----------------------- | ---------------------------------------- | ------------- | ------------- | ----------------------------------------------------------------- |
|
| ----------------------- | ---------------------------------------- | ------------- | ------------- | ----------------------------------------------------------------- |
|
||||||
| `streaming_encoding` | `--dataset.streaming_encoding` | `bool` | `True` | Enable real-time encoding during capture |
|
| `streaming_encoding` | `--dataset.streaming_encoding` | `bool` | `True` | Enable real-time encoding during capture |
|
||||||
| `vcodec` | `--dataset.camera_encoder_config.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder |
|
| `vcodec` | `--dataset.camera_encoder.vcodec` | `str` | `"libsvtav1"` | Video codec. `"auto"` detects best HW encoder |
|
||||||
| `encoder_threads` | `--dataset.encoder_threads` | `int \| None` | `None` (auto) | Threads per encoder instance. `None` will leave the vcoded decide |
|
| `encoder_threads` | `--dataset.encoder_threads` | `int \| None` | `None` (auto) | Threads per encoder instance. `None` will leave the vcoded decide |
|
||||||
| `encoder_queue_maxsize` | `--dataset.encoder_queue_maxsize` | `int` | `30` | Max buffered frames per camera (~1s at 30fps). Consumes RAM |
|
| `encoder_queue_maxsize` | `--dataset.encoder_queue_maxsize` | `int` | `30` | Max buffered frames per camera (~1s at 30fps). Consumes RAM |
|
||||||
|
|
||||||
@@ -84,13 +84,13 @@ Use HW encoding when:
|
|||||||
|
|
||||||
| Encoder | Platform | Hardware | CLI Value |
|
| Encoder | Platform | Hardware | CLI Value |
|
||||||
| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | ---------------------------------------------------------- |
|
| ------------------- | ------------- | ------------------------------------------------------------------------------------------------ | ---------------------------------------------------------- |
|
||||||
| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder_config.vcodec=h264_videotoolbox` |
|
| `h264_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=h264_videotoolbox` |
|
||||||
| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder_config.vcodec=hevc_videotoolbox` |
|
| `hevc_videotoolbox` | macOS | Apple Silicon / Intel | `--dataset.camera_encoder.vcodec=hevc_videotoolbox` |
|
||||||
| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder_config.vcodec=h264_nvenc` |
|
| `h264_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=h264_nvenc` |
|
||||||
| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder_config.vcodec=hevc_nvenc` |
|
| `hevc_nvenc` | Linux/Windows | NVIDIA GPU | `--dataset.camera_encoder.vcodec=hevc_nvenc` |
|
||||||
| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.camera_encoder_config.vcodec=h264_vaapi` |
|
| `h264_vaapi` | Linux | Intel/AMD GPU | `--dataset.camera_encoder.vcodec=h264_vaapi` |
|
||||||
| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.camera_encoder_config.vcodec=h264_qsv` |
|
| `h264_qsv` | Linux/Windows | Intel Quick Sync | `--dataset.camera_encoder.vcodec=h264_qsv` |
|
||||||
| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.camera_encoder_config.vcodec=auto` |
|
| `auto` | Any | Probes the system for available HW encoders. Falls back to `libsvtav1` if no HW encoder is found | `--dataset.camera_encoder.vcodec=auto` |
|
||||||
|
|
||||||
> [!NOTE]
|
> [!NOTE]
|
||||||
> In order to use the HW accelerated encoders you might need to upgrade your GPU drivers.
|
> In order to use the HW accelerated encoders you might need to upgrade your GPU drivers.
|
||||||
@@ -102,12 +102,12 @@ Use HW encoding when:
|
|||||||
|
|
||||||
| Symptom | Likely Cause | Fix |
|
| Symptom | Likely Cause | Fix |
|
||||||
| ------------------------------------------------------------------ | -------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
| ------------------------------------------------------------------ | -------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||||
| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.camera_encoder_config.vcodec=auto`) |
|
| System freezes or choppy robot movement or Rerun visualization lag | CPU starved (100% load usage) | Close other apps, reduce encoding throughput, lower `encoder_threads`, use `h264`, use `display_data=False`. If the CPU continues to be at 100% then it might be insufficient for your setup, consider `--dataset.streaming_encoding=false` or HW encoding (`--dataset.camera_encoder.vcodec=auto`) |
|
||||||
| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.camera_encoder_config.vcodec=auto`). |
|
| "Encoder queue full" warnings or dropped frames in dataset | Encoder can't keep up (Queue overflow) | If CPU is not at 100%: Increase `encoder_threads`, increase `encoder_queue_maxsize` or use HW encoding (`--dataset.camera_encoder.vcodec=auto`). |
|
||||||
| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding |
|
| High RAM usage | Queue filling faster than encoding | `encoder_threads` too low or CPU insufficient. Reduce `encoder_queue_maxsize` or use HW encoding |
|
||||||
| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows |
|
| Large video files | Using HW encoder or H.264 | Expected trade-off. Switch to `libsvtav1` if CPU allows |
|
||||||
| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` |
|
| `save_episode()` still slow | `streaming_encoding` is `False` | Set `--dataset.streaming_encoding=true` |
|
||||||
| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.camera_encoder_config.vcodec=auto` |
|
| Encoder thread crash | Codec not available or invalid settings | Check `vcodec` is installed, try `--dataset.camera_encoder.vcodec=auto` |
|
||||||
| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. |
|
| Recorded dataset is missing frames | CPU/GPU starvation or occasional load spikes | If ~5% of frames are missing, your system is likely overloaded — follow the recommendations above. If fewer frames are missing (~2%), they are probably due to occasional transient load spikes (often at startup) and can be considered expected. |
|
||||||
|
|
||||||
## 6. Recommended Configurations
|
## 6. Recommended Configurations
|
||||||
@@ -146,7 +146,7 @@ On very constrained systems, streaming encoding may compete too heavily with the
|
|||||||
# 2camsx 640x480x3 @30fps: Requires some tuning.
|
# 2camsx 640x480x3 @30fps: Requires some tuning.
|
||||||
|
|
||||||
# Use H.264, disable streaming, consider batching encoding
|
# Use H.264, disable streaming, consider batching encoding
|
||||||
lerobot-record --dataset.camera_encoder_config.vcodec=h264 --dataset.streaming_encoding=false ...
|
lerobot-record --dataset.camera_encoder.vcodec=h264 --dataset.streaming_encoding=false ...
|
||||||
```
|
```
|
||||||
|
|
||||||
## 7. Closing note
|
## 7. Closing note
|
||||||
|
|||||||
@@ -117,10 +117,10 @@ lerobot-edit-dataset \
|
|||||||
--repo_id lerobot/pusht_image \
|
--repo_id lerobot/pusht_image \
|
||||||
--operation.type convert_image_to_video \
|
--operation.type convert_image_to_video \
|
||||||
--operation.output_dir outputs/pusht_video \
|
--operation.output_dir outputs/pusht_video \
|
||||||
--operation.camera_encoder_config.vcodec libsvtav1 \
|
--operation.camera_encoder.vcodec libsvtav1 \
|
||||||
--operation.camera_encoder_config.pix_fmt yuv420p \
|
--operation.camera_encoder.pix_fmt yuv420p \
|
||||||
--operation.camera_encoder_config.g 2 \
|
--operation.camera_encoder.g 2 \
|
||||||
--operation.camera_encoder_config.crf 30
|
--operation.camera_encoder.crf 30
|
||||||
|
|
||||||
# Convert only specific episodes
|
# Convert only specific episodes
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -147,7 +147,7 @@ lerobot-edit-dataset \
|
|||||||
**Parameters:**
|
**Parameters:**
|
||||||
|
|
||||||
- `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`)
|
- `output_dir`: Custom output directory (optional - by default uses `new_repo_id` or `{repo_id}_video`)
|
||||||
- `camera_encoder_config`: Video encoder settings — all sub-fields accessible via `--operation.camera_encoder_config.<field>. See [Video Encoding Parameters](./video_encoding_parameters) for more details.
|
- `camera_encoder`: Video encoder settings — all sub-fields accessible via `--operation.camera_encoder.<field>. See [Video Encoding Parameters](./video_encoding_parameters) for more details.
|
||||||
- `episode_indices`: List of specific episodes to convert (default: all episodes)
|
- `episode_indices`: List of specific episodes to convert (default: all episodes)
|
||||||
- `num_workers`: Number of parallel workers for processing (default: 4)
|
- `num_workers`: Number of parallel workers for processing (default: 4)
|
||||||
|
|
||||||
|
|||||||
@@ -58,8 +58,8 @@ class DatasetRecordConfig:
|
|||||||
# Set to 1 for immediate encoding (default behavior), or higher for batched encoding
|
# Set to 1 for immediate encoding (default behavior), or higher for batched encoding
|
||||||
video_encoding_batch_size: int = 1
|
video_encoding_batch_size: int = 1
|
||||||
# Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys,
|
# Video encoder settings for camera MP4s (codec, quality, GOP, etc.). Tuned via CLI nested keys,
|
||||||
# e.g. ``--dataset.camera_encoder_config.vcodec=h264`` (see ``VideoEncoderConfig``).
|
# e.g. ``--dataset.camera_encoder.vcodec=h264`` (see ``VideoEncoderConfig``).
|
||||||
camera_encoder_config: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
||||||
# Enable streaming video encoding: encode frames in real-time during capture instead
|
# Enable streaming video encoding: encode frames in real-time during capture instead
|
||||||
# of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
|
# of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
|
||||||
streaming_encoding: bool = False
|
streaming_encoding: bool = False
|
||||||
|
|||||||
@@ -18,8 +18,8 @@ from logging import getLogger
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from lerobot import envs, policies # noqa: F401
|
from lerobot import envs, policies # noqa: F401
|
||||||
from lerobot.configs import parser
|
|
||||||
|
|
||||||
|
from . import parser
|
||||||
from .default import EvalConfig
|
from .default import EvalConfig
|
||||||
from .policies import PreTrainedConfig
|
from .policies import PreTrainedConfig
|
||||||
|
|
||||||
|
|||||||
@@ -27,12 +27,13 @@ from huggingface_hub import hf_hub_download
|
|||||||
from huggingface_hub.constants import CONFIG_NAME
|
from huggingface_hub.constants import CONFIG_NAME
|
||||||
from huggingface_hub.errors import HfHubHTTPError
|
from huggingface_hub.errors import HfHubHTTPError
|
||||||
|
|
||||||
from lerobot.configs.types import PolicyFeature
|
|
||||||
from lerobot.optim.optimizers import OptimizerConfig
|
from lerobot.optim.optimizers import OptimizerConfig
|
||||||
from lerobot.optim.schedulers import LRSchedulerConfig
|
from lerobot.optim.schedulers import LRSchedulerConfig
|
||||||
from lerobot.utils.device_utils import auto_select_torch_device, is_torch_device_available
|
from lerobot.utils.device_utils import auto_select_torch_device, is_torch_device_available
|
||||||
from lerobot.utils.hub import HubMixin
|
from lerobot.utils.hub import HubMixin
|
||||||
|
|
||||||
|
from .types import PolicyFeature
|
||||||
|
|
||||||
T = TypeVar("T", bound="RewardModelConfig")
|
T = TypeVar("T", bound="RewardModelConfig")
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|||||||
@@ -25,11 +25,11 @@ from huggingface_hub import hf_hub_download
|
|||||||
from huggingface_hub.errors import HfHubHTTPError
|
from huggingface_hub.errors import HfHubHTTPError
|
||||||
|
|
||||||
from lerobot import envs
|
from lerobot import envs
|
||||||
from lerobot.configs import parser
|
|
||||||
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
|
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
|
||||||
from lerobot.utils.hub import HubMixin
|
from lerobot.utils.hub import HubMixin
|
||||||
from lerobot.utils.sample_weighting import SampleWeightingConfig
|
from lerobot.utils.sample_weighting import SampleWeightingConfig
|
||||||
|
|
||||||
|
from . import parser
|
||||||
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
||||||
from .policies import PreTrainedConfig
|
from .policies import PreTrainedConfig
|
||||||
from .rewards import RewardModelConfig
|
from .rewards import RewardModelConfig
|
||||||
|
|||||||
@@ -514,7 +514,7 @@ class LeRobotDatasetMetadata:
|
|||||||
def update_video_info(
|
def update_video_info(
|
||||||
self,
|
self,
|
||||||
video_key: str | None = None,
|
video_key: str | None = None,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Populate per-feature video info in ``info.json``.
|
"""Populate per-feature video info in ``info.json``.
|
||||||
|
|
||||||
@@ -524,7 +524,7 @@ class LeRobotDatasetMetadata:
|
|||||||
Args:
|
Args:
|
||||||
video_key: If provided, only update this video key. Otherwise update
|
video_key: If provided, only update this video key. Otherwise update
|
||||||
all video keys in the dataset.
|
all video keys in the dataset.
|
||||||
camera_encoder_config: Encoder configuration used to produce the
|
camera_encoder: Encoder configuration used to produce the
|
||||||
videos. When provided, its fields are recorded as
|
videos. When provided, its fields are recorded as
|
||||||
``video.<field>`` entries alongside the stream-derived
|
``video.<field>`` entries alongside the stream-derived
|
||||||
``video.*`` entries (see :func:`get_video_info`).
|
``video.*`` entries (see :func:`get_video_info`).
|
||||||
@@ -537,7 +537,7 @@ class LeRobotDatasetMetadata:
|
|||||||
if not self.features[key].get("info", None):
|
if not self.features[key].get("info", None):
|
||||||
video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
|
video_path = self.root / self.video_path.format(video_key=key, chunk_index=0, file_index=0)
|
||||||
self.info.features[key]["info"] = get_video_info(
|
self.info.features[key]["info"] = get_video_info(
|
||||||
video_path, camera_encoder_config=camera_encoder_config
|
video_path, camera_encoder=camera_encoder
|
||||||
)
|
)
|
||||||
|
|
||||||
def update_chunk_settings(
|
def update_chunk_settings(
|
||||||
|
|||||||
@@ -96,7 +96,7 @@ def delete_episodes(
|
|||||||
episode_indices: list[int],
|
episode_indices: list[int],
|
||||||
output_dir: str | Path | None = None,
|
output_dir: str | Path | None = None,
|
||||||
repo_id: str | None = None,
|
repo_id: str | None = None,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> LeRobotDataset:
|
) -> LeRobotDataset:
|
||||||
"""Delete episodes from a LeRobotDataset and create a new dataset.
|
"""Delete episodes from a LeRobotDataset and create a new dataset.
|
||||||
|
|
||||||
@@ -105,7 +105,7 @@ def delete_episodes(
|
|||||||
episode_indices: List of episode indices to delete.
|
episode_indices: List of episode indices to delete.
|
||||||
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
camera_encoder_config: Video encoder settings used when re-encoding video segments
|
camera_encoder: Video encoder settings used when re-encoding video segments
|
||||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||||
"""
|
"""
|
||||||
if not episode_indices:
|
if not episode_indices:
|
||||||
@@ -139,7 +139,7 @@ def delete_episodes(
|
|||||||
|
|
||||||
video_metadata = None
|
video_metadata = None
|
||||||
if dataset.meta.video_keys:
|
if dataset.meta.video_keys:
|
||||||
video_metadata = _copy_and_reindex_videos(dataset, new_meta, episode_mapping, camera_encoder_config)
|
video_metadata = _copy_and_reindex_videos(dataset, new_meta, episode_mapping, camera_encoder)
|
||||||
|
|
||||||
data_metadata = _copy_and_reindex_data(dataset, new_meta, episode_mapping)
|
data_metadata = _copy_and_reindex_data(dataset, new_meta, episode_mapping)
|
||||||
|
|
||||||
@@ -161,7 +161,7 @@ def split_dataset(
|
|||||||
dataset: LeRobotDataset,
|
dataset: LeRobotDataset,
|
||||||
splits: dict[str, float | list[int]],
|
splits: dict[str, float | list[int]],
|
||||||
output_dir: str | Path | None = None,
|
output_dir: str | Path | None = None,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> dict[str, LeRobotDataset]:
|
) -> dict[str, LeRobotDataset]:
|
||||||
"""Split a LeRobotDataset into multiple smaller datasets.
|
"""Split a LeRobotDataset into multiple smaller datasets.
|
||||||
|
|
||||||
@@ -170,7 +170,7 @@ def split_dataset(
|
|||||||
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
splits: Either a dict mapping split names to episode indices, or a dict mapping
|
||||||
split names to fractions (must sum to <= 1.0).
|
split names to fractions (must sum to <= 1.0).
|
||||||
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
output_dir: Root directory where the split datasets will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id.
|
||||||
camera_encoder_config: Video encoder settings used when re-encoding video segments
|
camera_encoder: Video encoder settings used when re-encoding video segments
|
||||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||||
|
|
||||||
Examples:
|
Examples:
|
||||||
@@ -233,7 +233,7 @@ def split_dataset(
|
|||||||
video_metadata = None
|
video_metadata = None
|
||||||
if dataset.meta.video_keys:
|
if dataset.meta.video_keys:
|
||||||
video_metadata = _copy_and_reindex_videos(
|
video_metadata = _copy_and_reindex_videos(
|
||||||
dataset, new_meta, episode_mapping, camera_encoder_config
|
dataset, new_meta, episode_mapping, camera_encoder
|
||||||
)
|
)
|
||||||
|
|
||||||
data_metadata = _copy_and_reindex_data(dataset, new_meta, episode_mapping)
|
data_metadata = _copy_and_reindex_data(dataset, new_meta, episode_mapping)
|
||||||
@@ -590,7 +590,7 @@ def _keep_episodes_from_video_with_av(
|
|||||||
output_path: Path,
|
output_path: Path,
|
||||||
episodes_to_keep: list[tuple[int, int]],
|
episodes_to_keep: list[tuple[int, int]],
|
||||||
fps: float,
|
fps: float,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Keep only specified episodes from a video file using PyAV.
|
"""Keep only specified episodes from a video file using PyAV.
|
||||||
|
|
||||||
@@ -604,11 +604,11 @@ def _keep_episodes_from_video_with_av(
|
|||||||
Ranges are half-open intervals: [start_frame, end_frame), where start_frame
|
Ranges are half-open intervals: [start_frame, end_frame), where start_frame
|
||||||
is inclusive and end_frame is exclusive.
|
is inclusive and end_frame is exclusive.
|
||||||
fps: Frame rate of the video.
|
fps: Frame rate of the video.
|
||||||
camera_encoder_config: Video encoder settings
|
camera_encoder: Video encoder settings
|
||||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||||
"""
|
"""
|
||||||
if camera_encoder_config is None:
|
if camera_encoder is None:
|
||||||
camera_encoder_config = camera_encoder_defaults()
|
camera_encoder = camera_encoder_defaults()
|
||||||
from fractions import Fraction
|
from fractions import Fraction
|
||||||
|
|
||||||
import av
|
import av
|
||||||
@@ -632,13 +632,13 @@ def _keep_episodes_from_video_with_av(
|
|||||||
|
|
||||||
# Convert fps to Fraction for PyAV compatibility.
|
# Convert fps to Fraction for PyAV compatibility.
|
||||||
fps_fraction = Fraction(fps).limit_denominator(1000)
|
fps_fraction = Fraction(fps).limit_denominator(1000)
|
||||||
codec_options = camera_encoder_config.get_codec_options(as_strings=True)
|
codec_options = camera_encoder.get_codec_options(as_strings=True)
|
||||||
v_out = out.add_stream(camera_encoder_config.vcodec, rate=fps_fraction, options=codec_options)
|
v_out = out.add_stream(camera_encoder.vcodec, rate=fps_fraction, options=codec_options)
|
||||||
|
|
||||||
# PyAV type stubs don't distinguish video streams from audio/subtitle streams.
|
# PyAV type stubs don't distinguish video streams from audio/subtitle streams.
|
||||||
v_out.width = v_in.codec_context.width
|
v_out.width = v_in.codec_context.width
|
||||||
v_out.height = v_in.codec_context.height
|
v_out.height = v_in.codec_context.height
|
||||||
v_out.pix_fmt = camera_encoder_config.pix_fmt
|
v_out.pix_fmt = camera_encoder.pix_fmt
|
||||||
|
|
||||||
# Set time_base to match the frame rate for proper timestamp handling.
|
# Set time_base to match the frame rate for proper timestamp handling.
|
||||||
v_out.time_base = Fraction(1, int(fps))
|
v_out.time_base = Fraction(1, int(fps))
|
||||||
@@ -701,7 +701,7 @@ def _copy_and_reindex_videos(
|
|||||||
src_dataset: LeRobotDataset,
|
src_dataset: LeRobotDataset,
|
||||||
dst_meta: LeRobotDatasetMetadata,
|
dst_meta: LeRobotDatasetMetadata,
|
||||||
episode_mapping: dict[int, int],
|
episode_mapping: dict[int, int],
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> dict[int, dict]:
|
) -> dict[int, dict]:
|
||||||
"""Copy and filter video files, only re-encoding files with deleted episodes.
|
"""Copy and filter video files, only re-encoding files with deleted episodes.
|
||||||
|
|
||||||
@@ -713,14 +713,14 @@ def _copy_and_reindex_videos(
|
|||||||
src_dataset: Source dataset to copy from
|
src_dataset: Source dataset to copy from
|
||||||
dst_meta: Destination metadata object
|
dst_meta: Destination metadata object
|
||||||
episode_mapping: Mapping from old episode indices to new indices
|
episode_mapping: Mapping from old episode indices to new indices
|
||||||
camera_encoder_config: Video encoder settings used when re-encoding segments
|
camera_encoder: Video encoder settings used when re-encoding segments
|
||||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
dict mapping episode index to its video metadata (chunk_index, file_index, timestamps)
|
dict mapping episode index to its video metadata (chunk_index, file_index, timestamps)
|
||||||
"""
|
"""
|
||||||
if camera_encoder_config is None:
|
if camera_encoder is None:
|
||||||
camera_encoder_config = camera_encoder_defaults()
|
camera_encoder = camera_encoder_defaults()
|
||||||
if src_dataset.meta.episodes is None:
|
if src_dataset.meta.episodes is None:
|
||||||
src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
|
src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
|
||||||
|
|
||||||
@@ -809,7 +809,7 @@ def _copy_and_reindex_videos(
|
|||||||
dst_video_path,
|
dst_video_path,
|
||||||
episodes_to_keep_ranges,
|
episodes_to_keep_ranges,
|
||||||
src_dataset.meta.fps,
|
src_dataset.meta.fps,
|
||||||
camera_encoder_config,
|
camera_encoder,
|
||||||
)
|
)
|
||||||
|
|
||||||
cumulative_ts = 0.0
|
cumulative_ts = 0.0
|
||||||
@@ -1280,7 +1280,7 @@ def _estimate_frame_size_via_calibration(
|
|||||||
episode_indices: list[int],
|
episode_indices: list[int],
|
||||||
temp_dir: Path,
|
temp_dir: Path,
|
||||||
fps: int,
|
fps: int,
|
||||||
camera_encoder_config: VideoEncoderConfig,
|
camera_encoder: VideoEncoderConfig,
|
||||||
num_calibration_frames: int = 30,
|
num_calibration_frames: int = 30,
|
||||||
) -> float:
|
) -> float:
|
||||||
"""Estimate MB per frame by encoding a small calibration sample.
|
"""Estimate MB per frame by encoding a small calibration sample.
|
||||||
@@ -1294,7 +1294,7 @@ def _estimate_frame_size_via_calibration(
|
|||||||
episode_indices: List of episode indices being processed.
|
episode_indices: List of episode indices being processed.
|
||||||
temp_dir: Temporary directory for calibration files.
|
temp_dir: Temporary directory for calibration files.
|
||||||
fps: Frames per second for video encoding.
|
fps: Frames per second for video encoding.
|
||||||
camera_encoder_config: Video encoder settings used for calibration encoding.
|
camera_encoder: Video encoder settings used for calibration encoding.
|
||||||
num_calibration_frames: Number of frames to use for calibration (default: 30).
|
num_calibration_frames: Number of frames to use for calibration (default: 30).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
@@ -1330,7 +1330,7 @@ def _estimate_frame_size_via_calibration(
|
|||||||
imgs_dir=calibration_dir,
|
imgs_dir=calibration_dir,
|
||||||
video_path=calibration_video_path,
|
video_path=calibration_video_path,
|
||||||
fps=fps,
|
fps=fps,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_encoder,
|
||||||
overwrite=True,
|
overwrite=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1648,7 +1648,7 @@ def convert_image_to_video_dataset(
|
|||||||
dataset: LeRobotDataset,
|
dataset: LeRobotDataset,
|
||||||
output_dir: Path | None = None,
|
output_dir: Path | None = None,
|
||||||
repo_id: str | None = None,
|
repo_id: str | None = None,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
episode_indices: list[int] | None = None,
|
episode_indices: list[int] | None = None,
|
||||||
num_workers: int = 4,
|
num_workers: int = 4,
|
||||||
max_episodes_per_batch: int | None = None,
|
max_episodes_per_batch: int | None = None,
|
||||||
@@ -1663,7 +1663,7 @@ def convert_image_to_video_dataset(
|
|||||||
dataset: The source LeRobot dataset with images
|
dataset: The source LeRobot dataset with images
|
||||||
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
output_dir: Root directory where the edited dataset will be stored. If not specified, defaults to $HF_LEROBOT_HOME/repo_id. Equivalent to new_root in EditDatasetConfig.
|
||||||
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
repo_id: Edited dataset identifier. Equivalent to new_repo_id in EditDatasetConfig.
|
||||||
camera_encoder_config: Video encoder settings
|
camera_encoder: Video encoder settings
|
||||||
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
(``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`).
|
||||||
episode_indices: List of episode indices to convert (None = all episodes)
|
episode_indices: List of episode indices to convert (None = all episodes)
|
||||||
num_workers: Number of threads for parallel processing (default: 4)
|
num_workers: Number of threads for parallel processing (default: 4)
|
||||||
@@ -1673,8 +1673,8 @@ def convert_image_to_video_dataset(
|
|||||||
Returns:
|
Returns:
|
||||||
New LeRobotDataset with images encoded as videos
|
New LeRobotDataset with images encoded as videos
|
||||||
"""
|
"""
|
||||||
if camera_encoder_config is None:
|
if camera_encoder is None:
|
||||||
camera_encoder_config = camera_encoder_defaults()
|
camera_encoder = camera_encoder_defaults()
|
||||||
|
|
||||||
# Check that it's an image dataset
|
# Check that it's an image dataset
|
||||||
if len(dataset.meta.video_keys) > 0:
|
if len(dataset.meta.video_keys) > 0:
|
||||||
@@ -1700,8 +1700,8 @@ def convert_image_to_video_dataset(
|
|||||||
f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}"
|
f"Converting {len(episode_indices)} episodes with {len(img_keys)} cameras from {dataset.repo_id}"
|
||||||
)
|
)
|
||||||
logging.info(
|
logging.info(
|
||||||
f"Video codec: {camera_encoder_config.vcodec}, pixel format: {camera_encoder_config.pix_fmt}, "
|
f"Video codec: {camera_encoder.vcodec}, pixel format: {camera_encoder.pix_fmt}, "
|
||||||
f"GOP: {camera_encoder_config.g}, CRF: {camera_encoder_config.crf}"
|
f"GOP: {camera_encoder.g}, CRF: {camera_encoder.crf}"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Create new features dict, converting image features to video features
|
# Create new features dict, converting image features to video features
|
||||||
@@ -1772,7 +1772,7 @@ def convert_image_to_video_dataset(
|
|||||||
episode_indices=episode_indices,
|
episode_indices=episode_indices,
|
||||||
temp_dir=temp_dir,
|
temp_dir=temp_dir,
|
||||||
fps=fps,
|
fps=fps,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_encoder,
|
||||||
)
|
)
|
||||||
|
|
||||||
logging.info(f"Processing camera: {img_key}")
|
logging.info(f"Processing camera: {img_key}")
|
||||||
@@ -1814,7 +1814,7 @@ def convert_image_to_video_dataset(
|
|||||||
imgs_dir=imgs_dir,
|
imgs_dir=imgs_dir,
|
||||||
video_path=video_path,
|
video_path=video_path,
|
||||||
fps=fps,
|
fps=fps,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_encoder,
|
||||||
overwrite=True,
|
overwrite=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1861,7 +1861,7 @@ def convert_image_to_video_dataset(
|
|||||||
video_key=img_key, chunk_index=0, file_index=0
|
video_key=img_key, chunk_index=0, file_index=0
|
||||||
)
|
)
|
||||||
new_meta.info.features[img_key]["info"] = get_video_info(
|
new_meta.info.features[img_key]["info"] = get_video_info(
|
||||||
video_path, camera_encoder_config=camera_encoder_config
|
video_path, camera_encoder=camera_encoder
|
||||||
)
|
)
|
||||||
|
|
||||||
write_info(new_meta.info, new_meta.root)
|
write_info(new_meta.info, new_meta.root)
|
||||||
|
|||||||
@@ -67,7 +67,7 @@ def _encode_video_worker(
|
|||||||
episode_index: int,
|
episode_index: int,
|
||||||
root: Path,
|
root: Path,
|
||||||
fps: int,
|
fps: int,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
encoder_threads: int | None = None,
|
encoder_threads: int | None = None,
|
||||||
) -> Path:
|
) -> Path:
|
||||||
temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4"
|
temp_path = Path(tempfile.mkdtemp(dir=root)) / f"{video_key}_{episode_index:03d}.mp4"
|
||||||
@@ -77,7 +77,7 @@ def _encode_video_worker(
|
|||||||
img_dir,
|
img_dir,
|
||||||
temp_path,
|
temp_path,
|
||||||
fps,
|
fps,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_encoder,
|
||||||
encoder_threads=encoder_threads,
|
encoder_threads=encoder_threads,
|
||||||
overwrite=True,
|
overwrite=True,
|
||||||
)
|
)
|
||||||
@@ -96,7 +96,7 @@ class DatasetWriter:
|
|||||||
self,
|
self,
|
||||||
meta: LeRobotDatasetMetadata,
|
meta: LeRobotDatasetMetadata,
|
||||||
root: Path,
|
root: Path,
|
||||||
camera_encoder_config: VideoEncoderConfig | None,
|
camera_encoder: VideoEncoderConfig | None,
|
||||||
encoder_threads: int | None,
|
encoder_threads: int | None,
|
||||||
batch_encoding_size: int,
|
batch_encoding_size: int,
|
||||||
streaming_encoder: StreamingVideoEncoder | None = None,
|
streaming_encoder: StreamingVideoEncoder | None = None,
|
||||||
@@ -108,7 +108,7 @@ class DatasetWriter:
|
|||||||
meta: Dataset metadata instance (used for feature schema, chunk
|
meta: Dataset metadata instance (used for feature schema, chunk
|
||||||
settings, and episode persistence).
|
settings, and episode persistence).
|
||||||
root: Local dataset root directory.
|
root: Local dataset root directory.
|
||||||
camera_encoder_config: Video encoder settings applied to all cameras.
|
camera_encoder: Video encoder settings applied to all cameras.
|
||||||
``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`.
|
``None`` uses :func:`~lerobot.configs.camera_encoder_defaults`.
|
||||||
encoder_threads: Number of encoder threads (global). ``None``
|
encoder_threads: Number of encoder threads (global). ``None``
|
||||||
lets the codec decide.
|
lets the codec decide.
|
||||||
@@ -120,7 +120,7 @@ class DatasetWriter:
|
|||||||
"""
|
"""
|
||||||
self._meta = meta
|
self._meta = meta
|
||||||
self._root = root
|
self._root = root
|
||||||
self._camera_encoder_config = camera_encoder_config or camera_encoder_defaults()
|
self._camera_encoder = camera_encoder or camera_encoder_defaults()
|
||||||
self._encoder_threads = encoder_threads
|
self._encoder_threads = encoder_threads
|
||||||
self._batch_encoding_size = batch_encoding_size
|
self._batch_encoding_size = batch_encoding_size
|
||||||
self._streaming_encoder = streaming_encoder
|
self._streaming_encoder = streaming_encoder
|
||||||
@@ -293,7 +293,7 @@ class DatasetWriter:
|
|||||||
episode_index,
|
episode_index,
|
||||||
self._root,
|
self._root,
|
||||||
self._meta.fps,
|
self._meta.fps,
|
||||||
self._camera_encoder_config,
|
self._camera_encoder,
|
||||||
self._encoder_threads,
|
self._encoder_threads,
|
||||||
): video_key
|
): video_key
|
||||||
for video_key in self._meta.video_keys
|
for video_key in self._meta.video_keys
|
||||||
@@ -504,7 +504,7 @@ class DatasetWriter:
|
|||||||
|
|
||||||
# Update video info (only needed when first episode is encoded)
|
# Update video info (only needed when first episode is encoded)
|
||||||
if episode_index == 0:
|
if episode_index == 0:
|
||||||
self._meta.update_video_info(video_key, camera_encoder_config=self._camera_encoder_config)
|
self._meta.update_video_info(video_key, camera_encoder=self._camera_encoder)
|
||||||
write_info(self._meta.info, self._meta.root)
|
write_info(self._meta.info, self._meta.root)
|
||||||
|
|
||||||
metadata = {
|
metadata = {
|
||||||
@@ -577,7 +577,7 @@ class DatasetWriter:
|
|||||||
episode_index,
|
episode_index,
|
||||||
self._root,
|
self._root,
|
||||||
self._meta.fps,
|
self._meta.fps,
|
||||||
self._camera_encoder_config,
|
self._camera_encoder,
|
||||||
self._encoder_threads,
|
self._encoder_threads,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -58,7 +58,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
video_backend: str | None = None,
|
video_backend: str | None = None,
|
||||||
return_uint8: bool = False,
|
return_uint8: bool = False,
|
||||||
batch_encoding_size: int = 1,
|
batch_encoding_size: int = 1,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
encoder_threads: int | None = None,
|
encoder_threads: int | None = None,
|
||||||
streaming_encoding: bool = False,
|
streaming_encoding: bool = False,
|
||||||
encoder_queue_maxsize: int = 30,
|
encoder_queue_maxsize: int = 30,
|
||||||
@@ -177,7 +177,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
|
You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
|
||||||
batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
|
batch_encoding_size (int, optional): Number of episodes to accumulate before batch encoding videos.
|
||||||
Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
|
Set to 1 for immediate encoding (default), or higher for batched encoding. Defaults to 1.
|
||||||
camera_encoder_config (VideoEncoderConfig | None, optional): Video encoder settings for cameras
|
camera_encoder (VideoEncoderConfig | None, optional): Video encoder settings for cameras
|
||||||
(codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults`
|
(codec, quality, etc.). When ``None``, :func:`~lerobot.configs.video.camera_encoder_defaults`
|
||||||
is used by the writer.
|
is used by the writer.
|
||||||
encoder_threads (int | None, optional): Number of encoder threads (global). ``None`` lets the
|
encoder_threads (int | None, optional): Number of encoder threads (global). ``None`` lets the
|
||||||
@@ -250,14 +250,14 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
if streaming_encoding and len(self.meta.video_keys) > 0:
|
if streaming_encoding and len(self.meta.video_keys) > 0:
|
||||||
streaming_enc = self._build_streaming_encoder(
|
streaming_enc = self._build_streaming_encoder(
|
||||||
self.meta.fps,
|
self.meta.fps,
|
||||||
camera_encoder_config,
|
camera_encoder,
|
||||||
encoder_queue_maxsize,
|
encoder_queue_maxsize,
|
||||||
encoder_threads,
|
encoder_threads,
|
||||||
)
|
)
|
||||||
self.writer = DatasetWriter(
|
self.writer = DatasetWriter(
|
||||||
meta=self.meta,
|
meta=self.meta,
|
||||||
root=self.root,
|
root=self.root,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_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,
|
||||||
@@ -299,13 +299,13 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
@staticmethod
|
@staticmethod
|
||||||
def _build_streaming_encoder(
|
def _build_streaming_encoder(
|
||||||
fps: int,
|
fps: int,
|
||||||
camera_encoder_config: VideoEncoderConfig | None,
|
camera_encoder: VideoEncoderConfig | None,
|
||||||
encoder_queue_maxsize: int,
|
encoder_queue_maxsize: int,
|
||||||
encoder_threads: int | None,
|
encoder_threads: int | None,
|
||||||
) -> StreamingVideoEncoder:
|
) -> StreamingVideoEncoder:
|
||||||
return StreamingVideoEncoder(
|
return StreamingVideoEncoder(
|
||||||
fps=fps,
|
fps=fps,
|
||||||
camera_encoder_config=camera_encoder_config,
|
camera_encoder=camera_encoder,
|
||||||
queue_maxsize=encoder_queue_maxsize,
|
queue_maxsize=encoder_queue_maxsize,
|
||||||
encoder_threads=encoder_threads,
|
encoder_threads=encoder_threads,
|
||||||
)
|
)
|
||||||
@@ -622,7 +622,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
image_writer_threads: int = 0,
|
image_writer_threads: int = 0,
|
||||||
video_backend: str | None = None,
|
video_backend: str | None = None,
|
||||||
batch_encoding_size: int = 1,
|
batch_encoding_size: int = 1,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
metadata_buffer_size: int = 10,
|
metadata_buffer_size: int = 10,
|
||||||
streaming_encoding: bool = False,
|
streaming_encoding: bool = False,
|
||||||
encoder_queue_maxsize: int = 30,
|
encoder_queue_maxsize: int = 30,
|
||||||
@@ -653,7 +653,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
video_backend: Video decoding backend (used when reading back).
|
video_backend: Video decoding backend (used when reading back).
|
||||||
batch_encoding_size: Number of episodes to accumulate before
|
batch_encoding_size: Number of episodes to accumulate before
|
||||||
batch-encoding videos. ``1`` means encode immediately.
|
batch-encoding videos. ``1`` means encode immediately.
|
||||||
camera_encoder_config: 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.
|
||||||
encoder_threads: Number of encoder threads (global). ``None``
|
encoder_threads: Number of encoder threads (global). ``None``
|
||||||
lets the codec decide.
|
lets the codec decide.
|
||||||
@@ -698,12 +698,12 @@ 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(
|
||||||
fps, camera_encoder_config, encoder_queue_maxsize, encoder_threads
|
fps, camera_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_config=camera_encoder_config,
|
camera_encoder=camera_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,
|
||||||
@@ -726,7 +726,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
force_cache_sync: bool = False,
|
force_cache_sync: bool = False,
|
||||||
video_backend: str | None = None,
|
video_backend: str | None = None,
|
||||||
batch_encoding_size: int = 1,
|
batch_encoding_size: int = 1,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | 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,
|
||||||
@@ -754,7 +754,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
|||||||
video_backend: Video decoding backend for reading back data.
|
video_backend: Video decoding backend for reading back data.
|
||||||
batch_encoding_size: Number of episodes to accumulate before
|
batch_encoding_size: Number of episodes to accumulate before
|
||||||
batch-encoding videos.
|
batch-encoding videos.
|
||||||
camera_encoder_config: 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.
|
||||||
encoder_threads: Number of encoder threads (global). ``None``
|
encoder_threads: Number of encoder threads (global). ``None``
|
||||||
lets the codec decide.
|
lets the codec decide.
|
||||||
@@ -802,12 +802,12 @@ 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_config, encoder_queue_maxsize, encoder_threads
|
obj.meta.fps, camera_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_config=camera_encoder_config,
|
camera_encoder=camera_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,
|
||||||
|
|||||||
@@ -335,17 +335,17 @@ def encode_video_frames(
|
|||||||
imgs_dir: Path | str,
|
imgs_dir: Path | str,
|
||||||
video_path: Path | str,
|
video_path: Path | str,
|
||||||
fps: int,
|
fps: int,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
encoder_threads: int | None = None,
|
encoder_threads: int | None = None,
|
||||||
*,
|
*,
|
||||||
log_level: int | None = av.logging.WARNING,
|
log_level: int | None = av.logging.WARNING,
|
||||||
overwrite: bool = False,
|
overwrite: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
|
"""More info on ffmpeg arguments tuning on `benchmark/video/README.md`"""
|
||||||
if camera_encoder_config is None:
|
if camera_encoder is None:
|
||||||
camera_encoder_config = camera_encoder_defaults()
|
camera_encoder = camera_encoder_defaults()
|
||||||
vcodec = camera_encoder_config.vcodec
|
vcodec = camera_encoder.vcodec
|
||||||
pix_fmt = camera_encoder_config.pix_fmt
|
pix_fmt = camera_encoder.pix_fmt
|
||||||
|
|
||||||
video_path = Path(video_path)
|
video_path = Path(video_path)
|
||||||
imgs_dir = Path(imgs_dir)
|
imgs_dir = Path(imgs_dir)
|
||||||
@@ -367,7 +367,7 @@ def encode_video_frames(
|
|||||||
with Image.open(input_list[0]) as dummy_image:
|
with Image.open(input_list[0]) as dummy_image:
|
||||||
width, height = dummy_image.size
|
width, height = dummy_image.size
|
||||||
|
|
||||||
video_options = camera_encoder_config.get_codec_options(encoder_threads, as_strings=True)
|
video_options = camera_encoder.get_codec_options(encoder_threads, as_strings=True)
|
||||||
|
|
||||||
# Set logging level
|
# Set logging level
|
||||||
if log_level is not None:
|
if log_level is not None:
|
||||||
@@ -638,14 +638,14 @@ class StreamingVideoEncoder:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
fps: int,
|
fps: int,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
queue_maxsize: int = 30,
|
queue_maxsize: int = 30,
|
||||||
encoder_threads: int | None = None,
|
encoder_threads: int | None = None,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Args:
|
Args:
|
||||||
fps: Frames per second for the output videos.
|
fps: Frames per second for the output videos.
|
||||||
camera_encoder_config: Video encoder settings applied to all cameras.
|
camera_encoder: Video encoder settings applied to all cameras.
|
||||||
When ``None``, :func:`camera_encoder_defaults` is used.
|
When ``None``, :func:`camera_encoder_defaults` is used.
|
||||||
encoder_threads: Number of encoder threads (global setting).
|
encoder_threads: Number of encoder threads (global setting).
|
||||||
``None`` lets the codec decide.
|
``None`` lets the codec decide.
|
||||||
@@ -653,7 +653,7 @@ class StreamingVideoEncoder:
|
|||||||
back-pressure drops frames.
|
back-pressure drops frames.
|
||||||
"""
|
"""
|
||||||
self.fps = fps
|
self.fps = fps
|
||||||
self._camera_encoder_config = camera_encoder_config or camera_encoder_defaults()
|
self._camera_encoder = camera_encoder or camera_encoder_defaults()
|
||||||
self._encoder_threads = encoder_threads
|
self._encoder_threads = encoder_threads
|
||||||
self.queue_maxsize = queue_maxsize
|
self.queue_maxsize = queue_maxsize
|
||||||
|
|
||||||
@@ -686,15 +686,15 @@ class StreamingVideoEncoder:
|
|||||||
temp_video_dir = Path(tempfile.mkdtemp(dir=temp_dir))
|
temp_video_dir = Path(tempfile.mkdtemp(dir=temp_dir))
|
||||||
video_path = temp_video_dir / f"{video_key.replace('/', '_')}_streaming.mp4"
|
video_path = temp_video_dir / f"{video_key.replace('/', '_')}_streaming.mp4"
|
||||||
|
|
||||||
vcodec = self._camera_encoder_config.vcodec
|
vcodec = self._camera_encoder.vcodec
|
||||||
codec_options = self._camera_encoder_config.get_codec_options(
|
codec_options = self._camera_encoder.get_codec_options(
|
||||||
self._encoder_threads, as_strings=True
|
self._encoder_threads, as_strings=True
|
||||||
)
|
)
|
||||||
encoder_thread = _CameraEncoderThread(
|
encoder_thread = _CameraEncoderThread(
|
||||||
video_path=video_path,
|
video_path=video_path,
|
||||||
fps=self.fps,
|
fps=self.fps,
|
||||||
vcodec=vcodec,
|
vcodec=vcodec,
|
||||||
pix_fmt=self._camera_encoder_config.pix_fmt,
|
pix_fmt=self._camera_encoder.pix_fmt,
|
||||||
codec_options=codec_options,
|
codec_options=codec_options,
|
||||||
frame_queue=frame_queue,
|
frame_queue=frame_queue,
|
||||||
result_queue=result_queue,
|
result_queue=result_queue,
|
||||||
@@ -905,13 +905,13 @@ def get_audio_info(video_path: Path | str) -> dict:
|
|||||||
|
|
||||||
def get_video_info(
|
def get_video_info(
|
||||||
video_path: Path | str,
|
video_path: Path | str,
|
||||||
camera_encoder_config: VideoEncoderConfig | None = None,
|
camera_encoder: VideoEncoderConfig | None = None,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
"""Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``.
|
"""Build the ``video.*`` / ``audio.*`` info dict persisted in ``info.json``.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
video_path: Path to the encoded video file to probe.
|
video_path: Path to the encoded video file to probe.
|
||||||
camera_encoder_config: If provided, record the exact encoder settings used to encode this
|
camera_encoder: If provided, record the exact encoder settings used to encode this
|
||||||
video. Stream-derived values take precedence — encoder fields are only written for keys
|
video. Stream-derived values take precedence — encoder fields are only written for keys
|
||||||
not already populated from the video file itself.
|
not already populated from the video file itself.
|
||||||
"""
|
"""
|
||||||
@@ -946,8 +946,8 @@ def get_video_info(
|
|||||||
video_info.update(**get_audio_info(video_path))
|
video_info.update(**get_audio_info(video_path))
|
||||||
|
|
||||||
# Add additional encoder configuration if provided
|
# Add additional encoder configuration if provided
|
||||||
if camera_encoder_config is not None:
|
if camera_encoder is not None:
|
||||||
for field_name, field_value in asdict(camera_encoder_config).items():
|
for field_name, field_value in asdict(camera_encoder).items():
|
||||||
# vcodec is already populated from the video stream
|
# vcodec is already populated from the video stream
|
||||||
if field_name == "vcodec":
|
if field_name == "vcodec":
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -28,11 +28,12 @@ import torch.nn.functional as F # noqa: N812
|
|||||||
import torch.utils.checkpoint
|
import torch.utils.checkpoint
|
||||||
from torch import Tensor
|
from torch import Tensor
|
||||||
|
|
||||||
from lerobot.policies.eo1.configuration_eo1 import EO1Config
|
|
||||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
|
||||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||||
|
|
||||||
|
from ..pretrained import PreTrainedPolicy
|
||||||
|
from .configuration_eo1 import EO1Config
|
||||||
|
|
||||||
if TYPE_CHECKING or _transformers_available:
|
if TYPE_CHECKING or _transformers_available:
|
||||||
from transformers.activations import ACT2FN
|
from transformers.activations import ACT2FN
|
||||||
from transformers.models.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
|
from transformers.models.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
|
||||||
|
|||||||
@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||||
from lerobot.policies.eo1.configuration_eo1 import EO1Config
|
|
||||||
from lerobot.processor import (
|
from lerobot.processor import (
|
||||||
AddBatchDimensionProcessorStep,
|
AddBatchDimensionProcessorStep,
|
||||||
ComplementaryDataProcessorStep,
|
ComplementaryDataProcessorStep,
|
||||||
@@ -44,6 +43,8 @@ from lerobot.utils.constants import (
|
|||||||
)
|
)
|
||||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||||
|
|
||||||
|
from .configuration_eo1 import EO1Config
|
||||||
|
|
||||||
if TYPE_CHECKING or _transformers_available:
|
if TYPE_CHECKING or _transformers_available:
|
||||||
from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor
|
from transformers.models.qwen2_5_vl import Qwen2_5_VLProcessor
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -17,10 +17,11 @@ import logging
|
|||||||
import torch
|
import torch
|
||||||
from torch import Tensor, nn
|
from torch import Tensor, nn
|
||||||
|
|
||||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
|
||||||
from lerobot.rewards.pretrained import PreTrainedRewardModel
|
|
||||||
from lerobot.utils.constants import OBS_IMAGE, REWARD
|
from lerobot.utils.constants import OBS_IMAGE, REWARD
|
||||||
|
|
||||||
|
from ..pretrained import PreTrainedRewardModel
|
||||||
|
from .configuration_classifier import RewardClassifierConfig
|
||||||
|
|
||||||
|
|
||||||
class ClassifierOutput:
|
class ClassifierOutput:
|
||||||
"""Wrapper for classifier outputs with additional metadata."""
|
"""Wrapper for classifier outputs with additional metadata."""
|
||||||
|
|||||||
@@ -25,7 +25,8 @@ from lerobot.processor import (
|
|||||||
policy_action_to_transition,
|
policy_action_to_transition,
|
||||||
transition_to_policy_action,
|
transition_to_policy_action,
|
||||||
)
|
)
|
||||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
|
||||||
|
from .configuration_classifier import RewardClassifierConfig
|
||||||
|
|
||||||
|
|
||||||
def make_classifier_processor(
|
def make_classifier_processor(
|
||||||
|
|||||||
@@ -22,9 +22,10 @@ import torch
|
|||||||
|
|
||||||
from lerobot.configs.rewards import RewardModelConfig
|
from lerobot.configs.rewards import RewardModelConfig
|
||||||
from lerobot.processor import PolicyAction, PolicyProcessorPipeline
|
from lerobot.processor import PolicyAction, PolicyProcessorPipeline
|
||||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
|
||||||
from lerobot.rewards.pretrained import PreTrainedRewardModel
|
from .classifier.configuration_classifier import RewardClassifierConfig
|
||||||
from lerobot.rewards.sarm.configuration_sarm import SARMConfig
|
from .pretrained import PreTrainedRewardModel
|
||||||
|
from .sarm.configuration_sarm import SARMConfig
|
||||||
|
|
||||||
|
|
||||||
def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
|
def get_reward_model_class(name: str) -> type[PreTrainedRewardModel]:
|
||||||
|
|||||||
@@ -58,9 +58,10 @@ import torch
|
|||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from lerobot.datasets import LeRobotDataset
|
from lerobot.datasets import LeRobotDataset
|
||||||
from lerobot.rewards.sarm.modeling_sarm import SARMRewardModel
|
|
||||||
from lerobot.rewards.sarm.processor_sarm import make_sarm_pre_post_processors
|
from .modeling_sarm import SARMRewardModel
|
||||||
from lerobot.rewards.sarm.sarm_utils import normalize_stage_tau
|
from .processor_sarm import make_sarm_pre_post_processors
|
||||||
|
from .sarm_utils import normalize_stage_tau
|
||||||
|
|
||||||
|
|
||||||
def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:
|
def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:
|
||||||
|
|||||||
@@ -32,13 +32,14 @@ import torch.nn as nn
|
|||||||
import torch.nn.functional as F # noqa: N812
|
import torch.nn.functional as F # noqa: N812
|
||||||
from torch import Tensor
|
from torch import Tensor
|
||||||
|
|
||||||
from lerobot.rewards.pretrained import PreTrainedRewardModel
|
from lerobot.utils.constants import OBS_STR
|
||||||
from lerobot.rewards.sarm.configuration_sarm import SARMConfig
|
|
||||||
from lerobot.rewards.sarm.sarm_utils import (
|
from ..pretrained import PreTrainedRewardModel
|
||||||
|
from .configuration_sarm import SARMConfig
|
||||||
|
from .sarm_utils import (
|
||||||
normalize_stage_tau,
|
normalize_stage_tau,
|
||||||
pad_state_to_max_dim,
|
pad_state_to_max_dim,
|
||||||
)
|
)
|
||||||
from lerobot.utils.constants import OBS_STR
|
|
||||||
|
|
||||||
|
|
||||||
class StageTransformer(nn.Module):
|
class StageTransformer(nn.Module):
|
||||||
|
|||||||
@@ -58,15 +58,16 @@ from lerobot.processor import (
|
|||||||
policy_action_to_transition,
|
policy_action_to_transition,
|
||||||
transition_to_policy_action,
|
transition_to_policy_action,
|
||||||
)
|
)
|
||||||
from lerobot.rewards.sarm.configuration_sarm import SARMConfig
|
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
||||||
from lerobot.rewards.sarm.sarm_utils import (
|
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
||||||
|
|
||||||
|
from .configuration_sarm import SARMConfig
|
||||||
|
from .sarm_utils import (
|
||||||
apply_rewind_augmentation,
|
apply_rewind_augmentation,
|
||||||
compute_absolute_indices,
|
compute_absolute_indices,
|
||||||
find_stage_and_tau,
|
find_stage_and_tau,
|
||||||
pad_state_to_max_dim,
|
pad_state_to_max_dim,
|
||||||
)
|
)
|
||||||
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
|
|
||||||
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
|
|
||||||
|
|
||||||
|
|
||||||
class SARMEncodingProcessorStep(ProcessorStep):
|
class SARMEncodingProcessorStep(ProcessorStep):
|
||||||
|
|||||||
@@ -332,7 +332,7 @@ def build_rollout_context(
|
|||||||
cfg.dataset.repo_id,
|
cfg.dataset.repo_id,
|
||||||
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=cfg.dataset.camera_encoder,
|
||||||
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,
|
||||||
encoder_threads=cfg.dataset.encoder_threads,
|
encoder_threads=cfg.dataset.encoder_threads,
|
||||||
@@ -367,7 +367,7 @@ def build_rollout_context(
|
|||||||
image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera
|
image_writer_threads=cfg.dataset.num_image_writer_threads_per_camera
|
||||||
* len(robot.cameras if hasattr(robot, "cameras") else []),
|
* len(robot.cameras if hasattr(robot, "cameras") else []),
|
||||||
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=cfg.dataset.camera_encoder,
|
||||||
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,
|
||||||
encoder_threads=cfg.dataset.encoder_threads,
|
encoder_threads=cfg.dataset.encoder_threads,
|
||||||
|
|||||||
@@ -54,8 +54,8 @@ Delete episodes and re-encode video segments with h264:
|
|||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--operation.type delete_episodes \
|
--operation.type delete_episodes \
|
||||||
--operation.episode_indices "[0, 2, 5]" \
|
--operation.episode_indices "[0, 2, 5]" \
|
||||||
--operation.camera_encoder_config.vcodec h264 \
|
--operation.camera_encoder.vcodec h264 \
|
||||||
--operation.camera_encoder_config.crf 23
|
--operation.camera_encoder.crf 23
|
||||||
|
|
||||||
Split dataset by fractions (pusht_train, pusht_val):
|
Split dataset by fractions (pusht_train, pusht_val):
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -87,8 +87,8 @@ Split dataset and re-encode video segments with h264:
|
|||||||
--repo_id lerobot/pusht \
|
--repo_id lerobot/pusht \
|
||||||
--operation.type split \
|
--operation.type split \
|
||||||
--operation.splits '{"train": 0.8, "val": 0.2}' \
|
--operation.splits '{"train": 0.8, "val": 0.2}' \
|
||||||
--operation.camera_encoder_config.vcodec h264 \
|
--operation.camera_encoder.vcodec h264 \
|
||||||
--operation.camera_encoder_config.crf 23
|
--operation.camera_encoder.crf 23
|
||||||
|
|
||||||
Merge multiple datasets:
|
Merge multiple datasets:
|
||||||
lerobot-edit-dataset \
|
lerobot-edit-dataset \
|
||||||
@@ -208,8 +208,7 @@ from pathlib import Path
|
|||||||
|
|
||||||
import draccus
|
import draccus
|
||||||
|
|
||||||
from lerobot.configs import parser
|
from lerobot.configs import VideoEncoderConfig, camera_encoder_defaults, parser
|
||||||
from lerobot.configs.video import VideoEncoderConfig, camera_encoder_defaults
|
|
||||||
from lerobot.datasets import (
|
from lerobot.datasets import (
|
||||||
LeRobotDataset,
|
LeRobotDataset,
|
||||||
convert_image_to_video_dataset,
|
convert_image_to_video_dataset,
|
||||||
@@ -235,14 +234,14 @@ class OperationConfig(draccus.ChoiceRegistry, abc.ABC):
|
|||||||
@dataclass
|
@dataclass
|
||||||
class DeleteEpisodesConfig(OperationConfig):
|
class DeleteEpisodesConfig(OperationConfig):
|
||||||
episode_indices: list[int] | None = None
|
episode_indices: list[int] | None = None
|
||||||
camera_encoder_config: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
||||||
|
|
||||||
|
|
||||||
@OperationConfig.register_subclass("split")
|
@OperationConfig.register_subclass("split")
|
||||||
@dataclass
|
@dataclass
|
||||||
class SplitConfig(OperationConfig):
|
class SplitConfig(OperationConfig):
|
||||||
splits: dict[str, float | list[int]] | None = None
|
splits: dict[str, float | list[int]] | None = None
|
||||||
camera_encoder_config: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
||||||
|
|
||||||
|
|
||||||
@OperationConfig.register_subclass("merge")
|
@OperationConfig.register_subclass("merge")
|
||||||
@@ -269,7 +268,7 @@ class ModifyTasksConfig(OperationConfig):
|
|||||||
@dataclass
|
@dataclass
|
||||||
class ConvertImageToVideoConfig(OperationConfig):
|
class ConvertImageToVideoConfig(OperationConfig):
|
||||||
output_dir: str | None = None
|
output_dir: str | None = None
|
||||||
camera_encoder_config: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
|
||||||
episode_indices: list[int] | None = None
|
episode_indices: list[int] | None = None
|
||||||
num_workers: int = 4
|
num_workers: int = 4
|
||||||
max_episodes_per_batch: int | None = None
|
max_episodes_per_batch: int | None = None
|
||||||
@@ -371,7 +370,7 @@ def handle_delete_episodes(cfg: EditDatasetConfig) -> None:
|
|||||||
episode_indices=cfg.operation.episode_indices,
|
episode_indices=cfg.operation.episode_indices,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
repo_id=output_repo_id,
|
repo_id=output_repo_id,
|
||||||
camera_encoder_config=cfg.operation.camera_encoder_config,
|
camera_encoder=cfg.operation.camera_encoder,
|
||||||
)
|
)
|
||||||
|
|
||||||
logging.info(f"Dataset saved to {output_dir}")
|
logging.info(f"Dataset saved to {output_dir}")
|
||||||
@@ -403,7 +402,7 @@ def handle_split(cfg: EditDatasetConfig) -> None:
|
|||||||
dataset,
|
dataset,
|
||||||
splits=cfg.operation.splits,
|
splits=cfg.operation.splits,
|
||||||
output_dir=cfg.new_root,
|
output_dir=cfg.new_root,
|
||||||
camera_encoder_config=cfg.operation.camera_encoder_config,
|
camera_encoder=cfg.operation.camera_encoder,
|
||||||
)
|
)
|
||||||
|
|
||||||
for split_name, split_ds in split_datasets.items():
|
for split_name, split_ds in split_datasets.items():
|
||||||
@@ -574,7 +573,7 @@ def handle_convert_image_to_video(cfg: EditDatasetConfig) -> None:
|
|||||||
dataset=dataset,
|
dataset=dataset,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
repo_id=output_repo_id,
|
repo_id=output_repo_id,
|
||||||
camera_encoder_config=getattr(cfg.operation, "camera_encoder_config", None)
|
camera_encoder=getattr(cfg.operation, "camera_encoder", None)
|
||||||
or camera_encoder_defaults(),
|
or camera_encoder_defaults(),
|
||||||
episode_indices=getattr(cfg.operation, "episode_indices", None),
|
episode_indices=getattr(cfg.operation, "episode_indices", None),
|
||||||
num_workers=getattr(cfg.operation, "num_workers", 4),
|
num_workers=getattr(cfg.operation, "num_workers", 4),
|
||||||
|
|||||||
@@ -79,9 +79,9 @@ lerobot-record \\
|
|||||||
--dataset.single_task="Grab the cube" \\
|
--dataset.single_task="Grab the cube" \\
|
||||||
--dataset.streaming_encoding=true \\
|
--dataset.streaming_encoding=true \\
|
||||||
--dataset.encoder_threads=2 \\
|
--dataset.encoder_threads=2 \\
|
||||||
--dataset.camera_encoder_config.vcodec=h264 \\
|
--dataset.camera_encoder.vcodec=h264 \\
|
||||||
--dataset.camera_encoder_config.preset=fast \\
|
--dataset.camera_encoder.preset=fast \\
|
||||||
--dataset.camera_encoder_config.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\
|
--dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2} \\
|
||||||
--display_data=true
|
--display_data=true
|
||||||
```
|
```
|
||||||
"""
|
"""
|
||||||
@@ -398,7 +398,7 @@ def record(
|
|||||||
cfg.dataset.repo_id,
|
cfg.dataset.repo_id,
|
||||||
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=cfg.dataset.camera_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,
|
||||||
@@ -427,7 +427,7 @@ def record(
|
|||||||
image_writer_processes=cfg.dataset.num_image_writer_processes,
|
image_writer_processes=cfg.dataset.num_image_writer_processes,
|
||||||
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=cfg.dataset.camera_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,
|
||||||
@@ -441,7 +441,7 @@ def record(
|
|||||||
|
|
||||||
if not cfg.dataset.streaming_encoding:
|
if not cfg.dataset.streaming_encoding:
|
||||||
logging.info(
|
logging.info(
|
||||||
"Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.camera_encoder_config.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding"
|
"Streaming encoding is disabled. If you have capable hardware, consider enabling it for way faster episode saving. --dataset.streaming_encoding=true --dataset.encoder_threads=2 # --dataset.camera_encoder.vcodec=auto. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding"
|
||||||
)
|
)
|
||||||
|
|
||||||
with VideoEncodingManager(dataset):
|
with VideoEncodingManager(dataset):
|
||||||
|
|||||||
@@ -120,6 +120,18 @@ Usage examples
|
|||||||
--dataset.repo_id=user/rollout_sentry_data \\
|
--dataset.repo_id=user/rollout_sentry_data \\
|
||||||
--dataset.single_task="patrol" \\
|
--dataset.single_task="patrol" \\
|
||||||
--resume=true
|
--resume=true
|
||||||
|
|
||||||
|
# Rollout with custom video encoding parameters
|
||||||
|
lerobot-rollout \\
|
||||||
|
--strategy.type=base \\
|
||||||
|
--policy.path=lerobot/act_koch_real \\
|
||||||
|
--robot.type=koch_follower \\
|
||||||
|
--robot.port=/dev/ttyACM0 \\
|
||||||
|
--task="pick up cube" --duration=60 \\
|
||||||
|
--display_data=true \\
|
||||||
|
--dataset.camera_encoder.vcodec=h264 \\
|
||||||
|
--dataset.camera_encoder.preset=fast \\
|
||||||
|
--dataset.camera_encoder.extra_options={"tune": "film", "profile:v": "high", "bf": 2}
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
|||||||
@@ -25,9 +25,10 @@ from typing import Any
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from lerobot.transport import services_pb2
|
|
||||||
from lerobot.utils.transition import Transition
|
from lerobot.utils.transition import Transition
|
||||||
|
|
||||||
|
from . import services_pb2
|
||||||
|
|
||||||
# FIX for protobuf: Assign the enum to a variable and ignore the type error once
|
# FIX for protobuf: Assign the enum to a variable and ignore the type error once
|
||||||
TransferState = services_pb2.TransferState # type: ignore[attr-defined]
|
TransferState = services_pb2.TransferState # type: ignore[attr-defined]
|
||||||
|
|
||||||
|
|||||||
@@ -1247,7 +1247,7 @@ def test_convert_image_to_video_dataset(tmp_path):
|
|||||||
dataset=source_dataset,
|
dataset=source_dataset,
|
||||||
output_dir=output_dir,
|
output_dir=output_dir,
|
||||||
repo_id="lerobot/pusht_video",
|
repo_id="lerobot/pusht_video",
|
||||||
camera_encoder_config=VideoEncoderConfig(
|
camera_encoder=VideoEncoderConfig(
|
||||||
vcodec="libsvtav1",
|
vcodec="libsvtav1",
|
||||||
pix_fmt="yuv420p",
|
pix_fmt="yuv420p",
|
||||||
g=2,
|
g=2,
|
||||||
|
|||||||
@@ -53,8 +53,8 @@ def _make_frame(features: dict, task: str = "Dummy task") -> dict:
|
|||||||
# ── Existing encode_video_worker tests ───────────────────────────────
|
# ── Existing encode_video_worker tests ───────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
def test_encode_video_worker_forwards_camera_encoder_config(tmp_path):
|
def test_encode_video_worker_forwards_camera_encoder(tmp_path):
|
||||||
"""_encode_video_worker forwards camera_encoder_config to encode_video_frames."""
|
"""_encode_video_worker forwards camera_encoder to encode_video_frames."""
|
||||||
video_key = "observation.images.laptop"
|
video_key = "observation.images.laptop"
|
||||||
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
||||||
img_dir = tmp_path / Path(fpath).parent
|
img_dir = tmp_path / Path(fpath).parent
|
||||||
@@ -74,16 +74,16 @@ def test_encode_video_worker_forwards_camera_encoder_config(tmp_path):
|
|||||||
0,
|
0,
|
||||||
tmp_path,
|
tmp_path,
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=VideoEncoderConfig(vcodec="h264", preset=None),
|
camera_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
|
||||||
encoder_threads=4,
|
encoder_threads=4,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert captured_kwargs["camera_encoder_config"].vcodec == "h264"
|
assert captured_kwargs["camera_encoder"].vcodec == "h264"
|
||||||
assert captured_kwargs["encoder_threads"] == 4
|
assert captured_kwargs["encoder_threads"] == 4
|
||||||
|
|
||||||
|
|
||||||
def test_encode_video_worker_default_camera_encoder_config(tmp_path):
|
def test_encode_video_worker_default_camera_encoder(tmp_path):
|
||||||
"""_encode_video_worker passes None camera_encoder_config which encode_video_frames defaults."""
|
"""_encode_video_worker passes None camera_encoder which encode_video_frames defaults."""
|
||||||
video_key = "observation.images.laptop"
|
video_key = "observation.images.laptop"
|
||||||
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
|
||||||
img_dir = tmp_path / Path(fpath).parent
|
img_dir = tmp_path / Path(fpath).parent
|
||||||
@@ -100,7 +100,7 @@ def test_encode_video_worker_default_camera_encoder_config(tmp_path):
|
|||||||
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
|
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
|
||||||
_encode_video_worker(video_key, 0, tmp_path, fps=30)
|
_encode_video_worker(video_key, 0, tmp_path, fps=30)
|
||||||
|
|
||||||
assert captured_kwargs["camera_encoder_config"] is None
|
assert captured_kwargs["camera_encoder"] is None
|
||||||
assert captured_kwargs["encoder_threads"] is None
|
assert captured_kwargs["encoder_threads"] is None
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -179,7 +179,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.laptop"]
|
video_keys = [f"{OBS_IMAGES}.laptop"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -211,7 +211,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
|
video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -239,7 +239,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -267,7 +267,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -315,7 +315,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -352,7 +352,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
|
video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
@@ -391,7 +391,7 @@ class TestStreamingVideoEncoder:
|
|||||||
)
|
)
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=cfg,
|
camera_encoder=cfg,
|
||||||
encoder_threads=2,
|
encoder_threads=2,
|
||||||
)
|
)
|
||||||
assert encoder._encoder_threads == 2
|
assert encoder._encoder_threads == 2
|
||||||
@@ -430,7 +430,7 @@ class TestStreamingVideoEncoder:
|
|||||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||||
encoder = StreamingVideoEncoder(
|
encoder = StreamingVideoEncoder(
|
||||||
fps=30,
|
fps=30,
|
||||||
camera_encoder_config=self._make_encoder_config(
|
camera_encoder=self._make_encoder_config(
|
||||||
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
|
||||||
),
|
),
|
||||||
queue_maxsize=1,
|
queue_maxsize=1,
|
||||||
|
|||||||
@@ -337,7 +337,7 @@ def _encode_video(
|
|||||||
) -> Path:
|
) -> Path:
|
||||||
imgs_dir = path.parent / f"imgs_{path.stem}"
|
imgs_dir = path.parent / f"imgs_{path.stem}"
|
||||||
_write_frames(imgs_dir, num_frames=num_frames)
|
_write_frames(imgs_dir, num_frames=num_frames)
|
||||||
encode_video_frames(imgs_dir, path, fps=fps, camera_encoder_config=cfg, overwrite=True)
|
encode_video_frames(imgs_dir, path, fps=fps, camera_encoder=cfg, overwrite=True)
|
||||||
return path
|
return path
|
||||||
|
|
||||||
|
|
||||||
@@ -377,7 +377,7 @@ class TestGetVideoInfo:
|
|||||||
def test_merges_encoder_config_as_video_prefixed_entries(self):
|
def test_merges_encoder_config_as_video_prefixed_entries(self):
|
||||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||||
|
|
||||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder_config=cfg)
|
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
|
||||||
|
|
||||||
assert info["video.g"] == 2
|
assert info["video.g"] == 2
|
||||||
assert info["video.crf"] == 30
|
assert info["video.crf"] == 30
|
||||||
@@ -390,7 +390,7 @@ class TestGetVideoInfo:
|
|||||||
def test_stream_derived_keys_take_precedence_over_config(self):
|
def test_stream_derived_keys_take_precedence_over_config(self):
|
||||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
|
cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||||
|
|
||||||
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder_config=cfg)
|
info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
|
||||||
|
|
||||||
assert info["video.codec"] # populated from stream, not from config's vcodec
|
assert info["video.codec"] # populated from stream, not from config's vcodec
|
||||||
assert info["video.pix_fmt"] == "yuv420p"
|
assert info["video.pix_fmt"] == "yuv420p"
|
||||||
@@ -453,7 +453,7 @@ class TestEncodeVideoFrames:
|
|||||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
|
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
|
||||||
video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg)
|
video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg)
|
||||||
|
|
||||||
info = get_video_info(video_path, camera_encoder_config=cfg)
|
info = get_video_info(video_path, camera_encoder=cfg)
|
||||||
|
|
||||||
# Stream-derived
|
# Stream-derived
|
||||||
assert info["video.height"] == 64
|
assert info["video.height"] == 64
|
||||||
@@ -535,7 +535,7 @@ class TestEncoderConfigPersistence:
|
|||||||
def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
|
def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
|
||||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||||
dataset = empty_lerobot_dataset_factory(
|
dataset = empty_lerobot_dataset_factory(
|
||||||
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder_config=cfg
|
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
|
||||||
)
|
)
|
||||||
|
|
||||||
_add_frames(dataset, num_frames=4)
|
_add_frames(dataset, num_frames=4)
|
||||||
@@ -558,7 +558,7 @@ class TestEncoderConfigPersistence:
|
|||||||
def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory):
|
def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory):
|
||||||
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
|
||||||
dataset = empty_lerobot_dataset_factory(
|
dataset = empty_lerobot_dataset_factory(
|
||||||
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder_config=cfg
|
root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
|
||||||
)
|
)
|
||||||
|
|
||||||
_add_frames(dataset, num_frames=4)
|
_add_frames(dataset, num_frames=4)
|
||||||
|
|||||||
Reference in New Issue
Block a user