Improve FrameContainer and FrameSelector to be more memory efficient

bob/bio/video/test/test_utils.py

@@ -41,6 +41,12 @@ def test_frame_container():
     # test as_array method
     assert numpy.allclose(read.as_array(), test_data)
 
+    # check loading only a part of the hdf5
+    with bob.io.base.HDF5File(filename) as f:
+      # load only a subset of the FrameContainer
+      fc = bob.bio.video.FrameContainer().load(f, selection_style="spread", max_number_of_frames=10)
+    assert len(fc) == 10, len(fc)
+
   finally:
     if os.path.exists(filename):
       os.remove(filename)
@@ -113,3 +119,9 @@ def test_frame_selector():
   assert frames[1][0] == '6'
   assert numpy.allclose(frames[1][1], video_data[6])
   assert frames[1][2] is None
+
+  # test bob.io.video.reader support
+  path = bob.io.base.test_utils.datafile("testvideo.avi", __name__)
+  fs = bob.bio.video.FrameSelector(selection_style="spread", max_number_of_frames=20)
+  fc = fs(bob.io.video.reader(path))  # only loads 20 frames into memory
+  assert len(fc) == 20, len(fc)

bob/bio/video/utils/FrameSelector.py

@@ -10,84 +10,93 @@ import os
 import six
 
 import logging
-logger = logging.getLogger("bob.bio.video")
 
-from .FrameContainer import FrameContainer
+logger = logging.getLogger(__name__)
+
+from .FrameContainer import FrameContainer, select_frames
+
 
 class FrameSelector:
