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/FrameContainer.py

@@ -3,116 +3,227 @@
 
 import bob.bio.base
 import numpy
-
 import logging
-logger = logging.getLogger("bob.bio.video")
 
-class FrameContainer:
-  """A class for reading, manipulating and saving video content.
-  """
-
-  def __init__(self, hdf5 = None, load_function = bob.bio.base.load):
-    self._frames = []
-    if hdf5 is not None:
-      self.load(hdf5, load_function)
-
-  def __len__(self):
-    return len(self._frames)
-
-  def __iter__(self):
-    """Generator that returns the 3-tuple (frame_id, data, quality) for each frame."""
-    # don't sort
-    for frame in self._frames:
-      yield frame
-
-  def __getitem__(self, i):
-    """Indexer (mostly used in tests)."""
-    return self._frames[i]
-
-  def add(self, frame_id, frame, quality = None):
-    """Adds the frame with the given id and the given quality."""
-    self._frames.append((str(frame_id), frame, quality))
-
-  def load(self, hdf5, load_function = bob.bio.base.load):
-    self._frames = []
-    if hdf5.has_group("FrameIndexes"):
-      hdf5.cd("FrameIndexes")
-      indices = sorted(int(i) for i in hdf5.keys(relative=True))
-      frame_ids = [hdf5[str(i)] for i in indices]
-      hdf5.cd("..")
-    else:
-      frame_ids = hdf5.sub_groups(relative=True, recursive=False)
-
-    # Read content (frames) from HDF5File
-    for path in frame_ids:
-      # extract frame_id
-      if path[:6] == 'Frame_':
-        frame_id = str(path[6:])
-        hdf5.cd(path)
-        # Read data
-        data = load_function(hdf5)
-
-        # read quality, if present
-        quality = hdf5.read("FrameQuality") if hdf5.has_key("FrameQuality") else None
-        self.add(frame_id, data, quality)
-        hdf5.cd("..")
-    if not len(self):
-      raise IOError("Could not load data as a Frame Container from file %s" % hdf5.filename)
-
-  def save(self, hdf5, save_function = bob.bio.base.save):
-    """ Save the content to the given HDF5 File.
-    The contained data will be written using the given save_function."""
-    if not len(self):
-      logger.warn("Saving empty FrameContainer '%s'", hdf5.filename)
-    frame_ids = []
-    for frame_id, data, quality in self:
-      hdf5.create_group("Frame_%s" % frame_id)
-      hdf5.cd("Frame_%s" % frame_id)
-      frame_ids.append("Frame_%s" % frame_id)
-      save_function(data, hdf5)
-      if quality is not None:
-        hdf5.set("FrameQuality", quality)
-      hdf5.cd("..")
-    # save the order of frames too so we can load them correctly later
-    hdf5.create_group("FrameIndexes")
-    hdf5.cd("FrameIndexes")
-    for i, v in enumerate(frame_ids):
-      hdf5[str(i)] = v
-    hdf5.cd("..")
-
-  def is_similar_to(self, other):
-    if len(self) != len(other): return False
-    for a,b in zip(self, other):
-      if a[0] != b[0]: return False
-      if abs(a[2] - b[2]) > 1e-8: return False
-      if not numpy.allclose(a[1], b[1]): return False
-    return True
-
-  def as_array(self):
-    """Returns the data of frames as a numpy array.
+logger = logging.getLogger(__name__)
+
+
+def select_frames(count, max_number_of_frames, selection_style, step_size):
+    """Returns indices of the frames to be selected given the parameters.
+
+    Different selection styles are supported:
+
+    * first : The first frames are selected
+    * spread : Frames are selected to be taken from the whole video with equal spaces in
+      between.
+    * 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 selected unconditionally.
+
+    Parameters
+    ----------
+    count : int
+        Total number of frames that are available
+    max_number_of_frames : int
+        The maximum number of frames to be selected. Ignored when selection_style is
+        "all".
+    selection_style : str
+        One of (``first``, ``spread``, ``step``, ``all``). See above.
+    step_size : int
+        Only useful when ``selection_style`` is ``step``.
 
     Returns
     -------
