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

bob/bio/video/test/test_utils.py

+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)
+    )

bob/bio/video/utils.py

@@ -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 = []
-        for i, frame in enumerate(self.reader):
-            if i not in idx:
-                continue
-            video.append(frame)
-            if i == idx[-1]:
-                break
-
-        index = (slice(len(video)),) + index[1:]
-        return self.transform(np.asarray(video)[index])
+        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 real_frame_numbers:
+                    continue
+                # 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
+
+        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}"

conda/meta.yaml

@@ -53,6 +53,7 @@ test:
     - bob.db.atnt
     - bob.io.image
     - bob.db.youtube
+    - dask
 
 about:
   home: https://www.idiap.ch/software/bob/

test-requirements.txt

 bob.io.image
 bob.db.atnt
 bob.db.youtube
+dask