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


class File(VideoPadFile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
11
    """The file objects of the OULU-NPU dataset."""
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
12

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
13
14
15
    @property
    def frames(self):
        """Yields the frames of the biofile one by one.
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
16
17
18
19

        Yields
        ------
        :any:`numpy.array`
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
20
            A frame of the video. The size is (3, 1920, 1080).
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
21
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
22
23
        vfilename = self.make_path(directory=self.original_directory)
        for frame in reader(vfilename):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
24
25
26
            # crop frames to 720 x 1024
            h, w = numpy.shape(frame)[-2:]
            dh, dw = (h - 720) // 2, (w - 1024) // 2
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
27
28
29
30
31
            if dh != 0:
                frame = frame[:, dh:-dh, :]
            if dw != 0:
                frame = frame[:, :, dw:-dw]
            assert frame.shape == self.frame_shape, frame.shape
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
32
33
            yield frame

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
34
35
    @property
    def number_of_frames(self):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
36
37
38
39
40
41
42
        """Returns the number of frames in a video file.

        Returns
        -------
        int
            The number of frames.
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
43
44
        vfilename = self.make_path(directory=self.original_directory)
        return reader(vfilename).number_of_frames
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
45
46
47
48
49
50
51
52

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

        Returns
        -------
        (int, int, int)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
            The (#Channels, Height, Width) which is :any:`UVAD_FRAME_SHAPE`.
        """
        return UVAD_FRAME_SHAPE


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

    def __init__(self, original_directory=rc['bob.db.uvad.directory'],
                 name='uvad', pad_file_class=None,
                 original_extension=None, **kwargs):
        if pad_file_class is None:
            pad_file_class = File
        filelists_directory = resource_filename(__name__, 'lists')
        super(Database, self).__init__(
            filelists_directory=filelists_directory, name=name,
            original_directory=original_directory,
            pad_file_class=pad_file_class,
            original_extension=original_extension,
            **kwargs)

    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

    def frames(self, padfile):
        return padfile.frames

    def number_of_frames(self, padfile):
        return padfile.number_of_frames

    @property
    def frame_shape(self):
        return UVAD_FRAME_SHAPE

    def annotations(self, padfile):
        """Reads the annotations for the given padfile.

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

        Returns
        -------
        dict
            The annotations as a dictionary, e.g.:
            ``{'0': {'reye':(re_y,re_x), 'leye':(le_y,le_x)}, ...}``
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
106
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        path = padfile.make_path(
            directory=self.original_directory, extension=None)
        # the annotations are in the uvad/release_1/face-locations-v3 folder
        path = path.replace('release_1', 'release_1/face-locations-v3')
        path = path[:-3] + 'face'
        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