wehub-resource-sync
385 lines
15 KiB
Python
385 lines
15 KiB
Python
"""Tests for the low-level helpers in `rerun.experimental.dataloader._decoders`."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
from rerun.experimental.dataloader import Field
|
|
from rerun.experimental.dataloader._decoders import (
|
|
VideoFrameDecoder,
|
|
_avcc_to_annex_b,
|
|
_flatten_blob,
|
|
_h264_annex_b_has_idr,
|
|
_hevc_annex_b_has_irap,
|
|
_is_annex_b,
|
|
_is_av1_keyframe_packet,
|
|
_unwrap_to_numpy,
|
|
)
|
|
from rerun.experimental.dataloader._utils import _field_index_range, _prior_keyframe
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("data", "expected"),
|
|
[
|
|
(b"\x00\x00\x00\x01\xab\xcd", True), # 4-byte start code
|
|
(b"\x00\x00\x01\xab\xcd", True), # 3-byte short start code
|
|
(b"\x00\x00\x00\x01", True), # exactly the start code
|
|
(b"\x00\x00\x01", True), # exactly the short start code
|
|
(b"\x00\x00\x02\xab", False),
|
|
(b"\xab\xcd\xef\x01", False),
|
|
(b"", False),
|
|
(b"\x00", False),
|
|
(b"\x00\x00", False),
|
|
],
|
|
)
|
|
def test_is_annex_b(data: bytes, expected: bool) -> None:
|
|
assert _is_annex_b(data) is expected
|
|
|
|
|
|
def _make_obu_header(obu_type: int) -> int:
|
|
"""Return a byte whose OBU-type field (bits [3:6]) matches *obu_type*."""
|
|
return (obu_type & 0xF) << 3
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("obu_type", "expected"),
|
|
[
|
|
(1, True), # OBU_SEQUENCE_HEADER
|
|
(2, True), # OBU_TEMPORAL_DELIMITER
|
|
(3, False), # OBU_FRAME_HEADER
|
|
(6, False), # OBU_FRAME
|
|
(0, False),
|
|
(15, False),
|
|
],
|
|
)
|
|
def test_is_av1_keyframe_packet(obu_type: int, expected: bool) -> None:
|
|
sample = bytes([_make_obu_header(obu_type), 0x00, 0x00])
|
|
assert _is_av1_keyframe_packet(sample) is expected
|
|
|
|
|
|
def test_is_av1_keyframe_packet_empty() -> None:
|
|
assert _is_av1_keyframe_packet(b"") is False
|
|
|
|
|
|
def test_is_av1_keyframe_packet_ignores_low_bits() -> None:
|
|
# Low three bits (extension/has-size/reserved) must not affect detection.
|
|
header = _make_obu_header(1) | 0b111
|
|
assert _is_av1_keyframe_packet(bytes([header])) is True
|
|
|
|
|
|
def _avcc_encode(nal_units: list[bytes], nal_length_size: int = 4) -> bytes:
|
|
out = bytearray()
|
|
for unit in nal_units:
|
|
out.extend(len(unit).to_bytes(nal_length_size, "big"))
|
|
out.extend(unit)
|
|
return bytes(out)
|
|
|
|
|
|
def _h264_annex_b(nal_units: list[tuple[int, bytes]], use_4byte: bool = True) -> bytes:
|
|
"""Build an Annex B H.264 stream from `(nal_unit_type, payload)` pairs."""
|
|
start = b"\x00\x00\x00\x01" if use_4byte else b"\x00\x00\x01"
|
|
out = bytearray()
|
|
for nal_type, payload in nal_units:
|
|
out.extend(start)
|
|
# nal_ref_idc=3, forbidden_zero_bit=0
|
|
out.append((3 << 5) | (nal_type & 0x1F))
|
|
out.extend(payload)
|
|
return bytes(out)
|
|
|
|
|
|
def _hevc_annex_b(nal_units: list[tuple[int, bytes]], use_4byte: bool = True) -> bytes:
|
|
"""Build an Annex B HEVC stream from `(nal_unit_type, payload)` pairs."""
|
|
start = b"\x00\x00\x00\x01" if use_4byte else b"\x00\x00\x01"
|
|
out = bytearray()
|
|
for nal_type, payload in nal_units:
|
|
out.extend(start)
|
|
# forbidden_zero_bit=0, layer_id=0, temporal_id_plus1=1
|
|
out.append((nal_type & 0x3F) << 1)
|
|
out.append(0x01)
|
|
out.extend(payload)
|
|
return bytes(out)
|
|
|
|
|
|
@pytest.mark.parametrize("use_4byte", [True, False])
|
|
def test_h264_annex_b_has_idr_simple(use_4byte: bool) -> None:
|
|
sample = _h264_annex_b([(5, b"\xaa\xbb\xcc")], use_4byte=use_4byte)
|
|
assert _h264_annex_b_has_idr(sample) is True
|
|
|
|
|
|
def test_h264_annex_b_has_idr_after_sps_pps() -> None:
|
|
sample = _h264_annex_b([
|
|
(7, b"\x42\xc0\x1f"), # SPS
|
|
(8, b"\xce\x38\x80"), # PPS
|
|
(5, b"\x88\x84"), # IDR
|
|
])
|
|
assert _h264_annex_b_has_idr(sample) is True
|
|
|
|
|
|
def test_h264_annex_b_has_idr_after_aud_sei() -> None:
|
|
sample = _h264_annex_b([
|
|
(9, b"\x10"), # AUD
|
|
(6, b"\x01\x80"), # SEI
|
|
(5, b"\x88"), # IDR
|
|
])
|
|
assert _h264_annex_b_has_idr(sample) is True
|
|
|
|
|
|
def test_h264_annex_b_has_idr_p_slice_only() -> None:
|
|
sample = _h264_annex_b([(1, b"\xab\xcd\xef")])
|
|
assert _h264_annex_b_has_idr(sample) is False
|
|
|
|
|
|
def test_h264_annex_b_has_idr_empty() -> None:
|
|
assert _h264_annex_b_has_idr(b"") is False
|
|
|
|
|
|
@pytest.mark.parametrize("use_4byte", [True, False])
|
|
def test_hevc_annex_b_has_irap_simple(use_4byte: bool) -> None:
|
|
sample = _hevc_annex_b([(19, b"\xaa\xbb\xcc")], use_4byte=use_4byte)
|
|
assert _hevc_annex_b_has_irap(sample) is True
|
|
|
|
|
|
@pytest.mark.parametrize("nal_type", [16, 17, 18, 19, 20, 21, 22, 23])
|
|
def test_hevc_annex_b_has_irap_all_irap_types(nal_type: int) -> None:
|
|
sample = _hevc_annex_b([(nal_type, b"\xaa")])
|
|
assert _hevc_annex_b_has_irap(sample) is True
|
|
|
|
|
|
@pytest.mark.parametrize("nal_type", [1, 15, 24, 32, 35]) # TRAIL_R, RSV_VCL_*, just-out-of-IRAP, VPS, AUD
|
|
def test_hevc_annex_b_has_irap_non_irap_types(nal_type: int) -> None:
|
|
sample = _hevc_annex_b([(nal_type, b"\xaa")])
|
|
assert _hevc_annex_b_has_irap(sample) is False
|
|
|
|
|
|
def test_hevc_annex_b_has_irap_after_vps_sps_pps() -> None:
|
|
sample = _hevc_annex_b([
|
|
(32, b"\xaa"), # VPS
|
|
(33, b"\xbb"), # SPS
|
|
(34, b"\xcc"), # PPS
|
|
(19, b"\xdd"), # IDR_W_RADL
|
|
])
|
|
assert _hevc_annex_b_has_irap(sample) is True
|
|
|
|
|
|
def test_hevc_annex_b_has_irap_empty() -> None:
|
|
assert _hevc_annex_b_has_irap(b"") is False
|
|
|
|
|
|
def test_avcc_to_annex_b_single_unit() -> None:
|
|
unit = b"\x67\x42\xc0\x1f"
|
|
result = _avcc_to_annex_b(_avcc_encode([unit]))
|
|
assert result == b"\x00\x00\x00\x01" + unit
|
|
|
|
|
|
def test_avcc_to_annex_b_multiple_units() -> None:
|
|
units = [b"\x67\x42\xc0\x1f", b"\x68\xce\x38\x80", b"\x65\x88\x84"]
|
|
result = _avcc_to_annex_b(_avcc_encode(units))
|
|
expected = b"".join(b"\x00\x00\x00\x01" + u for u in units)
|
|
assert result == expected
|
|
|
|
|
|
def test_avcc_to_annex_b_length_size_2() -> None:
|
|
units = [b"\xaa\xbb", b"\xcc\xdd\xee"]
|
|
result = _avcc_to_annex_b(_avcc_encode(units, nal_length_size=2), nal_length_size=2)
|
|
expected = b"".join(b"\x00\x00\x00\x01" + u for u in units)
|
|
assert result == expected
|
|
|
|
|
|
def test_avcc_to_annex_b_truncated_stops_early() -> None:
|
|
# Well-formed first unit, then a length that claims more data than is left.
|
|
first = b"\x67\x42\xc0"
|
|
buf = len(first).to_bytes(4, "big") + first + (0xFF).to_bytes(4, "big") + b"\x00\x01"
|
|
result = _avcc_to_annex_b(buf)
|
|
assert result == b"\x00\x00\x00\x01" + first
|
|
|
|
|
|
def test_avcc_to_annex_b_empty() -> None:
|
|
assert _avcc_to_annex_b(b"") == b""
|
|
|
|
|
|
def test_unwrap_plain_numeric() -> None:
|
|
arr = pa.array([1.0, 2.0, 3.0], type=pa.float64())
|
|
np.testing.assert_array_equal(_unwrap_to_numpy(arr), np.array([1.0, 2.0, 3.0]))
|
|
|
|
|
|
def test_unwrap_list_float() -> None:
|
|
arr = pa.array([[1.0, 2.0], [3.0, 4.0, 5.0]], type=pa.list_(pa.float64()))
|
|
# Non-ragged requirement isn't enforced — the function returns the flattened values.
|
|
np.testing.assert_array_equal(_unwrap_to_numpy(arr), np.array([1.0, 2.0, 3.0, 4.0, 5.0]))
|
|
|
|
|
|
def test_unwrap_fixed_size_list() -> None:
|
|
arr = pa.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], type=pa.list_(pa.float32(), 3))
|
|
result = _unwrap_to_numpy(arr)
|
|
np.testing.assert_array_equal(result, np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.float32))
|
|
|
|
|
|
def test_unwrap_nested_list() -> None:
|
|
arr = pa.array([[[1.0, 2.0], [3.0]], [[4.0]]], type=pa.list_(pa.list_(pa.float64())))
|
|
np.testing.assert_array_equal(_unwrap_to_numpy(arr), np.array([1.0, 2.0, 3.0, 4.0]))
|
|
|
|
|
|
def test_unwrap_result_is_writeable() -> None:
|
|
# Torch requires writeable arrays downstream.
|
|
arr = pa.array([1, 2, 3], type=pa.int32())
|
|
result = _unwrap_to_numpy(arr)
|
|
assert result.flags.writeable
|
|
|
|
|
|
def test_flatten_blob_list_of_list_uint8_single_row() -> None:
|
|
arr = pa.array([[[1, 2, 3, 4]]], type=pa.list_(pa.list_(pa.uint8())))
|
|
result = _flatten_blob(arr, 0)
|
|
np.testing.assert_array_equal(result, np.array([1, 2, 3, 4], dtype=np.uint8))
|
|
|
|
|
|
def test_flatten_blob_list_of_list_uint8_concatenates_inner_rows() -> None:
|
|
# Row 0 has two inner lists, which should be concatenated.
|
|
arr = pa.array(
|
|
[[[1, 2], [3]], [[10, 20, 30]]],
|
|
type=pa.list_(pa.list_(pa.uint8())),
|
|
)
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 0), np.array([1, 2, 3], dtype=np.uint8))
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 1), np.array([10, 20, 30], dtype=np.uint8))
|
|
|
|
|
|
def test_flatten_blob_list_of_binary() -> None:
|
|
arr = pa.array([[b"hello"], [b"world!"]], type=pa.list_(pa.binary()))
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 0), np.frombuffer(b"hello", dtype=np.uint8))
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 1), np.frombuffer(b"world!", dtype=np.uint8))
|
|
|
|
|
|
def test_flatten_blob_list_of_large_binary() -> None:
|
|
arr = pa.array([[b"abc"], [b"defghi"]], type=pa.list_(pa.large_binary()))
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 0), np.frombuffer(b"abc", dtype=np.uint8))
|
|
np.testing.assert_array_equal(_flatten_blob(arr, 1), np.frombuffer(b"defghi", dtype=np.uint8))
|
|
|
|
|
|
def test_flatten_blob_binary_respects_offsets() -> None:
|
|
# The binary-path reads raw offsets, make sure subsequent rows don't leak into row 0.
|
|
arr = pa.array(
|
|
[[b"AAAA"], [b"BB"], [b"CCCCCC"]],
|
|
type=pa.list_(pa.binary()),
|
|
)
|
|
for row, expected in enumerate([b"AAAA", b"BB", b"CCCCCC"]):
|
|
np.testing.assert_array_equal(_flatten_blob(arr, row), np.frombuffer(expected, dtype=np.uint8))
|
|
|
|
|
|
def test_video_frame_decoder_returns_none_without_keyframe() -> None:
|
|
"""`decode` returns `None` when the prefetched window contains no keyframe."""
|
|
p_slice_only = _h264_annex_b([(1, b"\xab\xcd\xef\x01\x02\x03")])
|
|
raw = pa.chunked_array([pa.array([[p_slice_only]], type=pa.list_(pa.binary()))])
|
|
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=2)
|
|
assert decoder.decode(raw, 0, "seg") is None
|
|
|
|
|
|
def test_video_frame_decoder_is_keyframe_h264() -> None:
|
|
p_slice = _h264_annex_b([(1, b"\xab\xcd")])
|
|
idr = _h264_annex_b([(5, b"\x88")])
|
|
decoder = VideoFrameDecoder(codec="h264")
|
|
assert decoder._is_keyframe(p_slice) is False
|
|
assert decoder._is_keyframe(idr) is True
|
|
|
|
|
|
def test_video_frame_decoder_is_keyframe_hevc() -> None:
|
|
non_irap = _hevc_annex_b([(1, b"\xaa")])
|
|
irap = _hevc_annex_b([(19, b"\xaa")])
|
|
decoder = VideoFrameDecoder(codec="hevc")
|
|
assert decoder._is_keyframe(non_irap) is False
|
|
assert decoder._is_keyframe(irap) is True
|
|
|
|
|
|
def test_video_frame_decoder_is_keyframe_unknown_codec_returns_none() -> None:
|
|
assert VideoFrameDecoder(codec="vp9")._is_keyframe(b"\x00") is None
|
|
|
|
|
|
def test_video_frame_decoder_has_keyframe_h264() -> None:
|
|
p_slice = _h264_annex_b([(1, b"\xab\xcd")])
|
|
idr = _h264_annex_b([(7, b"\x42\xc0\x1f"), (8, b"\xce\x38"), (5, b"\x88")])
|
|
decoder = VideoFrameDecoder(codec="h264")
|
|
assert decoder._has_keyframe([]) is False
|
|
assert decoder._has_keyframe([p_slice]) is False
|
|
assert decoder._has_keyframe([p_slice, idr]) is True
|
|
|
|
|
|
def test_video_frame_decoder_has_keyframe_unknown_codec_trusts_decoder() -> None:
|
|
# Unknown codec: `_is_keyframe` returns None and `_has_keyframe` returns True so
|
|
# failures surface from the decoder rather than being swallowed as cold-start.
|
|
assert VideoFrameDecoder(codec="vp9")._has_keyframe([b"\x00"]) is True
|
|
|
|
|
|
def test_video_frame_decoder_derives_keyframe_path() -> None:
|
|
decoder = VideoFrameDecoder(codec="h264")
|
|
assert decoder.prior_keyframe_path("/camera:VideoStream:sample") == "/camera:VideoStream:is_keyframe"
|
|
assert (
|
|
decoder.prior_keyframe_path("/robot/cam_left:VideoStream:sample") == "/robot/cam_left:VideoStream:is_keyframe"
|
|
)
|
|
|
|
|
|
def test_video_frame_decoder_keyframe_path_no_separator() -> None:
|
|
# Defensive: a path with no `:` is non-canonical; return None rather than guessing.
|
|
assert VideoFrameDecoder(codec="h264").prior_keyframe_path("/just_an_entity") is None
|
|
|
|
|
|
def test_field_index_range_window_beats_anchor_and_heuristic() -> None:
|
|
field = Field(path="/camera:VideoStream:sample", decode=VideoFrameDecoder(codec="h264"), window=(-3, 5))
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=10)
|
|
# Anchor and heuristic must lose to the explicit window.
|
|
assert _field_index_range(100, field, decoder, prior_keyframe=42) == (97, 105)
|
|
|
|
|
|
def test_field_index_range_anchor_beats_heuristic_integer() -> None:
|
|
field = Field(path="/camera:VideoStream:sample", decode=VideoFrameDecoder(codec="h264"))
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=10)
|
|
assert _field_index_range(100, field, decoder, prior_keyframe=87) == (87, 100)
|
|
|
|
|
|
def test_field_index_range_anchor_beats_heuristic_timestamp() -> None:
|
|
field = Field(path="/camera:VideoStream:sample", decode=VideoFrameDecoder(codec="h264"))
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=30, fps_estimate=30.0)
|
|
target = np.datetime64(1_000_000_000, "ns")
|
|
result = _field_index_range(target, field, decoder, prior_keyframe=500_000_000)
|
|
assert result is not None
|
|
lo, hi = result
|
|
assert lo == np.datetime64(500_000_000, "ns")
|
|
assert hi == target
|
|
|
|
|
|
def test_field_index_range_falls_back_to_heuristic_when_anchor_missing() -> None:
|
|
# Simulates "no prior keyframe yet in this segment" — the prefetcher drops the
|
|
# entry and the field falls back to the decoder's heuristic context_range.
|
|
field = Field(path="/camera:VideoStream:sample", decode=VideoFrameDecoder(codec="h264"))
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=10)
|
|
assert _field_index_range(100, field, decoder, prior_keyframe=None) == (90, 100)
|
|
|
|
|
|
def test_field_index_range_default_kwarg_is_none() -> None:
|
|
# Existing call sites that don't pass `prior_keyframe` keep the same behavior.
|
|
field = Field(path="/camera:VideoStream:sample", decode=VideoFrameDecoder(codec="h264"))
|
|
decoder = VideoFrameDecoder(codec="h264", keyframe_interval=5)
|
|
assert _field_index_range(20, field, decoder) == (15, 20)
|
|
|
|
|
|
def test_prior_keyframe_none_or_empty_returns_none() -> None:
|
|
assert _prior_keyframe(None, 100) is None
|
|
assert _prior_keyframe(np.array([], dtype=np.int64), 100) is None
|
|
|
|
|
|
def test_prior_keyframe_target_before_first_returns_none() -> None:
|
|
assert _prior_keyframe(np.array([50, 100, 150], dtype=np.int64), 49) is None
|
|
|
|
|
|
def test_prior_keyframe_target_equals_keyframe_returns_keyframe() -> None:
|
|
assert _prior_keyframe(np.array([50, 100, 150], dtype=np.int64), 100) == 100
|
|
|
|
|
|
def test_prior_keyframe_target_between_returns_largest_leq() -> None:
|
|
kfs = np.array([50, 100, 150], dtype=np.int64)
|
|
assert _prior_keyframe(kfs, 99) == 50
|
|
assert _prior_keyframe(kfs, 149) == 100
|
|
|
|
|
|
def test_prior_keyframe_target_after_last_returns_last() -> None:
|
|
assert _prior_keyframe(np.array([50, 100, 150], dtype=np.int64), 9999) == 150
|