Mercurial > hg > nsaunier > traffic-intelligence

comparison trafficintelligence/storage.py @ 1037:6a6c37eb3a74

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
added function to load prototype assignments
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Tue, 03 Jul 2018 23:58:16 -0400
parents 933588568bec
children 5621e4ad2428
comparison
equal deleted inserted replaced
1036:0d7e5e290ea3 1037:6a6c37eb3a74
19 'truck':3} 19 'truck':3}
20 20
21 tableNames = {'feature':'positions', 21 tableNames = {'feature':'positions',
22 'object': 'objects', 22 'object': 'objects',
23 'objectfeatures': 'positions'} 23 'objectfeatures': 'positions'}
24
25 assignmentTableNames = {'feature':'positions',
26 'object': 'objects',
27 'objectfeatures': 'positions'}
28 24
29 ######################### 25 #########################
30 # Sqlite 26 # Sqlite
31 ######################### 27 #########################
32 28
591 cursor.execute('UPDATE prototypes SET nmatchings = {} WHERE prototype_filename = \"{}\" AND prototype_id = {} AND trajectory_type = \"{}\"'.format(nMatchings, p.getFilename(), p.getNum(), p.getTrajectoryType())) 587 cursor.execute('UPDATE prototypes SET nmatchings = {} WHERE prototype_filename = \"{}\" AND prototype_id = {} AND trajectory_type = \"{}\"'.format(nMatchings, p.getFilename(), p.getNum(), p.getTrajectoryType()))
592 except sqlite3.OperationalError as error: 588 except sqlite3.OperationalError as error:
593 printDBError(error) 589 printDBError(error)
594 connection.commit() 590 connection.commit()
595 591
592 def prototypeAssignmentNames(objectType):
593 tableName = objectType+'s_prototypes'
594 if objectType == 'feature':
595 #tableName = 'features_prototypes'
596 objectIdColumnName = 'trajectory_id'
597 elif objectType == 'object':
598 #tableName = 'objects_prototypes'
599 objectIdColumnName = 'object_id'
600 return tableName, objectIdColumnName
601
596 def savePrototypeAssignmentsToSqlite(filename, objectNumbers, objectType, labels, prototypes): 602 def savePrototypeAssignmentsToSqlite(filename, objectNumbers, objectType, labels, prototypes):
597 with sqlite3.connect(filename) as connection: 603 with sqlite3.connect(filename) as connection:
598 cursor = connection.cursor() 604 cursor = connection.cursor()
599 try: 605 try:
600 if objectType == 'feature': 606 tableName, objectIdColumnName = prototypeAssignmentNames(objectType)
601 tableName = 'features_prototypes'
602 objectIdColumnName = 'trajectory_id'
603 elif objectType == 'object':
604 tableName = 'objects_prototypes'
605 objectIdColumnName = 'object_id'
606 cursor.execute('CREATE TABLE IF NOT EXISTS '+tableName+' ('+objectIdColumnName+' INTEGER, prototype_filename VARCHAR, prototype_id INTEGER, trajectory_type VARCHAR CHECK (trajectory_type IN (\"feature\", \"object\")), PRIMARY KEY('+objectIdColumnName+', prototype_filename, prototype_id, trajectory_type))') 607 cursor.execute('CREATE TABLE IF NOT EXISTS '+tableName+' ('+objectIdColumnName+' INTEGER, prototype_filename VARCHAR, prototype_id INTEGER, trajectory_type VARCHAR CHECK (trajectory_type IN (\"feature\", \"object\")), PRIMARY KEY('+objectIdColumnName+', prototype_filename, prototype_id, trajectory_type))')
607 for objNum, label in zip(objectNumbers, labels): 608 for objNum, label in zip(objectNumbers, labels):
608 if label >=0: 609 if label >=0:
609 proto = prototypes[label] 610 proto = prototypes[label]
610 cursor.execute('INSERT INTO '+tableName+' VALUES(?,?,?,?)', (objNum, proto.getFilename(), proto.getNum(), proto.getTrajectoryType())) 611 cursor.execute('INSERT INTO '+tableName+' VALUES(?,?,?,?)', (objNum, proto.getFilename(), proto.getNum(), proto.getTrajectoryType()))
611 except sqlite3.OperationalError as error: 612 except sqlite3.OperationalError as error:
612 printDBError(error) 613 printDBError(error)
613 connection.commit() 614 connection.commit()
614 615
616 def loadPrototypeAssignmentsFromSqlite(filename, objectType):
617 with sqlite3.connect(filename) as connection:
618 cursor = connection.cursor()
619 try:
620 tableName, objectIdColumnName = prototypeAssignmentNames(objectType)
621 cursor.execute('SELECT * FROM '+tableName)
622 prototypeAssignments = {}
623 for row in cursor:
624 p = moving.Prototype(row[1], row[2], row[3])
625 if p in prototypeAssignments:
626 prototypeAssignments[p].append(row[0])
627 else:
628 prototypeAssignments[p] = [row[0]]
629 return prototypeAssignments
630 except sqlite3.OperationalError as error:
631 printDBError(error)
632
615 def loadPrototypesFromSqlite(filename, withTrajectories = True): 633 def loadPrototypesFromSqlite(filename, withTrajectories = True):
616 'Loads prototype ids and matchings (if stored)' 634 'Loads prototype ids and matchings (if stored)'
617 prototypes = [] 635 prototypes = []
618 with sqlite3.connect(filename) as connection: 636 with sqlite3.connect(filename) as connection:
619 cursor = connection.cursor() 637 cursor = connection.cursor()