create.py 7.12 KB
Newer Older
1
2
3
#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
# Laurent El Shafey <Laurent.El-Shafey@idiap.ch>
4
# Fri 20 May 17:00:50 2011
5
6
7
8
9
10
11
12

"""This script creates the BANCA database in a single pass.
"""

import os

from .models import *

13
def add_files(session, imagedir, verbose):
14
  """Add files (and clients) to the BANCA database."""
15

16
  def add_file(session, filename, client_dict, verbose):
17
18
19
20
21
22
23
24
25
26
    """Parse a single filename and add it to the list.
       Also add a client entry if not already in the database."""

    v = os.path.splitext(os.path.basename(filename))[0].split('_')
    if not (v[0] in client_dict):
      if (v[2] == 'wm'):
        v[2] = 'world'
      session.add(Client(int(v[0]), v[1], v[2], v[5]))
      client_dict[v[0]] = True
    session_id = int(v[3].split('s')[1])
27
    if verbose: print "Adding file '%s'..." %(os.path.basename(filename).split('.')[0], )
28
    session.add(File(int(v[0]), os.path.basename(filename).split('.')[0], v[4], v[6], session_id))
29

30
  file_list = os.listdir(imagedir)
31
  client_dict = {}
32
  for filename in file_list:
33
    add_file(session, os.path.join(imagedir, filename), client_dict, verbose)
34

35
def add_subworlds(session, verbose):
36
  """Adds splits in the world set, based on the client ids"""
37
38
39
40

  # one third and two thirds
  snames = ["onethird", "twothirds"]
  slist = [ [9003, 9005, 9027, 9033, 9035, 9043, 9049, 9053, 9055, 9057],
41
42
            [9001, 9007, 9009, 9011, 9013, 9015, 9017, 9019, 9021, 9023,
             9025, 9029, 9031, 9037, 9039, 9041, 9045, 9047, 9051, 9059] ]
43
44
45
46
47
48
49
50
51
52
53
54
  for k in range(len(snames)):
    if verbose: print "Adding subworld '%s'" %(snames[k], )
    su = Subworld(snames[k])
    session.add(su)
    session.flush()
    session.refresh(su)
    l = slist[k]
    for c_id in l:
      if verbose: print "Adding client '%d' to subworld '%s'..." %(c_id, snames[k])
      su.clients.append(session.query(Client).filter(Client.id == c_id).first())

def add_protocols(session, verbose):
55
  """Adds protocols"""
56
57
58
59

  # 1. DEFINITIONS
  # Numbers in the lists correspond to session identifiers
  protocol_definitions = {}
60

61
  # Protocol Mc
62
63
64
65
66
  enrol = [1]
  probe_c = [2, 3, 4]
  probe_i = [1, 2, 3, 4]
  protocol_definitions['Mc'] = [enrol, probe_c, probe_i]

67
  # Protocol Md
68
69
70
71
  enrol = [5]
  probe_c = [6, 7, 8]
  probe_i = [5, 6, 7, 8]
  protocol_definitions['Md'] = [enrol, probe_c, probe_i]
72

73
  # Protocol Ma
74
75
76
77
  enrol = [9]
  probe_c = [10, 11, 12]
  probe_i = [9, 10, 11, 12]
  protocol_definitions['Ma'] = [enrol, probe_c, probe_i]
78

79
  # Protocol Ud
80
81
82
83
  enrol = [1]
  probe_c = [6, 7, 8]
  probe_i = [5, 6, 7, 8]
  protocol_definitions['Ud'] = [enrol, probe_c, probe_i]
84

85
86
87
88
89
  # Protocol Ua
  enrol = [1]
  probe_c = [10, 11, 12]
  probe_i = [9, 10, 11, 12]
  protocol_definitions['Ua'] = [enrol, probe_c, probe_i]
90

91
  # Protocol P
92
93
94
95
  enrol = [1]
  probe_c = [2, 3, 4, 6, 7, 8, 10, 11, 12]
  probe_i = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
  protocol_definitions['P'] = [enrol, probe_c, probe_i]
96

97
  # Protocol G
98
99
100
101
  enrol = [1, 5, 9]
  probe_c = [2, 3, 4, 6, 7, 8, 10, 11, 12]
  probe_i = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
  protocol_definitions['G'] = [enrol, probe_c, probe_i]
102

