Commit 86eef97a authored by Amir MOHAMMADI's avatar Amir MOHAMMADI
Browse files

[VideoAsArray] Simplify the code for __getitem__ and add more tests

parent a63dc69f
Pipeline #46058 passed with stage
in 4 minutes and 37 seconds
import time
from bob.bio.base.test.utils import is_library_available
from bob.io.base.test_utils import datafile
from bob.io.video import reader
import numpy as np
import bob.bio.video
......@@ -11,16 +14,52 @@ def test_video_as_array():
video = bob.bio.video.VideoAsArray(path, selection_style="all")
assert len(video) == 83, len(video)
assert video.indices == range(83), video.indices
assert video.shape == (83, 3, 480, 640), video.shape
video = bob.bio.video.VideoAsArray(path, selection_style="spread", max_number_of_frames=3)
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)
video = bob.bio.video.VideoAsArray(
path, max_number_of_frames=3
)
assert len(video) == 3, len(video)
assert video.indices == [13, 41, 69], video.indices
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)
def test_video_like_container():
path = datafile("testvideo.avi", "bob.bio.video.test")
video = bob.bio.video.VideoAsArray(path, selection_style="spread", max_number_of_frames=3)
video = bob.bio.video.VideoAsArray(
path, selection_style="spread", max_number_of_frames=3
)
container = bob.bio.video.VideoLikeContainer(video, video.indices)
container_path = datafile("video_like.hdf5", "bob.bio.video.test")
......@@ -31,4 +70,6 @@ def test_video_like_container():
loaded_container = bob.bio.video.VideoLikeContainer.load(container_path)
np.testing.assert_equal(np.array(container.data), np.array(loaded_container.data))
np.testing.assert_equal(np.array(container.indices), np.array(loaded_container.indices))
np.testing.assert_equal(
np.array(container.indices), np.array(loaded_container.indices)
)
......@@ -135,30 +135,61 @@ class VideoAsArray:
def __getitem__(self, index):
# logger.debug("Getting frame %s from %s", index, self.path)
# In this method, someone is requesting indices thinking this video has
# the shape of self.shape but self.shape is determined through
# select_frames parameters. What we want to do here is to translate
# ``index`` to real indices of the video file given that we want to load
# only the selected frames. List of the selected frames are stored in
# self.indices
# If only one frame is requested, first translate the index to the real
# frame number in the video file and load that
if isinstance(index, int):
idx = self.indices[index]
return self.transform([self.reader[idx]])[0]
if not (isinstance(index, tuple) and len(index) == self.ndim):
if not (
isinstance(index, tuple)
and len(index) == self.ndim
and all(isinstance(idx, slice) for idx in index)
):
raise NotImplementedError(f"Indxing like {index} is not supported yet!")
# dask.array.from_array sometimes requests empty arrays
if all(i == slice(0, 0) for i in index):
return np.array([], dtype=self.dtype)
if self.selection_style == "all":
return self.transform(np.asarray(self.reader.load())[index])
idx = self.indices[index[0]]
video = []
def _frames_generator():
# read the frames one by one and yield them
real_frame_numbers = self.indices[index[0]]
for i, frame in enumerate(self.reader):
if i not in idx:
if i not in real_frame_numbers:
continue
video.append(frame)
if i == idx[-1]:
# make sure arrays are loaded in C order because we reshape them
# by C order later. Also, index into the frames here
frame = np.ascontiguousarray(frame)[index[1:]]
# return a tuple of flat array to match what is expected by
# field_dtype
yield (frame.ravel(),)
if i == real_frame_numbers[-1]:
break
index = (slice(len(video)),) + index[1:]
return self.transform(np.asarray(video)[index])
iterable = _frames_generator()
# compute the final shape given self.shape and index
# see https://stackoverflow.com/a/36188683/1286165
shape = [len(range(*idx.indices(dim))) for idx, dim in zip(index, self.shape)]
# field_dtype contains information about dtype and shape of each frame
# numpy black magic: https://stackoverflow.com/a/12473478/1286165 allows
# us to yield frame by frame in _frames_generator which greatly speeds
# up loading the video
field_dtype = [("", (self.dtype, (np.prod(shape[1:]),)))]
total_number_of_frames = shape[0]
video = np.fromiter(iterable, field_dtype, total_number_of_frames)
# view the array as self.dtype to remove the field_dtype
video = np.reshape(video.view(self.dtype), shape, order="C")
return self.transform(video)
def __repr__(self):
return f"{self.reader!r} {self.dtype!r} {self.ndim!r} {self.shape!r} {self.indices!r}"
......
......@@ -53,6 +53,7 @@ test:
- bob.db.atnt
- bob.io.image
- bob.db.youtube
- dask
about:
home: https://www.idiap.ch/software/bob/
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment