Commit 997d95c2 authored by Amir MOHAMMADI's avatar Amir MOHAMMADI

Improve FrameContainer and FrameSelector to be more memory efficient

parent fe242495
......@@ -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)
This diff is collapsed.
......@@ -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)
)
from .FrameContainer import FrameContainer, load_compressed, save_compressed
from .FrameContainer import FrameContainer, load_compressed, save_compressed, select_frames
from .FrameSelector import FrameSelector
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