query.py 3.5 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1
from pkg_resources import resource_filename
2 3
from bob.pad.base.database import FileListPadDatabase
from bob.pad.face.database import VideoPadFile
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
4
from bob.extension import rc
5
from bob.io.video import reader
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
6
from . import OULUNPU_FRAME_SHAPE
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
7 8


9
class File(VideoPadFile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
10
    """The file objects of the OULU-NPU dataset."""
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45

    @property
    def frames(self):
        """Yields the frames of the biofile one by one.

        Yields
        ------
        :any:`numpy.array`
            A frame of the video. The size is (3, 1920, 1080).
        """
        vfilename = self.make_path(directory=self.original_directory)
        return iter(reader(vfilename))

    @property
    def number_of_frames(self):
        """Returns the number of frames in a video file.

        Returns
        -------
        int
            The number of frames.
        """
        vfilename = self.make_path(directory=self.original_directory)
        return reader(vfilename).number_of_frames

    @property
    def frame_shape(self):
        """Returns the size of each frame in this database.

        Returns
        -------
        (int, int, int)
            The (#Channels, Height, Width) which is :any:`OULUNPU_FRAME_SHAPE`.
        """
        return OULUNPU_FRAME_SHAPE
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
46 47 48 49 50


class Database(FileListPadDatabase):
    """The database interface for the OULU-NPU dataset."""

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
51
    def __init__(self, original_directory=rc['bob.db.oulunpu.directory'],
52
                 name='oulunpu', pad_file_class=None,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
53
                 original_extension=".avi", **kwargs):
54 55
        if pad_file_class is None:
            pad_file_class = File
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
56 57 58 59
        filelists_directory = resource_filename(__name__, 'lists')
        super(Database, self).__init__(
            filelists_directory=filelists_directory, name=name,
            original_directory=original_directory,
60
            pad_file_class=pad_file_class,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
61
            original_extension=original_extension,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
62
            **kwargs)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
63

64 65 66 67 68 69 70 71 72
    def objects(self, groups=None, protocol=None, purposes=None,
                model_ids=None, classes=None, **kwargs):
        files = super(Database, self).objects(
            groups=groups, protocol=protocol, purposes=purposes,
            model_ids=model_ids, classes=classes, **kwargs)
        for f in files:
            f.original_directory = self.original_directory
        return files

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
73
    def frames(self, padfile):
74
        return padfile.frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
75

76 77
    def number_of_frames(self, padfile):
        return padfile.number_of_frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
78

79 80 81
    @property
    def frame_shape(self):
        return OULUNPU_FRAME_SHAPE
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
82

83 84
    def annotations(self, padfile):
        """Reads the annotations for the given padfile.
85 86 87 88

        Parameters
        ----------
        padfile : :any:`File`
89
            The file object for which the annotations should be read.
90 91 92

        Returns
        -------
93 94 95
        dict
            The annotations as a dictionary, e.g.:
            ``{'0': {'reye':(re_y,re_x), 'leye':(le_y,le_x)}, ...}``
96
        """
97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
        path = padfile.make_path(
            directory=self.original_directory, extension='.txt')
        annotations = {}
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                num_frame, x_eye_left, y_eye_left, x_eye_right, y_eye_right = \
                    line.split(',')
                annotations[num_frame] = {
                    'reye': (int(y_eye_right), int(x_eye_right)),
                    'leye': (int(y_eye_left), int(x_eye_left)),
                }
        return annotations