mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-08 10:32:00 +00:00
Pepijn
122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
import json
|
|
import struct
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from lerobot.datasets.episode_video_streaming import assert_hf_hub_range_cache_branch
|
|
from lerobot.datasets.mp4 import (
|
|
_box,
|
|
_co64,
|
|
_dinf,
|
|
_hdlr,
|
|
_mdhd,
|
|
_mvhd,
|
|
_stco,
|
|
_stsc_one_sample_per_chunk,
|
|
_stss,
|
|
_stsz,
|
|
_stts,
|
|
_tkhd,
|
|
_vmhd,
|
|
parse_mp4_index,
|
|
synthesize_mp4,
|
|
)
|
|
|
|
|
|
def _minimal_mp4(sample_offsets: list[int], *, use_co64: bool = False) -> bytes:
|
|
ftyp = _box(b"ftyp", b"isom\0\0\2\0isomiso2mp41")
|
|
sizes = np.array([10, 10, 10], dtype=np.int64)
|
|
durations = np.array([1000, 1000, 1000], dtype=np.int64)
|
|
stsd_body = struct.pack(">II", 0, 1) + struct.pack(">I4s", 16, b"avc1") + b"\0" * 8
|
|
offsets = _co64(sample_offsets) if use_co64 else _stco(sample_offsets)
|
|
stbl = _box(
|
|
b"stbl",
|
|
_box(b"stsd", stsd_body)
|
|
+ _stts(durations)
|
|
+ _stsc_one_sample_per_chunk(len(sizes))
|
|
+ _stsz(sizes)
|
|
+ offsets
|
|
+ _stss(np.array([1], dtype=np.int64)),
|
|
)
|
|
minf = _box(b"minf", _vmhd() + _dinf() + stbl)
|
|
mdia = _box(b"mdia", _mdhd(1000, 3000) + _hdlr() + minf)
|
|
trak = _box(b"trak", _tkhd(1, 3000, 64, 48) + mdia)
|
|
moov = _box(b"moov", _mvhd(1000, 3000, 2) + trak)
|
|
mdat_payload_start = 10_000
|
|
free_size = mdat_payload_start - 8 - len(ftyp) - len(moov)
|
|
assert free_size >= 8
|
|
free = _box(b"free", b"\0" * (free_size - 8))
|
|
return ftyp + moov + free + _box(b"mdat", b"x" * 128)
|
|
|
|
|
|
def test_episode_slice_uses_min_max_sample_offsets_for_reordered_chunks():
|
|
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
|
|
|
|
sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
|
|
|
|
assert sample_slice.byte_offset == 10_000
|
|
assert sample_slice.byte_length == 60
|
|
assert sample_slice.sample_lo == 0
|
|
assert sample_slice.sample_hi == 2
|
|
|
|
|
|
def test_synthesized_mp4_rebases_one_chunk_per_sample_offsets():
|
|
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
|
|
sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
|
|
|
|
mini = synthesize_mp4(mp4, sample_slice, b"x" * sample_slice.byte_length)
|
|
mini_index = parse_mp4_index("mini.mp4", mini)
|
|
|
|
expected = np.array([0, 50, 25], dtype=np.int64) + mini_index.mdat_payload_offset
|
|
np.testing.assert_array_equal(mini_index.sample_offsets, expected)
|
|
np.testing.assert_array_equal(mini_index.sample_sizes, np.array([10, 10, 10]))
|
|
|
|
|
|
def test_parser_accepts_co64_chunk_offsets():
|
|
mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025], use_co64=True))
|
|
|
|
np.testing.assert_array_equal(mp4.sample_offsets, np.array([10_000, 10_050, 10_025]))
|
|
|
|
|
|
def test_hf_hub_branch_assertion_accepts_requested_revision(monkeypatch):
|
|
class FakeDist:
|
|
def read_text(self, name):
|
|
assert name == "direct_url.json"
|
|
return json.dumps(
|
|
{
|
|
"url": "https://github.com/huggingface/huggingface_hub.git",
|
|
"vcs_info": {"requested_revision": "feat/hffs-cache-cdn-range-reads"},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"lerobot.datasets.episode_video_streaming.metadata.distribution", lambda _: FakeDist()
|
|
)
|
|
|
|
assert_hf_hub_range_cache_branch()
|
|
|
|
|
|
def test_hf_hub_branch_assertion_rejects_plain_install(monkeypatch):
|
|
class FakeDist:
|
|
def read_text(self, name):
|
|
assert name == "direct_url.json"
|
|
return json.dumps({"url": "https://github.com/huggingface/huggingface_hub.git"})
|
|
|
|
monkeypatch.setattr(
|
|
"lerobot.datasets.episode_video_streaming.metadata.distribution", lambda _: FakeDist()
|
|
)
|
|
|
|
with pytest.raises(AssertionError):
|
|
assert_hf_hub_range_cache_branch()
|