query.py 3.56 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

    @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).
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
21 22
        vfilename = self.make_path(
            directory=self.original_directory, extension='.avi')
23 24 25 26 27 28 29 30 31 32 33
        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.
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
34 35
        vfilename = self.make_path(
            directory=self.original_directory, extension='.avi')
36 37 38 39 40 41 42 43 44 45 46 47
        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
48 49 50 51 52


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

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

66 67 68 69 70 71 72 73 74
    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
75
    def frames(self, padfile):
76
        return padfile.frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
77

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

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

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

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

        Returns
        -------
95 96 97
        dict
            The annotations as a dictionary, e.g.:
            ``{'0': {'reye':(re_y,re_x), 'leye':(le_y,le_x)}, ...}``
98
        """
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
        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