create.py 7.39 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 14
def nodot(item):
  """Can be used to ignore hidden files, starting with the . character."""
15
  return item[0] != '.'
16

17
def add_files(session, imagedir, verbose):
18
  """Add files (and clients) to the BANCA database."""
19

20
  def add_file(session, subdir, filename, client_dict, verbose):
21 22 23 24 25 26 27 28 29 30
    """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])
Laurent EL SHAFEY's avatar
Laurent EL SHAFEY committed
31
    base_path = os.path.join(subdir, os.path.basename(filename).split('.')[0])
32 33
    if verbose: print "Adding file '%s'..." %(base_path, )
    session.add(File(int(v[0]), base_path, v[4], v[6], session_id))
34

35
  subdir_list = filter(nodot, os.listdir(imagedir))
36
  client_dict = {}
37 38 39
  for subdir in subdir_list:
    file_list = filter(nodot, os.listdir(os.path.join(imagedir, subdir)))
    for filename in file_list:
40
      add_file(session, subdir, os.path.join(imagedir, filename), client_dict, verbose)
41

42
def add_subworlds(session, verbose):
43
  """Adds splits in the world set, based on the client ids"""
44 45 46 47

  # one third and two thirds
  snames = ["onethird", "twothirds"]
  slist = [ [9003, 9005, 9027, 9033, 9035, 9043, 9049, 9053, 9055, 9057],
48 49
            [9001, 9007, 9009, 9011, 9013, 9015, 9017, 9019, 9021, 9023,
             9025, 9029, 9031, 9037, 9039, 9041, 9045, 9047, 9051, 9059] ]
50 51 52 53 54 55 56 57 58 59 60 61
  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):
62
  """Adds protocols"""
63 64 65 66

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

68
  # Protocol Mc
69 70 71 72 73
  enrol = [1]
  probe_c = [2, 3, 4]
  probe_i = [1, 2, 3, 4]
  protocol_definitions['Mc'] = [enrol, probe_c, probe_i]

74
  # Protocol Md
75 76 77 78
  enrol = [5]
  probe_c = [6, 7, 8]
  probe_i = [5, 6, 7, 8]
  protocol_definitions['Md'] = [enrol, probe_c, probe_i]
79

80
  # Protocol Ma
81 82 83 84
  enrol = [9]
  probe_c = [10, 11, 12]
  probe_i = [9, 10, 11, 12]
  protocol_definitions['Ma'] = [enrol, probe_c, probe_i]
85

86
  # Protocol Ud
87 88 89 90
  enrol = [1]
  probe_c = [6, 7, 8]
  probe_i = [5, 6, 7, 8]
  protocol_definitions['Ud'] = [enrol, probe_c, probe_i]
91

92 93 94 95 96
  # Protocol Ua
  enrol = [1]
  probe_c = [10, 11, 12]
  probe_i = [9, 10, 11, 12]
  protocol_definitions['Ua'] = [enrol, probe_c, probe_i]
97

98
  # Protocol P
99 100 101 102
  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]
103

104
  # Protocol G
105 106 107 108
  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]
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
  # 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 = []
136
      if(key == 1 or key == 3):
137
        session_list = protocol_definitions[proto][0]
138
      elif(key == 2):
139 140
        session_list = protocol_definitions[proto][1]
        session_list_i = protocol_definitions[proto][2]
141
      elif(key == 4):
142 143 144 145 146 147 148 149 150 151 152 153 154
        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).\
155
                filter(and_(File.session_id == sid, File.client_id == File.claimed_id)).\
156 157 158 159 160 161 162 163 164
                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).\
165
                filter(and_(File.session_id == sid, File.client_id != File.claimed_id)).\
166 167 168 169
                order_by(File.id)
          for k in q:
            if verbose: print "    Adding protocol file '%s'..." % (k.path)
            pu.files.append(k)
170 171 172 173 174


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

175
  from bob.db.utils import create_engine_try_nolock
176

177
  engine = create_engine_try_nolock(args.type, args.files[0], echo=(args.verbose >= 2))
178
  Base.metadata.create_all(engine)
179 180 181 182 183 184 185

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

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

186
  from bob.db.utils import session_try_nolock
187 188 189

  dbfile = args.files[0]

190
  if args.recreate:
191 192 193 194 195 196 197 198 199
    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)
200
  s = session_try_nolock(args.type, args.files[0], echo=(args.verbose >= 2))
201 202 203
  add_files(s, args.imagedir, args.verbose)
  add_subworlds(s, args.verbose)
  add_protocols(s, args.verbose)
204 205 206 207 208 209 210 211 212 213
  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
214
  parser.add_argument('-v', '--verbose', action='count',
215
      help="Do SQL operations in a verbose way")
216
  parser.add_argument('-D', '--imagedir', action='store', metavar='DIR',
217
      default='/idiap/group/biometric/databases/banca/english/images/images/',
218
      help="Change the relative path to the directory containing the images of the BANCA database (defaults to %(default)s)")
219

220
  parser.set_defaults(func=create) #action