-    numpy.ndarray
-        The frames are returned as an array with the shape of (n_frames, ...)
-        like a video.
+    range
+        A range of frames to be selected.
+
+    Raises
+    ------
+    ValueError
+        If ``selection_style`` is not one of the supported ones.
     """
-    def _reader(frame):
-      # Each frame is assumed to be an image here. We make it a single frame
-      # video here by expanding its dimensions. This way it can be used with
-      # the vstack_features function.
-      return frame[1][None, ...]
-    return bob.bio.base.vstack_features(_reader, self._frames, same_size=True)
+    if selection_style == "first":
+        # get the first frames (limited by all frames)
+        indices = range(0, min(count, max_number_of_frames))
+    elif selection_style == "spread":
+        # get frames lineraly spread over all frames
+        indices = bob.bio.base.selected_indices(count, max_number_of_frames)
+    elif selection_style == "step":
+        indices = range(step_size // 2, count, step_size)[:max_number_of_frames]
+    elif selection_style == "all":
+        indices = range(0, count)
+    else:
+        raise ValueError(f"Invalid selection style: {selection_style}")
+
+    return indices
+
+
+class FrameContainer:
+    """A class for reading, manipulating and saving video content.
+    """
+
+    def __init__(self, hdf5=None, load_function=bob.bio.base.load, **kwargs):
+        super().__init__(**kwargs)
+        self._frames = []
+        if hdf5 is not None:
+            self.load(hdf5, load_function)
+
+    def __len__(self):
+        return len(self._frames)
+
+    def __iter__(self):
+        """Generator that returns the 3-tuple (frame_id, data, quality) for each frame."""
+        # don't sort
+        for frame in self._frames:
+            yield frame
+
+    def __getitem__(self, i):
+        """Indexer (mostly used in tests)."""
+        return self._frames[i]
+
+    def add(self, frame_id, frame, quality=None):
+        """Adds the frame with the given id and the given quality."""
+        self._frames.append((str(frame_id), frame, quality))
+
+    def load(
+        self,
+        hdf5,
+        load_function=bob.bio.base.load,
+        selection_style="all",
+        max_number_of_frames=20,
+        step_size=10,
+    ):
+        """Loads a previously saved FrameContainer into the current FrameContainer.
+
+        Parameters
+        ----------
+        hdf5 : :any:`bob.io.base.HDF5File`
+            An opened HDF5 file to load the data form
+        load_function : ``callable``, ``optional``
+            the function to be used on the hdf5 object to load each frame
+        selection_style : str, ``optional``
+            See :any:`select_frames`
+        max_number_of_frames : int, ``optional``
+            See :any:`select_frames`
+        step_size : int, ``optional``
+            See :any:`select_frames`
+
+        Returns
+        -------
+        object
+            returns itself.
+
+        Raises
+        ------
+        IOError
+            If no frames can be loaded from the hdf5 file.
+        ValueError
+            If the selection_style is all and you are trying to load an old format
+            FrameContainer.
+        """
+        self._frames = []
+        if hdf5.has_group("FrameIndexes"):
+            hdf5.cd("FrameIndexes")
+            indices = sorted(int(i) for i in hdf5.keys(relative=True))
+            indices = select_frames(
+                count=len(indices),
+                max_number_of_frames=max_number_of_frames,
+                selection_style=selection_style,
+                step_size=step_size,
+            )
+            frame_ids = [hdf5[str(i)] for i in indices]
+            hdf5.cd("..")
+        else:
+            if selection_style != "all":
+                raise ValueError(
+                    "selection_style must be all when loading FrameContainers with "
+                    "the old format. Try re-writing the FrameContainers again "
+                    "to avoid this."
+                )
+            frame_ids = hdf5.sub_groups(relative=True, recursive=False)
+
+        # Read content (frames) from HDF5File
+        for path in frame_ids:
+            # extract frame_id
+            if path[:6] == "Frame_":
+                frame_id = str(path[6:])
+                hdf5.cd(path)
+                # Read data
+                data = load_function(hdf5)
+
+                # read quality, if present
+                quality = hdf5["FrameQuality"] if "FrameQuality" in hdf5 else None
+                self.add(frame_id, data, quality)
+                hdf5.cd("..")
+        if not len(self):
+            raise IOError(
+                "Could not load data as a Frame Container from file %s" % hdf5.filename
+            )
+        return self
+
+    def save(self, hdf5, save_function=bob.bio.base.save):
+        """ Save the content to the given HDF5 File.
+        The contained data will be written using the given save_function."""
+        if not len(self):
+            logger.warn("Saving empty FrameContainer '%s'", hdf5.filename)
+        frame_ids = []
+        for frame_id, data, quality in self:
+            hdf5.create_group("Frame_%s" % frame_id)
+            hdf5.cd("Frame_%s" % frame_id)
+            frame_ids.append("Frame_%s" % frame_id)
+            save_function(data, hdf5)
+            if quality is not None:
+                hdf5.set("FrameQuality", quality)
+            hdf5.cd("..")
+        # save the order of frames too so we can load them correctly later
+        hdf5.create_group("FrameIndexes")
+        hdf5.cd("FrameIndexes")
+        for i, v in enumerate(frame_ids):
+            hdf5[str(i)] = v
+        hdf5.cd("..")
+
+    def is_similar_to(self, other):
+        if len(self) != len(other):
+            return False
+        for a, b in zip(self, other):
+            if a[0] != b[0]:
+                return False
+            if abs(a[2] - b[2]) > 1e-8:
+                return False
+            if not numpy.allclose(a[1], b[1]):
+                return False
+        return True
+
+    def as_array(self):
+        """Returns the data of frames as a numpy array.
+
+        Returns
+        -------
+        numpy.ndarray
+            The frames are returned as an array with the shape of (n_frames, ...)
+            like a video.
+        """
+
+        def _reader(frame):
+            # Each frame is assumed to be an image here. We make it a single frame
+            # video here by expanding its dimensions. This way it can be used with
+            # the vstack_features function.
+            return frame[1][None, ...]
+
+        return bob.bio.base.vstack_features(_reader, self._frames, same_size=True)
+
 
 def save_compressed(frame_container, filename, save_function, create_link=True):
-  hdf5 = bob.bio.base.open_compressed(filename, 'w')
-  frame_container.save(hdf5, save_function)
-  bob.bio.base.close_compressed(filename, hdf5, create_link=create_link)
-  del hdf5
+    hdf5 = bob.bio.base.open_compressed(filename, "w")
+    frame_container.save(hdf5, save_function)
+    bob.bio.base.close_compressed(filename, hdf5, create_link=create_link)
+    del hdf5
+
 
 def load_compressed(filename, load_function):
-  hdf5 = bob.bio.base.open_compressed(filename, 'r')
-  fc = FrameContainer(hdf5, load_function)
-  bob.bio.base.close_compressed(filename, hdf5)
-  del hdf5
-  return fc
+    hdf5 = bob.bio.base.open_compressed(filename, "r")
+    fc = FrameContainer(hdf5, load_function)
+    bob.bio.base.close_compressed(filename, hdf5)
+    del hdf5
+    return 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