query_bio.py 4.54 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1
2
#!/usr/bin/env python

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
3
from bob.bio.spear.database import AudioBioFile
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
4
5
from bob.bio.video.database import VideoBioFile
from bob.bio.video.utils import FrameSelector
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
6
import bob.bio.base
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
7
8
9
import bob.io.base
import bob.io.video
import numpy as np
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
10
from bob.extension import rc
11
from .common import swan_bio_file_metadata, read_audio
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
12
13


14
15
class SwanBioFile(object):
    """A base class for SWAN bio files which can handle the metadata."""
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
16

17
18
19
20
21
22
    def __init__(self, **kwargs):
        super(SwanBioFile, self).__init__(**kwargs)
        (
            self.client, self.session, self.nrecording,
            self.device, self.modality
        ) = swan_bio_file_metadata(self.path)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
23

24
25

class SwanAudioBioFile(AudioBioFile, SwanBioFile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
26
27
    """SwanAudioBioFile are video files actually"""

28
29
30
31
32
    def __init__(self, **kwargs):
        super(SwanAudioBioFile, self).__init__(**kwargs)

    def load(self, directory=None, extension=None):
        if extension is None:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
33
34
            video_path = self.make_path(directory, extension)
            rate, audio = read_audio(video_path)
35
            return rate, np.cast['float'](audio)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
36
37
38
39
        else:
            return super(SwanAudioBioFile, self).load(directory, extension)


40
class SwanVideoBioFile(VideoBioFile, SwanBioFile):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
41
42
    """SwanVideoBioFile are video files actually"""

43
44
45
    def __init__(self, **kwargs):
        super(SwanVideoBioFile, self).__init__(**kwargs)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
46
47
48
49
    def swap(self, data):
        # rotate the video or image since SWAN videos are not upright!
        return np.swapaxes(data, -2, -1)

50
51
52
    def load(self, directory=None, extension=None,
             frame_selector=FrameSelector(selection_style='all')):
        if extension is None:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
53
54
            video_path = self.make_path(directory or self.original_directory,
                                        extension)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
55
            video = bob.io.base.load(video_path)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
56
            video = self.swap(video)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
57
58
59
60
61
            return frame_selector(video)
        else:
            return super(SwanVideoBioFile, self).load(
                directory, extension, frame_selector)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
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
106
    def frames(self):
        """Yields the frames of the biofile one by one.

        Parameters
        ----------
        biofile : :any:`SwanVideoBioFile`
            The high-level bio file

        Yields
        ------
        :any:`numpy.array`
            A frame of the video. The size is (3, 1280, 720).
        """
        vfilename = self.make_path(directory=self.original_directory)
        reader = bob.io.video.reader(vfilename)
        for frame in reader:
            yield self.swap(frame)

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

        Parameters
        ----------
        biofile : :any:`SwanVideoBioFile`
            The high-level bio file

        Returns
        -------
        int
            The number of frames.
        """
        vfilename = self.make_path(directory=self.original_directory)
        return bob.io.video.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 (3, 1920, 1080).
        """
        return (3, 1280, 720)

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
107

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
108
109
110
111
112
class Database(bob.bio.base.database.FileListBioDatabase):
    """Wrapper class for the SWAN database for speaker recognition
    (http://www.idiap.ch/dataset/swan). This class defines a simple protocol
    for training, dev and and by splitting the audio files of the database in
    three main parts.
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
113
114
    """

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
115
    def __init__(self, original_directory=rc['bob.db.swan.directory'],
116
                 bio_file_class=SwanAudioBioFile, name='swan', **kwargs):
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
117
118
        # call base class constructor
        from pkg_resources import resource_filename
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
119
        folder = resource_filename(__name__, 'lists')
120
        super(Database, self).__init__(
121
            folder, name=name, bio_file_class=bio_file_class,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
122
            original_directory=original_directory,
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
123
124
            training_depends_on_protocol=True,
            models_depend_on_protocol=True,
125
126
            **kwargs
        )
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140

    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 load_frames(biofile, directory, extension):
    for frame in biofile.frames():
        yield frame