Mercurial Hosting > traffic-intelligence
comparison trafficintelligence/storage.py @ 1037:6a6c37eb3a74
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() |