diff python/storage.py @ 491:343cfd185ca6

minor changes and reaarrangements
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 16 Apr 2014 17:43:53 -0400
parents f6415f012640
children 935430b1d408
line wrap: on
line diff
--- a/python/storage.py	Fri Apr 11 17:47:55 2014 -0400
+++ b/python/storage.py	Wed Apr 16 17:43:53 2014 -0400
@@ -17,6 +17,20 @@
 # Sqlite
 #########################
 
+# utils
+def printDBError(error):
+    print('DB Error: {}'.format(error))
+
+def dropTables(connection, tableNames):
+    'deletes the table with names in tableNames'
+    try:
+        cursor = connection.cursor()
+        for tableName in tableNames:
+            cursor.execute('DROP TABLE IF EXISTS '+tableName)
+    except sqlite3.OperationalError as error:
+        printDBError(error)
+
+# IO to sqlite
 def writeTrajectoriesToSqlite(objects, outFilename, trajectoryType, objectNumbers = -1):
     """
     This function writers trajectories to a specified sqlite file
@@ -69,7 +83,7 @@
     try:
         cursor.execute('SELECT * from prototypes order by prototype_id, trajectory_id_matched')
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
         return []
 
     for row in cursor:
@@ -116,7 +130,7 @@
         else:
             print('no trajectory type was chosen')
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
         return []
 
     objId = -1
@@ -185,7 +199,7 @@
                 obj.userType = userTypes[obj.getNum()]
              
         except sqlite3.OperationalError as error:
-            utils.printDBError(error)
+            printDBError(error)
             return []
 
     connection.close()
@@ -193,16 +207,20 @@
 
 def removeFromSqlite(filename, dataType):
     'Removes some tables in the filename depending on type of data'
-    connection = sqlite3.connect(filename)
-    if dataType == 'object':
-        utils.dropTables(connection, ['objects', 'objects_features'])
-    elif dataType == 'interaction':
-        utils.dropTables(connection, ['interactions', 'indicators'])
-    elif dataType == 'bb':
-        utils.dropTables(connection, ['bounding_boxes'])
+    import os
+    if os.path.isfile(filename):
+        connection = sqlite3.connect(filename)
+        if dataType == 'object':
+            dropTables(connection, ['objects', 'objects_features'])
+        elif dataType == 'interaction':
+            dropTables(connection, ['interactions', 'indicators'])
+        elif dataType == 'bb':
+            dropTables(connection, ['bounding_boxes'])
+        else:
+            print('Unknown data type {} to delete from database'.format(dataType))
+        connection.close()
     else:
-        print('Unknown data type {} to delete from database'.format(dataType))
-    connection.close()
+        print('{} does not exist'.format(filename))
 
 def createInteractionTable(cursor):
     cursor.execute('CREATE TABLE IF NOT EXISTS interactions (id INTEGER PRIMARY KEY, object_id1 INTEGER, object_id2 INTEGER, first_frame_number INTEGER, last_frame_number INTEGER, FOREIGN KEY(object_id1) REFERENCES objects(id), FOREIGN KEY(object_id2) REFERENCES objects(id))')
@@ -225,7 +243,7 @@
         for inter in interactions:
             saveInteraction(cursor, inter)
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
     connection.commit()
     connection.close()
 
@@ -248,7 +266,7 @@
                 if indicator != None:
                     saveIndicator(cursor, inter.getNum(), indicator)
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
     connection.commit()
     connection.close()
 
@@ -287,7 +305,7 @@
             interactions.append(events.Interaction(interactionNum, moving.TimeInterval(row[3],row[4]), roadUserNumbers[0], roadUserNumbers[1]))
             interactions[-1].indicators = tmpIndicators
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
         return []
     connection.close()
     return interactions
@@ -306,7 +324,7 @@
               '(SELECT OF.object_id, P.frame_number, P.x_coordinate as x, P.y_coordinate as y, P.x_coordinate*{}+P.y_coordinate*{}+{} as w from positions P, objects_features OF where P.trajectory_id = OF.trajectory_id)) '.format(invHomography[0,0], invHomography[0,1], invHomography[0,2], invHomography[1,0], invHomography[1,1], invHomography[1,2], invHomography[2,0], invHomography[2,1], invHomography[2,2])+
               'GROUP BY object_id, frame_number')
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
     connection.commit()
     connection.close()
 
@@ -324,7 +342,7 @@
                 #if row[0] != objId:
                 boundingBoxes.setdefault(row[1], []).append([moving.Point(row[2], row[3]), moving.Point(row[4], row[5])])
     except sqlite3.OperationalError as error:
-        utils.printDBError(error)
+        printDBError(error)
         return boundingBoxes
     connection.close()
     return boundingBoxes