File size: 877 Bytes
1e2bb2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | import numpy as np
import torch
from PIL import Image
from stimulus_synthesis.media.normalize import videos_to_b_t_c_h_w
def test_videos_to_b_t_c_h_w_from_pil_frames():
frames = [
Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)),
Image.fromarray(np.full((8, 8, 3), 255, dtype=np.uint8)),
]
video = videos_to_b_t_c_h_w([frames], size=4, num_frames=3)
assert video.shape == (1, 3, 3, 4, 4)
assert torch.all(video >= 0)
assert torch.all(video <= 1)
def test_videos_to_b_t_c_h_w_from_tensor_thwc():
tensor = torch.zeros(2, 8, 8, 3)
video = videos_to_b_t_c_h_w([tensor], size=4)
assert video.shape == (1, 2, 3, 4, 4)
def test_videos_to_b_t_c_h_w_from_numpy_bthwc():
array = np.zeros((1, 2, 8, 8, 3), dtype=np.float32)
video = videos_to_b_t_c_h_w([array], size=4)
assert video.shape == (1, 2, 3, 4, 4)
|