Wrapper.py 6.88 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1 2 3
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :

4 5 6 7 8 9
import bob.bio.base
import bob.io.base

from .. import utils


10 11
class Wrapper(bob.bio.base.preprocessor.Preprocessor):
    """Wrapper class to run image preprocessing algorithms on video data.
12

13 14
    This class provides functionality to read original video data from several databases.
    So far, the video content from :py:class:`bob.db.mobio` and the image list content from :py:class:`bob.db.youtube` are supported.
15

16 17 18
    Furthermore, frames are extracted from these video data, and a ``preprocessor`` algorithm is applied on all selected frames.
    The preprocessor can either be provided as a registered resource, i.e., one of :ref:`bob.bio.face.preprocessors`, or an instance of a preprocessing class.
    Since most of the databases do not provide annotations for all frames of the videos, commonly the preprocessor needs to apply face detection.
19

20 21
    The ``frame_selector`` can be chosen to select some frames from the video.
    By default, a few frames spread over the whole video sequence are selected.
22

23 24 25
    The ``quality_function`` is used to assess the quality of the frame.
    If no ``quality_function`` is given, the quality is based on the face detector, or simply left as ``None``.
    So far, the quality of the frames are not used, but it is foreseen to select frames based on quality.
26

27
    **Parameters:**
28

29 30
    preprocessor : str or :py:class:`bob.bio.base.preprocessor.Preprocessor` instance
      The preprocessor to be used to preprocess the frames.
31

32 33
    frame_selector : :py:class:`bob.bio.video.FrameSelector`
      A frame selector class to define, which frames of the video to use.
34

35 36 37 38
    quality_function : function or ``None``
      A function assessing the quality of the preprocessed image.
      If ``None``, no quality assessment is performed.
      If the preprocessor contains a ``quality`` attribute, this is taken instead.
39

40 41 42 43 44
    compressed_io : bool
      Use compression to write the resulting preprocessed HDF5 files.
      This is experimental and might cause trouble.
      Use this flag with care.
    """
45

46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
    def __init__(self,
                 preprocessor='landmark-detect',
                 frame_selector=utils.FrameSelector(),
                 quality_function=None,
                 compressed_io=False,
                 read_original_data=lambda biofile, directory, extension: biofile.load(directory, extension)
                 ):

        # load preprocessor configuration
        if isinstance(preprocessor, str):
            self.preprocessor = bob.bio.base.load_resource(preprocessor, "preprocessor")
        elif isinstance(preprocessor, bob.bio.base.preprocessor.Preprocessor):
            self.preprocessor = preprocessor
        else:
            raise ValueError("The given preprocessor could not be interpreted")
61

62 63 64 65 66 67 68
        bob.bio.base.preprocessor.Preprocessor.__init__(
            self,
            preprocessor=preprocessor,
            frame_selector=frame_selector,
            compressed_io=compressed_io,
            read_original_data=read_original_data
        )
69

70 71 72
        self.frame_selector = frame_selector
        self.quality_function = quality_function
        self.compressed_io = compressed_io
73

74 75 76
    def _check_data(self, frames):
        """Checks if the given video is in the desired format."""
        assert isinstance(frames, utils.FrameContainer)
77

78 79
    def __call__(self, frames, annotations=None):
        """__call__(frames, annotations = None) -> preprocessed
80

81
        Preprocesses the given frames using the desired ``preprocessor``.
82

83
        Faces are extracted for all frames in the given frame container, using the ``preprocessor`` specified in the constructor.
84

85 86 87
        If given, the annotations need to be in a dictionary.
        The key is either the frame number (for video data) or the image name (for image list data).
        The value is another dictionary, building the relation between facial landmark names and their location, e.g. ``{'leye' : (le_y, le_x), 'reye' : (re_y, re_x)}``
88

89 90
        The annotations for the according frames, if present, are passed to the preprocessor.
        Please assure that your database interface provides the annotations in the desired format.
91

92
        **Parameters:**
93

94 95
        frames : :py:class:`bob.bio.video.FrameContainer`
          The pre-selected frames, as returned by :py:meth:`read_original_data`.
96

97 98
        annotations : dict or ``None``
          The annotations for the frames, if any.
99

100
        **Returns:**
101

102 103 104 105 106
        preprocessed : :py:class:`bob.bio.video.FrameContainer`
          A frame container that contains the preprocessed frames.
        """
        if not isinstance(frames, utils.FrameContainer):
            frames = self.frame_selector(frames)
107

108 109
        annots = None
        fc = utils.FrameContainer()
110

111 112 113
        for index, frame, _ in frames:
            # if annotations are given, we take them
            if annotations is not None: annots = annotations[index]
114

115 116 117 118 119 120 121 122 123 124 125 126
            # preprocess image (by default: detect a face)
            preprocessed = self.preprocessor(frame, annots)
            if preprocessed is not None:
                # compute the quality of the detection
                if self.quality_function is not None:
                    quality = self.quality_function(preprocessed)
                elif hasattr(self.preprocessor, 'quality'):
                    quality = self.preprocessor.quality
                else:
                    quality = None
                # add image to frame container
                fc.add(index, preprocessed, quality)
127

128
        return fc
129 130


131 132
    def read_data(self, filename):
        """read_data(filename) -> frames
133

134 135
        Reads the preprocessed data from file and returns them in a frame container.
        The preprocessors ``read_data`` function is used to read the data for each frame.
136

137
        **Parameters:**
138

139 140
        filename : str
          The name of the preprocessed data file.
141

142
        **Returns:**
143

144 145 146 147 148 149 150
        frames : :py:class:`bob.bio.video.FrameContainer`
          The read frames, stored in a frame container.
        """
        if self.compressed_io:
            return utils.load_compressed(filename, self.preprocessor.read_data)
        else:
            return utils.FrameContainer(bob.io.base.HDF5File(filename), self.preprocessor.read_data)
151

152 153
    def write_data(self, frames, filename):
        """Writes the preprocessed data to file.
154

155
        The preprocessors ``write_data`` function is used to write the data for each frame.
156

157
        **Parameters:**
158

159 160
        frames : :py:class:`bob.bio.video.FrameContainer`
          The preprocessed frames, as returned by the :py:meth:`__call__` function.
161

162 163 164 165
        filename : str
          The name of the preprocessed data file to write.
        """
        self._check_data(frames)
166

167 168 169 170
        if self.compressed_io:
            return utils.save_compressed(frames, filename, self.preprocessor.write_data)
        else:
            frames.save(bob.io.base.HDF5File(filename, 'w'), self.preprocessor.write_data)