Mercurial Hosting > traffic-intelligence
diff python/storage.py @ 920:499154254f37
improved prototype loading
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Wed, 05 Jul 2017 16:30:04 -0400 |
parents | 7b3f2e0a2652 |
children | 630934595871 |
line wrap: on
line diff
--- a/python/storage.py Wed Jul 05 13:16:47 2017 -0400 +++ b/python/storage.py Wed Jul 05 16:30:04 2017 -0400 @@ -54,7 +54,7 @@ elif dataType == 'pois': dropTables(connection, ['gaussians2d', 'objects_pois']) elif dataType == 'prototype': - dropTables(connection, ['prototypes', 'prototype_positions', 'prototype_velocities']) + dropTables(connection, ['prototypes']) else: print('Unknown data type {} to delete from database'.format(dataType)) connection.close() @@ -592,7 +592,7 @@ # saving and loading for scene interpretation: POIs and Prototypes ######################### -def savePrototypesToSqlite(filename, prototypeIndices, trajectoryType, objects = None, nMatchings = None, dbFilenames = None): +def savePrototypesToSqlite(filename, prototypeIndices, trajectoryType, nMatchings = None, dbFilenames = None): '''save the prototype indices if objects is not None, the trajectories are also saved in prototype_positions and _velocities (prototypeIndices have to be in objects @@ -604,7 +604,7 @@ connection = sqlite3.connect(filename) cursor = connection.cursor() try: - cursor.execute('CREATE TABLE IF NOT EXISTS prototypes (id INTEGER, dbfilename VARCHAR, trajectory_type VARCHAR CHECK (trajectory_type IN (\"feature\", \"object\")), nmatchings INTEGER, positions_id INTEGER, PRIMARY KEY (id, dbfilename))') + cursor.execute('CREATE TABLE IF NOT EXISTS prototypes (id INTEGER, dbfilename VARCHAR, trajectory_type VARCHAR CHECK (trajectory_type IN (\"feature\", \"object\")), nmatchings INTEGER, PRIMARY KEY (id, dbfilename))') for i, protoId in enumerate(prototypeIndices): if nMatchings is not None: n = nMatchings[i] @@ -614,14 +614,7 @@ dbfn = dbFilenames[i] else: dbfn = filename - cursor.execute('INSERT INTO prototypes (id, dbfilename, trajectory_type, nmatchings, positions_id) VALUES ({},\"{}\",\"{}\",{}, {})'.format(protoId, dbfn, trajectoryType, n, i)) - if objects is not None: # save positions and velocities - features = [] - for i, o in enumerate(objects): - f = copy(o) - f.num = i - features.append(f) - saveTrajectoriesToTable(connection, features, 'feature', 'prototype') + cursor.execute('INSERT INTO prototypes (id, dbfilename, trajectory_type, nmatchings) VALUES (?,?,?,?)', (protoId, dbfn, trajectoryType, n)) except sqlite3.OperationalError as error: printDBError(error) connection.commit() @@ -630,7 +623,7 @@ def savePrototypeAssignments(filename, objects): pass -def loadPrototypesFromSqlite(filename): +def loadPrototypesFromSqlite(filename, withTrajectories = True): 'Loads prototype ids and matchings (if stored)' connection = sqlite3.connect(filename) cursor = connection.cursor() @@ -638,7 +631,7 @@ dbFilenames = [] trajectoryTypes = [] nMatchings = [] - trajectoryNumbers = [] + objects = [] try: cursor.execute('SELECT * FROM prototypes') for row in cursor: @@ -647,12 +640,15 @@ trajectoryTypes.append(row[2]) if row[3] is not None: nMatchings.append(row[3]) - if row[4] is not None: - trajectoryNumbers.append(row[4]) - if tableExists(connection, 'prototype_positions'): # load prototypes trajectories - objects = loadTrajectoriesFromSqlite(filename, 'feature', trajectoryNumbers, tablePrefix = 'prototype') - else: - objects = None + if withTrajectories: + loadingInformation = {} + for dbfn, trajType, protoId in zip(dbFilenames, trajectoryTypes, prototypeIndices): + if (dbfn, trajType) in loadingInformation: + loadingInformation[(dbfn, trajType)].append(protoId) + else: + loadingInformation[(dbfn, trajType)] = [protoId] + for k, v in loadingInformation.iteritems(): + objects += loadTrajectoriesFromSqlite(k[0], k[1], v) except sqlite3.OperationalError as error: printDBError(error) connection.close()