query.py 7.74 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):
Pedro TOME's avatar
Pedro TOME committed
17 18 19 20 21 22
  """The dataset class opens and maintains a connection opened to the Database.

  It provides many different ways to probe for the characteristics of the data
  and for the data itself inside the database.
  """

23

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

André Anjos's avatar
André Anjos committed
28

29 30
  def protocol_names(self):
    """Returns a list of all supported protocols"""
André Anjos's avatar
André Anjos committed
31

32
    return tuple([k.name for k in self.query(Protocol).order_by(Protocol.name)])
Pedro TOME's avatar
Pedro TOME committed
33 34


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

38
    return Subset.purpose_choices
André Anjos's avatar
André Anjos committed
39 40


41 42
  def groups(self):
    """Returns a list of all supported groups"""
André Anjos's avatar
André Anjos committed
43

44
    return Subset.group_choices
Pedro TOME's avatar
Pedro TOME committed
45

André Anjos's avatar
André Anjos committed
46

47 48
  def genders(self):
    """Returns a list of all supported gender values"""
Pedro TOME's avatar
Pedro TOME committed
49

50
    return Client.gender_choices
André Anjos's avatar
André Anjos committed
51 52


53 54
  def finger_names(self):
    """Returns a list of all supported finger name values"""
55

56
    return Finger.name_choices
57

Pedro TOME's avatar
Pedro TOME committed
58

59 60
  def sessions(self):
    """Returns a list of all supported session values"""
Pedro TOME's avatar
Pedro TOME committed
61

62
    return File.session_choices
Pedro TOME's avatar
Pedro TOME committed
63 64


65 66
  def file_from_model_id(self, model_id):
    """Returns the file in the database given a ``model_id``"""
André Anjos's avatar
André Anjos committed
67

68
    return self.query(File).filter(File.model_id == model_id).one()
André Anjos's avatar
André Anjos committed
69 70


71 72
  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
73

74
    return self.file_from_model_id(model_id).unique_finger_name
Pedro TOME's avatar
Pedro TOME committed
75 76


77 78
  def model_ids(self, protocol=None, groups=None):
    """Returns a set of models for a given protocol/group
Pedro TOME's avatar
Pedro TOME committed
79

André Anjos's avatar
André Anjos committed
80 81
    Parameters:

82 83
      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
84

85 86 87 88 89
      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
90 91 92 93


    Returns:

94 95
      list: A list of string corresponding model identifiers with the specified
      filtering criteria
Pedro TOME's avatar
Pedro TOME committed
96 97 98

    """

99 100 101 102 103
    protocols = None
    if protocol:
      valid_protocols = self.protocol_names()
      protocols = self.check_parameters_for_validity(protocol, "protocol",
                                                     valid_protocols)
Pedro TOME's avatar
Pedro TOME committed
104

105 106 107 108
    if groups:
      valid_groups = self.groups()
      groups = self.check_parameters_for_validity(groups, "group",
                                                  valid_groups)
André Anjos's avatar
André Anjos committed
109

110
    retval = self.query(File)
Pedro TOME's avatar
Pedro TOME committed
111

112 113
    joins = []
    filters = []
Pedro TOME's avatar
Pedro TOME committed
114

115 116
    subquery = self.query(Subset)
    subfilters = []
Pedro TOME's avatar
Pedro TOME committed
117

118 119 120
    if protocols:
      subquery = subquery.join(Protocol)
      subfilters.append(Protocol.name.in_(protocols))
André Anjos's avatar
André Anjos committed
121

122 123
    if groups:
      subfilters.append(Subset.group.in_(groups))
André Anjos's avatar
André Anjos committed
124

125
    subfilters.append(Subset.purpose == 'enroll')
André Anjos's avatar
André Anjos committed
126

127 128
    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
129

130
    retval = retval.join(*joins).filter(*filters).distinct().order_by('id')
Pedro TOME's avatar
Pedro TOME committed
131

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


135 136 137
  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
138

Pedro TOME's avatar
Pedro TOME committed
139

140
    Parameters:
André Anjos's avatar
André Anjos committed
141

142 143
      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
144

145 146 147 148 149 150 151 152
      groups (:py:class:`str`, :py:class:`list`, optional): One or more of the
        supported groups. If not set, returns data from all groups

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

      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
153

154 155 156 157 158 159 160 161
      genders (:py:class:`str`, :py:class:`list`, optional): If set, limit
        output using the provided gender identifiers

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

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


André Anjos's avatar
André Anjos committed
164 165
    Returns:

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

Pedro TOME's avatar
Pedro TOME committed
169
    """
170

171 172 173 174 175
    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
176

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

182 183 184 185
    if purposes:
      valid_purposes = self.purposes()
      purposes = self.check_parameters_for_validity(purposes, "purpose",
                                                    valid_purposes)
Pedro TOME's avatar
Pedro TOME 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
    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
196

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

202 203 204 205
    if finger_names:
      valid_finger_names = self.finger_names()
      finger_names = self.check_parameters_for_validity(finger_names,
          "finger_names", valid_finger_names)
André Anjos's avatar
André Anjos committed
206

207 208 209 210
    if sessions:
      valid_sessions = self.sessions()
      sessions = self.check_parameters_for_validity(sessions, "sessions",
                                                    valid_sessions)
Pedro TOME's avatar
Pedro TOME committed
211

212
    retval = self.query(File)
Pedro TOME's avatar
Pedro TOME committed
213

214 215
    joins = []
    filters = []
André Anjos's avatar
André Anjos committed
216

217
    if protocols or groups or purposes:
Pedro TOME's avatar
Pedro TOME committed
218

219 220
      subquery = self.query(Subset)
      subfilters = []
Pedro TOME's avatar
Pedro TOME committed
221

222 223 224
      if protocols:
        subquery = subquery.join(Protocol)
        subfilters.append(Protocol.name.in_(protocols))
Pedro TOME's avatar
Pedro TOME committed
225

226 227 228 229
      if groups:
        subfilters.append(Subset.group.in_(groups))
      if purposes:
        subfilters.append(Subset.purpose.in_(purposes))
Pedro TOME's avatar
Pedro TOME committed
230

231
      subsets = subquery.filter(*subfilters)
Pedro TOME's avatar
Pedro TOME committed
232

233
      filters.append(File.subsets.any(Subset.id.in_([k.id for k in subsets])))
André Anjos's avatar
André Anjos committed
234

235 236
    if genders or finger_names:
      joins.append(Finger)
Pedro TOME's avatar
Pedro TOME committed
237

238 239 240 241
      if genders:
        fingers = self.query(Finger).join(
            Client).filter(Client.gender.in_(genders))
        filters.append(Finger.id.in_([k.id for k in fingers]))
André Anjos's avatar
André Anjos committed
242

243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259
      if finger_names:
        filters.append(Finger.name.in_(finger_names))

    if sessions:
      filters.append(File.session.in_(sessions))

    if model_ids:
      filters.append(File.model_id.in_(model_ids))

    # 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')
260

261
    return list(retval)