query.py 8.36 KB
Newer Older
Pedro TOME's avatar
Pedro TOME committed
1
2
3
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :

4
"""Dataset interface allowing the user to query the UTFVP database"""
Pedro TOME's avatar
Pedro TOME committed
5
6
7
8

import six
from .models import *
from .driver import Interface
9
from sqlalchemy import and_, not_
Pedro TOME's avatar
Pedro TOME committed
10

André Anjos's avatar
André Anjos committed
11
import bob.db.base
Pedro TOME's avatar
Pedro TOME committed
12
13
14

SQLITE_FILE = Interface().files()[0]

15

André Anjos's avatar
André Anjos committed
16
class Database(bob.db.base.SQLiteDatabase):
17
    """The dataset class opens and maintains a connection opened to the Database.
Pedro TOME's avatar
Pedro TOME committed
18

19
20
21
    It provides many different ways to probe for the characteristics of the data
    and for the data itself inside the database.
    """
Pedro TOME's avatar
Pedro TOME committed
22

23
24
25
26
    def __init__(self, original_directory=None, original_extension=None):
        super(Database, self).__init__(
            SQLITE_FILE, File, original_directory, original_extension
        )
André Anjos's avatar
André Anjos committed
27

28
29
    def protocol_names(self):
        """Returns a list of all supported protocols"""
Pedro TOME's avatar
Pedro TOME committed
30

31
        return tuple([k.name for k in self.query(Protocol).order_by(Protocol.name)])
André Anjos's avatar
André Anjos committed
32

33
34
    def purposes(self):
        """Returns a list of all supported purposes"""
André Anjos's avatar
André Anjos committed
35

36
        return Subset.purpose_choices
37

38
39
    def groups(self):
        """Returns a list of all supported groups"""
40

41
        return Subset.group_choices
Pedro TOME's avatar
Pedro TOME committed
42

43
44
    def genders(self):
        """Returns a list of all supported gender values"""
Pedro TOME's avatar
Pedro TOME committed
45

46
        return Client.gender_choices
Pedro TOME's avatar
Pedro TOME committed
47

48
49
    def finger_names(self):
        """Returns a list of all supported finger name values"""
Pedro TOME's avatar
Pedro TOME committed
50

51
        return Finger.name_choices
André Anjos's avatar
André Anjos committed
52

53
54
    def sessions(self):
        """Returns a list of all supported session values"""
André Anjos's avatar
André Anjos committed
55

56
        return File.session_choices
André Anjos's avatar
André Anjos committed
57

58
59
    def file_from_model_id(self, model_id):
        """Returns the file in the database given a ``model_id``"""
Pedro TOME's avatar
Pedro TOME committed
60

61
        return self.query(File).filter(File.model_id == model_id).one()
Pedro TOME's avatar
Pedro TOME committed
62

63
64
    def finger_name_from_model_id(self, model_id):
        """Returns the unique finger name in the database given a ``model_id``"""
Pedro TOME's avatar
Pedro TOME committed
65

66
        return self.file_from_model_id(model_id).unique_finger_name
Pedro TOME's avatar
Pedro TOME committed
67

68
69
    def model_ids(self, protocol=None, groups=None):
        """Returns a set of models for a given protocol/group
André Anjos's avatar
André Anjos committed
70

71
        Parameters:
André Anjos's avatar
André Anjos committed
72

73
74
          protocol (:py:class:`str`, :py:class:`list`, optional): One or more of
            the supported protocols.  If not set, returns data from all protocols
André Anjos's avatar
André Anjos committed
75

76
77
78
79
80
          groups (:py:class:`str`, :py:class:`list`, optional): One or more of the
            supported groups. If not set, returns data from all groups. Notice this
            parameter should either not set or set to ``dev``. Otherwise, this
            method will return an empty list given we don't have a test set, only a
            development set.
André Anjos's avatar
André Anjos committed
81
82


83
        Returns:
Pedro TOME's avatar
Pedro TOME committed
84

85
86
          list: A list of string corresponding model identifiers with the specified
          filtering criteria
Pedro TOME's avatar
Pedro TOME committed
87

88
        """