103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
  # 2. ADDITIONS TO THE SQL DATABASE
  protocolPurpose_list = [('world', 'train'), ('dev', 'enrol'), ('dev', 'probe'), ('eval', 'enrol'), ('eval', 'probe')]
  for proto in protocol_definitions:
    p = Protocol(proto)
    # Add protocol
    if verbose: print "Adding protocol %s..." % (proto)
    session.add(p)
    session.flush()
    session.refresh(p)

    # Add protocol purposes
    for key in range(len(protocolPurpose_list)):
      purpose = protocolPurpose_list[key]
      pu = ProtocolPurpose(p.id, purpose[0], purpose[1])
      if verbose: print "  Adding protocol purpose ('%s','%s')..." % (purpose[0], purpose[1])
      session.add(pu)
      session.flush()
      session.refresh(pu)

       # Add files attached with this protocol purpose
      client_group = ""
      if(key == 0): client_group = "world"
      elif(key == 1 or key == 2): client_group = "g1"
      elif(key == 3 or key == 4): client_group = "g2"
      session_list = []
      session_list_i = []
129
      if(key == 1 or key == 3):
130
        session_list = protocol_definitions[proto][0]
131
      elif(key == 2):
132
133
        session_list = protocol_definitions[proto][1]
        session_list_i = protocol_definitions[proto][2]
134
      elif(key == 4):
135
136
137
138
139
140
141
142
143
144
145
146
147
        session_list = protocol_definitions[proto][1]
        session_list_i = protocol_definitions[proto][2]

      # Adds 'regular' files (i.e. 'world', 'enrol', 'probe client')
      if not session_list:
        q = session.query(File).join(Client).filter(Client.sgroup == client_group).\
              order_by(File.id)
        for k in q:
          if verbose: print "    Adding protocol file '%s'..." % (k.path)
          pu.files.append(k)
      else:
        for sid in session_list:
          q = session.query(File).join(Client).filter(Client.sgroup == client_group).\
148
                filter(and_(File.session_id == sid, File.client_id == File.claimed_id)).\
149
150
151
152
153
154
155
156
157
                order_by(File.id)
          for k in q:
            if verbose: print "    Adding protocol file '%s'..." % (k.path)
            pu.files.append(k)

      # Adds impostors if required
      if session_list_i:
        for sid in session_list_i:
          q = session.query(File).join(Client).filter(Client.sgroup == client_group).\
158
                filter(and_(File.session_id == sid, File.client_id != File.claimed_id)).\
159
160
161
162
                order_by(File.id)
          for k in q:
            if verbose: print "    Adding protocol file '%s'..." % (k.path)
            pu.files.append(k)
163
164
165
166
167


def create_tables(args):
  """Creates all necessary tables (only to be used at the first time)"""

168
  from bob.db.utils import create_engine_try_nolock
169

170
  engine = create_engine_try_nolock(args.type, args.files[0], echo=(args.verbose >= 2))
171
  Base.metadata.create_all(engine)
172
173
174
175
176
177
178

# Driver API
# ==========

def create(args):
  """Creates or re-creates this database"""

179
  from bob.db.utils import session_try_nolock
180
181
182

  dbfile = args.files[0]

183
  if args.recreate:
184
185
186
187
188
189
190
191
192
    if args.verbose and os.path.exists(dbfile):
      print('unlinking %s...' % dbfile)
    if os.path.exists(dbfile): os.unlink(dbfile)

  if not os.path.exists(os.path.dirname(dbfile)):
    os.makedirs(os.path.dirname(dbfile))

  # the real work...
  create_tables(args)
193
  s = session_try_nolock(args.type, args.files[0], echo=(args.verbose >= 2))
194
195
196
  add_files(s, args.imagedir, args.verbose)
  add_subworlds(s, args.verbose)
  add_protocols(s, args.verbose)
197
198
199
200
201
202
203
204
205
206
  s.commit()
  s.close()

def add_command(subparsers):
  """Add specific subcommands that the action "create" can use"""

  parser = subparsers.add_parser('create', help=create.__doc__)

  parser.add_argument('-R', '--recreate', action='store_true', default=False,
      help="If set, I'll first erase the current database")
André Anjos's avatar
André Anjos committed
207
  parser.add_argument('-v', '--verbose', action='count',
208
      help="Do SQL operations in a verbose way")
209
  parser.add_argument('-D', '--imagedir', action='store', metavar='DIR',
210
211
      default='/idiap/group/vision/visidiap/databases/banca/english/images_gray',
      help="Change the relative path to the directory containing the images of the BANCA database (defaults to %(default)s)")
212

213
  parser.set_defaults(func=create) #action