-  """A class for selecting frames from videos.
-  In total, up to ``max_number_of_frames`` is selected (unless selection style is ``all``
-
-  Different selection styles are supported:
-
-  * first : The first frames are selected
-  * spread : Frames are selected to be taken from the whole video
-  * step : Frames are selected every ``step_size`` indices, starting at ``step_size/2`` **Think twice if you want to have that when giving FrameContainer data!**
-  * all : All frames are stored unconditionally
-  * quality (only valid for FrameContainer data) : Select the frames based on the highest internally stored quality value
-  """
-
-  def __init__(self,
-      max_number_of_frames = 20,
-      selection_style = "spread",
-      step_size = 10
-  ):
-    if selection_style not in ('first', 'spread', 'step', 'all'):
-      raise ValueError("Unknown selection style '%s', choose one of ('first', 'spread', 'step', 'all')" % selection_style)
-    self.selection = selection_style
-    self.max_frames = max_number_of_frames
-    self.step = step_size
-
-  def __call__(self, data, load_function = bob.io.base.load):
-    """Selects frames and returns them in a FrameContainer.
-    Different ``data`` parameters are accepted:
-
-    * :py:class:`FrameContainer` : frames are selected from the given frame container
-    * ``str`` : A video file to read and select frames from
-    * ``[str]`` : A list of image names to select from
-    * ``numpy.array`` (3D or 4D): A video to select frames from
-
-    When giving ``str`` or ``[str]`` data, the given ``load_function`` is used to read the data from file.
+    """A class for selecting frames from videos.
+    In total, up to ``max_number_of_frames`` is selected (unless selection style is ``all``
+
+    Different selection styles are supported:
+
+    * first : The first frames are selected
+    * spread : Frames are selected to be taken from the whole video
+    * step : Frames are selected every ``step_size`` indices, starting at ``step_size/2`` **Think twice if you want to have that when giving FrameContainer data!**
+    * all : All frames are stored unconditionally
+    * quality (only valid for FrameContainer data) : Select the frames based on the highest internally stored quality value
     """
-    # if given a string, first load the video
-    if isinstance(data, six.string_types):
-      logger.debug("Loading video file '%s'", data)
-      data = load_function(data)
-
-    # first, get the indices
-    count = len(data)
-    if self.selection == 'first':
-      # get the first frames (limited by all frames)
-      indices = range(0, min(count, self.max_frames))
-    elif self.selection == 'spread':
-      # get frames lineraly spread over all frames
-      indices = bob.bio.base.selected_indices(count, self.max_frames)
-    elif self.selection == 'step':
-      indices = range(self.step//2, count, self.step)[:self.max_frames]
-    elif self.selection == 'all':
-      indices = range(0, count)
-
-    # now, iterate through the data
-    fc = FrameContainer()
-    if isinstance(data, FrameContainer):
-      indices = set(indices)
-      # frame container data, just copy
-      for i, frame in enumerate(data):
-        if i in indices:
-          fc.add(*frame)
-    elif isinstance(data, numpy.ndarray):
-      # select video frames
-      for i in indices:
-        fc.add(i, data[i])
-    elif isinstance(data, list):
-      for i in indices:
-        # load image
-        image = load_function(data[i])
-        # save image name as well
-        fc.add(os.path.basename(data[i]), image)
-
-    return fc
-
-  def __str__(self):
-    """Writes the parameters of the FrameSelector as a string."""
-    return "FrameSelector(max_number_of_frames=%d, selection_style='%s', step_size=%d)" % (self.max_frames, self.selection, self.step)
+
+    def __init__(self, max_number_of_frames=20, selection_style="spread", step_size=10):
+        if selection_style not in ("first", "spread", "step", "all"):
+            raise ValueError(
+                "Unknown selection style '%s', choose one of ('first', 'spread', 'step', 'all')"
+                % selection_style
+            )
+        self.selection = selection_style
+        self.max_frames = max_number_of_frames
+        self.step = step_size
+
+    def __call__(self, data, load_function=bob.io.base.load):
+        """Selects frames and returns them in a FrameContainer.
+        Different ``data`` parameters are accepted:
+
+        * :py:class:`FrameContainer` : frames are selected from the given frame container
+        * ``str`` : A video file to read and select frames from
+        * ``[str]`` : A list of image names to select from
+        * ``numpy.array`` (3D or 4D): A video to select frames from
+        * ``bob.io.video.reader`` : An instance of bob.io.video.reader.
+
+        When giving ``str`` or ``[str]`` data, the given ``load_function`` is used to read the data from file.
+        """
+        # if given a string, first load the video
+        if isinstance(data, six.string_types):
+            logger.debug("Loading video file '%s'", data)
+            data = load_function(data)
+
+        # first, get the indices
+        if isinstance(data, bob.io.video.reader):
+            count = data.number_of_frames
+        else:
+            count = len(data)
+
+        indices = select_frames(
+            count=count,
+            max_number_of_frames=self.max_frames,
+            selection_style=self.selection,
+            step_size=self.step,
+        )
+
+        # now, iterate through the data
+        fc = FrameContainer()
+        if isinstance(data, FrameContainer):
+            indices = set(indices)
+            # frame container data, just copy
+            for i, frame in enumerate(data):
+                if i in indices:
+                    fc.add(*frame)
+        elif isinstance(data, bob.io.video.reader):
+            for i, frame in enumerate(data):
+                if i in indices:
+                    fc.add(i, frame)
+        elif isinstance(data, numpy.ndarray):
+            # select video frames
+            for i in indices:
+                fc.add(i, data[i])
+        elif isinstance(data, list):
+            for i in indices:
+                # load image
+                image = load_function(data[i])
+                # save image name as well
+                fc.add(os.path.basename(data[i]), image)
+
+        return fc
+
+    def __str__(self):
+        """Writes the parameters of the FrameSelector as a string."""
+        return (
+            "FrameSelector(max_number_of_frames=%d, selection_style='%s', step_size=%d)"
+            % (self.max_frames, self.selection, self.step)
+        )

bob/bio/video/utils/__init__.py

-from .FrameContainer import FrameContainer, load_compressed, save_compressed
+from .FrameContainer import FrameContainer, load_compressed, save_compressed, select_frames
 from .FrameSelector import FrameSelector