diff python/storage.py @ 805:180b6b0231c0

added saving/loading points of interests
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Thu, 09 Jun 2016 15:36:21 -0400
parents 30bd0f2223b7
children 21f10332c72b
line wrap: on
line diff
--- a/python/storage.py	Tue May 31 17:07:23 2016 -0400
+++ b/python/storage.py	Thu Jun 09 15:36:21 2016 -0400
@@ -384,6 +384,8 @@
             dropTables(connection, ['interactions', 'indicators'])
         elif dataType == 'bb':
             dropTables(connection, ['bounding_boxes'])
+        elif dataType == 'pois':
+            dropTables(connection, ['gaussians2d'])
         else:
             print('Unknown data type {} to delete from database'.format(dataType))
         connection.close()
@@ -506,6 +508,71 @@
     return boundingBoxes
 
 #########################
+# saving and loading for scene interpretation
+#########################
+
+def savePOIs(filename, gmm, gmmType, gmmId):
+    '''Saves a Gaussian mixture model (of class sklearn.mixture.GMM)
+    gmmType is a type of GMM, learnt either from beginnings or ends of trajectories'''
+    connection = sqlite3.connect(filename)
+    cursor = connection.cursor()
+    if gmmType not in ['beginning', 'end']:
+        print('Unknown POI type {}. Exiting'.format(gmmType))
+        import sys
+        sys.exit()
+    try:
+        cursor.execute('CREATE TABLE IF NOT EXISTS gaussians2d (id INTEGER, type VARCHAR, x_center REAL, y_center REAL, covar00 REAL, covar01 REAL, covar10 REAL, covar11 REAL, covariance_type VARCHAR, weight, mixture_id INTEGER, PRIMARY KEY(id, mixture_id))')
+        for i in xrange(gmm.n_components):
+            cursor.execute('INSERT INTO gaussians2d VALUES({}, \'{}\', {}, {}, {}, {}, {}, {}, \'{}\', {}, {})'.format(i, gmmType, gmm.means_[i][0], gmm.means_[i][1], gmm.covars_[i][0,0], gmm.covars_[i][0,1], gmm.covars_[i][1,0], gmm.covars_[i][1,1], gmm.covariance_type, gmm.weights_[i], gmmId))
+        connection.commit()
+    except sqlite3.OperationalError as error:
+        printDBError(error)
+    connection.close()
+
+def loadPOIs(filename):
+    'Loads all 2D Gaussians in the database'
+    from sklearn import mixture # todo if not avalaible, load data in duck-typed class with same fields
+    connection = sqlite3.connect(filename)
+    cursor = connection.cursor()
+    pois = []
+    try:
+        cursor.execute('SELECT * from gaussians2d')
+        gmmId = None
+        gmm = []
+        for row in cursor:
+            if gmmId is None or row[10] != gmmId:
+                if len(gmm) > 0:
+                    tmp = mixture.GMM(len(gmm), covarianceType)
+                    tmp.means_ = array([gaussian['mean'] for gaussian in gmm])
+                    tmp.covars_ = array([gaussian['covar'] for gaussian in gmm])
+                    tmp.weights_ = array([gaussian['weight'] for gaussian in gmm])
+                    tmp.gmmTypes = [gaussian['type'] for gaussian in gmm]
+                    pois.append(tmp)
+                gaussian = {'type': row[1],
+                            'mean': row[2:4],
+                            'covar': array(row[4:8]).reshape(2,2),
+                            'weight': row[9]}
+                gmm = [gaussian]
+                covarianceType = row[8]
+                gmmId = row[10]
+            else:
+                gmm.append({'type': row[1],
+                            'mean': row[2:4],
+                            'covar': array(row[4:8]).reshape(2,2),
+                            'weight': row[9]})
+        if len(gmm) > 0:
+            tmp = mixture.GMM(len(gmm), covarianceType)
+            tmp.means_ = array([gaussian['mean'] for gaussian in gmm])
+            tmp.covars_ = array([gaussian['covar'] for gaussian in gmm])
+            tmp.weights_ = array([gaussian['weight'] for gaussian in gmm])
+            tmp.gmmTypes = [gaussian['type'] for gaussian in gmm]
+            pois.append(tmp)
+    except sqlite3.OperationalError as error:
+        printDBError(error)
+    connection.close()
+    return pois
+    
+#########################
 # saving and loading for scene interpretation (Mohamed Gomaa Mohamed's PhD)
 #########################