test_utils.py 2.71 KB
Newer Older
1
import tempfile
2
import time
3 4 5

import bob.bio.video
import numpy as np
6
from bob.bio.base.test.utils import is_library_available
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
7
from bob.io.base.test_utils import datafile
8
from bob.io.video import reader
9 10 11 12

regenerate_refs = False


Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
13 14
def test_video_as_array():
    path = datafile("testvideo.avi", "bob.bio.video.test")
15

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
16 17 18
    video = bob.bio.video.VideoAsArray(path, selection_style="all")
    assert len(video) == 83, len(video)
    assert video.indices == range(83), video.indices
19
    assert video.shape == (83, 3, 480, 640), video.shape
20

21 22 23 24 25 26 27 28
    video_slice = video[1:2, 1:-1, 1:-1, 1:-1]
    assert video_slice.shape == (1, 1, 478, 638), video_slice.shape
    # test the slice against the video loaded by bob.io.video directly
    video = reader(path)[1]
    video = video[1:-1, 1:-1, 1:-1]
    video = video[None, ...]
    np.testing.assert_allclose(video, video_slice)

29
    video = bob.bio.video.VideoAsArray(path, max_number_of_frames=3)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
30 31
    assert len(video) == 3, len(video)
    assert video.indices == [13, 41, 69], video.indices
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54
    assert video.shape == (3, 3, 480, 640), video.shape


@is_library_available("dask")
def test_video_as_array_vs_dask():
    import dask

    path = datafile("testvideo.avi", "bob.bio.video.test")
    start = time.time()
    video = bob.bio.video.VideoAsArray(path, selection_style="all")
    video = dask.array.from_array(video, (20, 1, 480, 640))
    video = video.compute()
    load_time = time.time() - start

    start = time.time()
    reference = reader(path).load()
    load_time2 = time.time() - start
    # Here, we're also chunking each frame, but normally we would only chunk the first axis.
    print(
        f"FYI: It took {load_time:.2f} s to load the video with dask and {load_time2:.2f} s "
        "to load directly. The slower loading with dask is expected."
    )
    np.testing.assert_allclose(reference, video)
55 56


Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
57 58
def test_video_like_container():
    path = datafile("testvideo.avi", "bob.bio.video.test")
59

60 61 62
    video = bob.bio.video.VideoAsArray(
        path, selection_style="spread", max_number_of_frames=3
    )
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
63
    container = bob.bio.video.VideoLikeContainer(video, video.indices)
64

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
65
    container_path = datafile("video_like.hdf5", "bob.bio.video.test")
66

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
67 68
    if regenerate_refs:
        container.save(container_path)
69

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
70
    loaded_container = bob.bio.video.VideoLikeContainer.load(container_path)
71
    assert container == loaded_container
72

73 74 75 76 77 78 79 80 81
    # test saving and loading None arrays
    with tempfile.NamedTemporaryFile(suffix=".pkl") as f:
        data = [None] * 10 + [1]
        indices = range(11)
        frame_container = bob.bio.video.VideoLikeContainer(data, indices)
        frame_container.save(f.name)

        loaded = bob.bio.video.VideoLikeContainer.load(f.name)
        assert loaded == frame_container