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

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

import os

from .models import *

13
def add_files(session, imagedir, verbose):
14
15
  """Add files (and clients) to the BANCA database."""
 
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
29
30
31
32
    session.add(File(int(v[0]), os.path.basename(filename).split('.')[0], v[4], v[6], session_id))
  
  file_list = os.listdir(imagedir)
  client_dict = {} 
  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
41
42
43
44
45
46
47
48
49
50
51
52
53
54

  # one third and two thirds
  snames = ["onethird", "twothirds"]
  slist = [ [9003, 9005, 9027, 9033, 9035, 9043, 9049, 9053, 9055, 9057],
            [9001, 9007, 9009, 9011, 9013, 9015, 9017, 9019, 9021, 9023, 
             9025, 9029, 9031, 9037, 9039, 9041, 9045, 9047, 9051, 9059] ] 
  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
60

  # 1. DEFINITIONS
  # Numbers in the lists correspond to session identifiers
  protocol_definitions = {}
 
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
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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
  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]
   
  # 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 = []
      if(key == 1 or key == 3): 
        session_list = protocol_definitions[proto][0]
      elif(key == 2): 
        session_list = protocol_definitions[proto][1]
        session_list_i = protocol_definitions[proto][2]
      elif(key == 4): 
        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).\
                filter(and_(File.session_id == sid, File.real_id == File.claimed_id)).\
                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).\
                filter(and_(File.session_id == sid, File.real_id != File.claimed_id)).\
                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
183
184
185
186
187
188
189
190
191
192

  dbfile = args.files[0]

  if args.recreate: 
    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
212
213
      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)")
  
  parser.set_defaults(func=create) #action