* Add a transform function for VideoAsArray

* Don't load VideoLikeContainer checkpoints into memory
* More documentation

bob/bio/video/utils.py

@@ -7,7 +7,9 @@ import numpy as np
 logger = logging.getLogger(__name__)
 
 
-def select_frames(count, max_number_of_frames=None, selection_style=None, step_size=None):
+def select_frames(
+    count, max_number_of_frames=None, selection_style=None, step_size=None
+):
     """Returns indices of the frames to be selected given the parameters.
 
     Different selection styles are supported:
@@ -67,14 +69,35 @@ def select_frames(count, max_number_of_frames=None, selection_style=None, step_s
 
 
 class VideoAsArray:
+    """A memory efficient class to load only select video frames
+    It also supports efficient conversion to dask arrays.
+    """
     def __init__(
         self,
         path,
         selection_style=None,
         max_number_of_frames=None,
         step_size=None,
+        transform=None,
         **kwargs,
     ):
+        """init
+
+        Parameters
+        ----------
+        path : str
+            Path to the video file
+        selection_style : str, optional
+            See :any:`select_frames`, by default None
+        max_number_of_frames : int, optional
+            See :any:`select_frames`, by default None
+        step_size : int, optional
+            See :any:`select_frames`, by default None
+        transform : callable, optional
+            A function that transforms the loaded video. This function should
+            not change the video shape or its dtype. For example, you may flip
+            the frames horizontally using this function, by default None
+        """
         super().__init__(**kwargs)
         self.path = path
         self.reader = bob.io.video.reader(self.path)
@@ -89,6 +112,10 @@ class VideoAsArray:
         )
         self.indices = indices
         self.shape = (len(indices),) + shape[1:]
+        if transform is None:
+            def transform(x):
+                return x
+        self.transform = transform
 
     def __getstate__(self):
         d = self.__dict__.copy()
@@ -106,7 +133,7 @@ class VideoAsArray:
         # logger.debug("Getting frame %s from %s", index, self.path)
         if isinstance(index, int):
             idx = self.indices[index]
-            return self.reader[idx]
+            return self.transform([self.reader[idx]])[0]
 
         if not (isinstance(index, tuple) and len(index) == self.ndim):
             raise NotImplementedError(f"Indxing like {index} is not supported yet!")
@@ -115,7 +142,7 @@ class VideoAsArray:
             return np.array([], dtype=self.dtype)
 
         if self.selection_style == "all":
-            return np.asarray(self.reader.load())[index]
+            return self.transform(np.asarray(self.reader.load())[index])
 
         idx = self.indices[index[0]]
         video = []
@@ -127,7 +154,7 @@ class VideoAsArray:
                 break
 
         index = (slice(len(video)),) + index[1:]
-        return np.asarray(video)[index]
+        return self.transform(np.asarray(video)[index])
 
     def __repr__(self):
         return f"{self.reader!r} {self.dtype!r} {self.ndim!r} {self.shape!r} {self.indices!r}"
@@ -156,8 +183,10 @@ class VideoLikeContainer:
 
     @classmethod
     def load(cls, file):
-        with h5py.File(file, mode="r") as f:
-            data = np.array(f["data"])
-            indices = np.array(f["indices"])
+        # weak closing of the hdf5 file so we don't load all the data into
+        # memory https://docs.h5py.org/en/stable/high/file.html#closing-files
+        f = h5py.File(file, mode="r")
+        data = f["data"]
+        indices = f["indices"]
         self = cls(data=data, indices=indices)
         return self