query.py 4.13 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
from bob.db.base.annotations import read_annotation_file
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
5
6
from bob.pad.base.database import FileListPadDatabase
from bob.pad.face.database import VideoPadFile
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
7
from pkg_resources import resource_filename
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
8
9
10
11
import numpy


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

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

        Yields
        ------
        :any:`numpy.array`
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
21
            A frame of the video. The size is (3, 1920, 1080).
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
22
        """
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
23
24
        vfilename = self.make_path(directory=self.original_directory)
        for frame in reader(vfilename):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
25
26
27
            # 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
28
            if dh != 0:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
29
                frame = frame[:, : -2 * dh, :]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
30
            if dw != 0:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
31
                frame = frame[:, :, : -2 * dw]
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
32
            assert frame.shape == self.frame_shape, frame.shape
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
33
34
            yield frame

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

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

    @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
54
55
56
57
            The (#Channels, Height, Width) which is :any:`UVAD_FRAME_SHAPE`.
        """
        return UVAD_FRAME_SHAPE

58
59
    @property
    def annotations(self):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
60
61
62
63
        # The annotations shipped with uvad are totally wrong. We load custom
        # annotations here
        path = self.make_path(self.annotation_directory, self.annotation_extension)
        annotations = read_annotation_file(path, annotation_type=self.annotation_type)
64
65
        return annotations

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
66
67
68
69

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

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
70
71
72
73
74
75
76
77
78
79
80
    def __init__(
        self,
        original_directory=rc["bob.db.uvad.directory"],
        annotation_directory=rc["bob.db.uvad.annotation_dir"],
        annotation_extension=".json",
        annotation_type="json",
        name="uvad",
        pad_file_class=None,
        original_extension=None,
        **kwargs
    ):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
81
82
        if pad_file_class is None:
            pad_file_class = File
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
83
        filelists_directory = resource_filename(__name__, "lists")
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
84
        super(Database, self).__init__(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
85
86
            filelists_directory=filelists_directory,
            name=name,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
87
88
89
            original_directory=original_directory,
            pad_file_class=pad_file_class,
            original_extension=original_extension,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
90
91
92
93
94
95
96
97
98
99
100
101
            **kwargs
        )

    def objects(
        self,
        groups=None,
        protocol=None,
        purposes=None,
        model_ids=None,
        classes=None,
        **kwargs
    ):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
102
        files = super(Database, self).objects(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
103
104
105
106
107
108
109
            groups=groups,
            protocol=protocol,
            purposes=purposes,
            model_ids=model_ids,
            classes=classes,
            **kwargs
        )
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
110
111
        for f in files:
            f.original_directory = self.original_directory
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
112
113
114
            f.annotation_directory = self.annotation_directory
            f.annotation_extension = self.annotation_extension
            f.annotation_type = self.annotation_type
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
        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
140
        """
141
        return padfile.annotations