Pedro TOME's avatar
Pedro TOME committed
89

90
91
92
93
94
95
        protocols = None
        if protocol:
            valid_protocols = self.protocol_names()
            protocols = self.check_parameters_for_validity(
                protocol, "protocol", valid_protocols
            )
André Anjos's avatar
André Anjos committed
96

97
98
99
        if groups:
            valid_groups = self.groups()
            groups = self.check_parameters_for_validity(groups, "group", valid_groups)
Pedro TOME's avatar
Pedro TOME committed
100

101
        retval = self.query(File)
Pedro TOME's avatar
Pedro TOME committed
102

103
104
        joins = []
        filters = []
Pedro TOME's avatar
Pedro TOME committed
105

106
107
        subquery = self.query(Subset)
        subfilters = []
André Anjos's avatar
André Anjos committed
108

109
110
111
        if protocols:
            subquery = subquery.join(Protocol)
            subfilters.append(Protocol.name.in_(protocols))
André Anjos's avatar
André Anjos committed
112

113
114
        if groups:
            subfilters.append(Subset.group.in_(groups))
André Anjos's avatar
André Anjos committed
115

116
        subfilters.append(Subset.purpose == "enroll")
Pedro TOME's avatar
Pedro TOME committed
117

118
119
        subsets = subquery.filter(*subfilters)
        filters.append(File.subsets.any(Subset.id.in_([k.id for k in subsets])))
Pedro TOME's avatar
Pedro TOME committed
120

121
        retval = retval.join(*joins).filter(*filters).distinct().order_by("id")
Pedro TOME's avatar
Pedro TOME committed
122

123
        return sorted(set([k.model_id for k in retval.distinct()]))
Pedro TOME's avatar
Pedro TOME committed
124

125
126
127
128
129
130
131
132
133
134
135
    def objects(
        self,
        protocol=None,
        groups=None,
        purposes=None,
        model_ids=None,
        genders=None,
        finger_names=None,
        sessions=None,
    ):
        """Returns objects filtered by criteria
André Anjos's avatar
André Anjos committed
136

Pedro TOME's avatar
Pedro TOME committed
137

138
        Parameters:
André Anjos's avatar
André Anjos committed
139

140
141
          protocol (:py:class:`str`, :py:class:`list`, optional): One or more of
            the supported protocols. If not set, returns data from all protocols
Pedro TOME's avatar
Pedro TOME committed
142

143
144
          groups (:py:class:`str`, :py:class:`list`, optional): One or more of the
            supported groups. If not set, returns data from all groups
145

146
147
          purposes (:py:class:`str`, :py:class:`list`, optional): One or more of
            the supported purposes. If not set, returns data for all purposes
148

149
150
          model_ids (:py:class:`str`, :py:class:`list`, optional): If set, limit
            output using the provided model identifiers
André Anjos's avatar
André Anjos committed
151

152
153
          genders (:py:class:`str`, :py:class:`list`, optional): If set, limit
            output using the provided gender identifiers
154

155
156
          finger_names (:py:class:`str`, :py:class:`list`, optional): If set, limit
            output using the provided finger name identifier
157

158
159
          sessions (:py:class:`str`, :py:class:`list`, optional): If set, limit
            output using the provided session identifiers
Pedro TOME's avatar
Pedro TOME committed
160
161


162
        Returns:
André Anjos's avatar
André Anjos committed
163

164
165
          list: A list of :py:class:`File` objects corresponding to the filtering
          criteria.
André Anjos's avatar
André Anjos committed
166

167
        """
168

169
        protocols = None
André Anjos's avatar
André Anjos committed
170

171
172
173
174
175
        if protocol:
            valid_protocols = self.protocol_names()
            protocols = self.check_parameters_for_validity(
                protocol, "protocol", valid_protocols
            )
