query.py 4.93 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
6
from bob.db.base.annotations import read_annotation_file
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
7
from . import OULUNPU_FRAME_SHAPE
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
8
9


10
class File(VideoPadFile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
11
    """The file objects of the OULU-NPU dataset."""
12
13
14
15
16
17
18
19
20
21

    @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).
        """
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.
        """
34
        vfilename = self.make_path(directory=self.original_directory, extension=".avi")
35
36
37
38
39
40
41
42
43
44
45
46
        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
47

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
48
49
    @property
    def annotations(self):
50
51
52
53
54
        """Reads the annotations.

        If the file object has an attribute of annotation_directory, it will read
        annotations from there instead of loading annotations that are shipped with the
        database.
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
55
56
57
58
59
60
61

        Returns
        -------
        dict
            The annotations as a dictionary, e.g.:
            ``{'0': {'reye':(re_y,re_x), 'leye':(le_y,le_x)}, ...}``
        """
62
63
64
65
66
67
68
69
        if (
            hasattr(self, "annotation_directory")
            and self.annotation_directory is not None
        ):
            path = self.make_path(self.annotation_directory, extension=".json")
            return read_annotation_file(path, annotation_type="json")

        path = self.make_path(directory=self.original_directory, extension=".txt")
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
70
71
72
73
74
75
        annotations = {}
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
76
77
78
                num_frame, x_eye_right, y_eye_right, x_eye_left, y_eye_left = line.split(
                    ","
                )
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
79
                annotations[num_frame] = {
80
81
                    "reye": (int(y_eye_right), int(x_eye_right)),
                    "leye": (int(y_eye_left), int(x_eye_left)),
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
82
83
84
                }
        return annotations

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
85
86
87
88

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

89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    def __init__(
        self,
        original_directory=rc["bob.db.oulunpu.directory"],
        name="oulunpu",
        pad_file_class=None,
        original_extension=".avi",
        annotation_directory=None,
        **kwargs
    ):
        """Summary

        Parameters
        ----------
        original_directory : TYPE, optional
            Description
        name : str, optional
            Description
        pad_file_class : None, optional
            Description
        original_extension : str, optional
            Description
        annotation_directory : str
            If provided, the annotations will be read from this directory
            instead of the default annotations that are provided.
        **kwargs
            Description
        """
116
117
        if pad_file_class is None:
            pad_file_class = File
118
        filelists_directory = resource_filename(__name__, "lists")
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
119
        super(Database, self).__init__(
120
121
            filelists_directory=filelists_directory,
            name=name,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
122
            original_directory=original_directory,
123
            pad_file_class=pad_file_class,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
124
            original_extension=original_extension,
125
            annotation_directory=annotation_directory,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
126
            training_depends_on_protocol=True,
127
128
129
130
131
132
133
134
135
136
137
138
            **kwargs
        )

    def objects(
        self,
        groups=None,
        protocol=None,
        purposes=None,
        model_ids=None,
        classes=None,
        **kwargs
    ):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
139
        """Returns the requested samples."""
140
        files = super(Database, self).objects(
141
142
143
144
145
146
147
            groups=groups,
            protocol=protocol,
            purposes=purposes,
            model_ids=model_ids,
            classes=classes,
            **kwargs
        )
148
149
        for f in files:
            f.original_directory = self.original_directory
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
150
            f.original_extension = self.original_extension
151
152
            f.annotation_directory = self.annotation_directory

153
154
        return files

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
155
    def frames(self, padfile):
156
        return padfile.frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
157

158
159
    def number_of_frames(self, padfile):
        return padfile.number_of_frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
160

161
162
163
    @property
    def frame_shape(self):
        return OULUNPU_FRAME_SHAPE
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
164

165
    def annotations(self, padfile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
166
        return padfile.annotations