wehub-resource-sync
36 lines
1010 B
Python
36 lines
1010 B
Python
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
|
|
|
|
def test_read_videos():
|
|
uri = "s3://anonymous@ray-example-data/basketball.mp4"
|
|
ds = ray.data.read_videos(uri, include_timestamps=True)
|
|
|
|
assert ds.count() == 333
|
|
assert ds.schema().names == ["frame", "frame_index", "frame_timestamp"]
|
|
|
|
frame_indices = ds.select_columns(["frame_index"]).take_all()
|
|
assert sorted(frame_indices, key=lambda item: item["frame_index"]) == [
|
|
{"frame_index": i} for i in range(333)
|
|
]
|
|
|
|
frame_timestamps = ds.select_columns(["frame_timestamp"]).take_all()
|
|
for t in frame_timestamps:
|
|
assert isinstance(t["frame_timestamp"], np.ndarray)
|
|
assert t["frame_timestamp"].shape[0] == 2
|
|
|
|
frame_type, frame_index_type, _ = ds.schema().types
|
|
|
|
assert frame_type.shape == (720, 1280, 3)
|
|
assert frame_type.value_type == pa.uint8()
|
|
assert frame_index_type == pa.int64()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|