Pedro TOME's avatar
Pedro TOME committed
176

177
178
179
        if groups:
            valid_groups = self.groups()
            groups = self.check_parameters_for_validity(groups, "group", valid_groups)
Pedro TOME's avatar
Pedro TOME committed
180

181
182
183
184
185
        if purposes:
            valid_purposes = self.purposes()
            purposes = self.check_parameters_for_validity(
                purposes, "purpose", valid_purposes
            )
André Anjos's avatar
André Anjos committed
186

187
188
189
190
        # if only asking for 'probes', then ignore model_ids as all of our
        # protocols do a full probe-model scan
        if purposes and len(purposes) == 1 and "probe" in purposes:
            model_ids = None
André Anjos's avatar
André Anjos committed
191

192
193
194
195
196
        if model_ids:
            valid_model_ids = self.model_ids(protocol, groups)
            model_ids = self.check_parameters_for_validity(
                model_ids, "model_ids", valid_model_ids
            )
André Anjos's avatar
André Anjos committed
197

198
199
200
201
202
        if genders:
            valid_genders = self.genders()
            genders = self.check_parameters_for_validity(
                genders, "genders", valid_genders
            )
André Anjos's avatar
André Anjos committed
203

204
205
206
207
208
        if finger_names:
            valid_finger_names = self.finger_names()
            finger_names = self.check_parameters_for_validity(
                finger_names, "finger_names", valid_finger_names
            )
Pedro TOME's avatar
Pedro TOME committed
209

210
211
212
213
214
        if sessions:
            valid_sessions = self.sessions()
            sessions = self.check_parameters_for_validity(
                sessions, "sessions", valid_sessions
            )
Pedro TOME's avatar
Pedro TOME committed
215

216
        retval = self.query(File)
André Anjos's avatar
André Anjos committed
217

218
219
        joins = []
        filters = []
Pedro TOME's avatar
Pedro TOME committed
220

221
        if protocols or groups or purposes:
Pedro TOME's avatar
Pedro TOME committed
222

223
224
            subquery = self.query(Subset)
            subfilters = []
Pedro TOME's avatar
Pedro TOME committed
225

226
227
228
            if protocols:
                subquery = subquery.join(Protocol)
                subfilters.append(Protocol.name.in_(protocols))
Pedro TOME's avatar
Pedro TOME committed
229

230
231
232
233
            if groups:
                subfilters.append(Subset.group.in_(groups))
            if purposes:
                subfilters.append(Subset.purpose.in_(purposes))
Pedro TOME's avatar
Pedro TOME committed
234

235
            subsets = subquery.filter(*subfilters)
André Anjos's avatar
André Anjos committed
236

237
            filters.append(File.subsets.any(Subset.id.in_([k.id for k in subsets])))
Pedro TOME's avatar
Pedro TOME committed
238

239
        #import ipdb
André Anjos's avatar
André Anjos committed
240

241
242
243
        #ipdb.set_trace()
        joins.append(Finger)
        if genders or finger_names:
244

245
246
247
248
249
            if genders:
                fingers = (
                    self.query(Finger).join(Client).filter(Client.gender.in_(genders))
                )
                filters.append(Finger.id.in_([k.id for k in fingers]))
250

251
252
            if finger_names:
                filters.append(Finger.name.in_(finger_names))
253

254
255
        if sessions:
            filters.append(File.session.in_(sessions))
256

257
258
        if model_ids:
            filters.append(File.model_id.in_(model_ids))
259

260
261
262
263
264
265
266
267
268
269
270
271
272
273
        # special case for 1vsall protocol: if only one model id given, returns
        # all but the sample for the model id in the list
        if (
            model_ids
            and len(model_ids) == 1
            and protocols
            and len(protocols) == 1
            and protocols[0] == "1vsall"
        ):
            filters.append(~self.file_from_model_id(model_ids[0]))

        retval = retval.join(*joins).filter(*filters).distinct().order_by("id")

        return list(retval)