admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(depth encoding): updating and cleaning video/depth encoding tests

Browse Source
This commit is contained in:
CarolinePascal
2026-06-12 19:17:59 +02:00
parent 291d2e982c
commit 6b2d5cec4c
3 changed files with 325 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/fixtures/dataset_factories.py
Vendored
+5 -8
View File
@@ -38,6 +38,8 @@ from lerobot.datasets.utils import (
DEFAULT_VIDEO_PATH,
DatasetInfo,
)
from lerobot.datasets.video_utils import encode_video_frames
from lerobot.utils.constants import DEFAULT_FEATURES
from tests.fixtures.constants import (
DEFAULT_FPS,
DUMMY_CAMERA_FEATURES,
@@ -45,13 +47,9 @@ from tests.fixtures.constants import (
DUMMY_REPO_ID,
DUMMY_ROBOT_TYPE,
)
from lerobot.datasets.video_utils import encode_video_frames
from lerobot.utils.constants import DEFAULT_FEATURES
def add_frames(
dataset: LeRobotDataset, num_frames: int
) -> None:
def add_frames(dataset: LeRobotDataset, num_frames: int) -> None:
"""Append ``num_frames`` synthetic frames to ``dataset``.
Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for
@@ -59,9 +57,8 @@ def add_frames(
and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp,
frame_index, ...) are auto-populated by ``add_frame`` and skipped here.
"""
if video_keys is None:
video_keys = dataset.meta.video_keys
depth_keys = set(dataset.meta.depth_keys)
video_keys = dataset.meta.video_keys
depth_keys = dataset.meta.depth_keys
# Smooth gradient base reused per (H, W) to keep depth frames cheap to
# encode (HEVC Main 12 hates white noise).
_depth_base_cache: dict[tuple[int, int], np.ndarray] = {}