diff python/storage.py @ 921:630934595871

work in progress with prototype class
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 05 Jul 2017 18:01:43 -0400
parents 499154254f37
children acb5379c5fd7
line wrap: on
line diff
--- a/python/storage.py	Wed Jul 05 16:30:04 2017 -0400
+++ b/python/storage.py	Wed Jul 05 18:01:43 2017 -0400
@@ -592,29 +592,14 @@
 # saving and loading for scene interpretation: POIs and Prototypes
 #########################
 
-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
-    objects will be saved as features, with the centroid trajectory as if it is a feature)
-    nMatchings, if not None, is a list of the number of matches
-    dbFilenames, if not None, is a list of the DB filenames
-
-    The order of prototypeIndices, objects, nMatchings and dbFilenames should be consistent'''
+def savePrototypesToSqlite(filename, prototypes):
+    '''save the prototypes (a prototype is defined by a filename, a number and type'''
     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, PRIMARY KEY (id, dbfilename))')
-        for i, protoId in enumerate(prototypeIndices):
-            if nMatchings is not None:
-                n = nMatchings[i]
-            else:
-                n = 'NULL'
-            if dbFilenames is not None:
-                dbfn = dbFilenames[i]
-            else:
-                dbfn = filename
-            cursor.execute('INSERT INTO prototypes (id, dbfilename, trajectory_type, nmatchings) VALUES (?,?,?,?)', (protoId, dbfn, trajectoryType, n))
+        cursor.execute('CREATE TABLE IF NOT EXISTS prototypes (dbfilename VARCHAR, id INTEGER, trajectory_type VARCHAR CHECK (trajectory_type IN (\"feature\", \"object\")), nmatchings INTEGER, PRIMARY KEY (id, dbfilename))')
+        for p in prototypes:
+            cursor.execute('INSERT INTO prototypes (dbfilename, id, trajectory_type, nmatchings) VALUES (?,?,?,?)', (p.getFilename(), p.getNum(), p.getTrajectoryType(), p.getNMatchings()))
     except sqlite3.OperationalError as error:
         printDBError(error)
     connection.commit()
@@ -627,34 +612,31 @@
     'Loads prototype ids and matchings (if stored)'
     connection = sqlite3.connect(filename)
     cursor = connection.cursor()
-    prototypeIndices = []
-    dbFilenames = []
-    trajectoryTypes = []
-    nMatchings = []
+    prototypes = []
     objects = []
     try:
         cursor.execute('SELECT * FROM prototypes')
         for row in cursor:
-            prototypeIndices.append(row[0])
-            dbFilenames.append(row[1])
-            trajectoryTypes.append(row[2])
-            if row[3] is not None:
-                nMatchings.append(row[3])
+            prototypes.append(moving.Prototype(row[0], row[1], row[2], row[3]))
         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)
+            for p in prototypes:
+                p.setMovingObject(loadTrajectoriesFromSqlite(p.getFilename(), p.getTrajectoryType(), [p.getNum()])[0])
+            # loadingInformation = {} # complicated slightly optimized
+            # for p in prototypes:
+            #     dbfn = p.getFilename()
+            #     trajType = p.getTrajectoryType()
+            #     if (dbfn, trajType) in loadingInformation:
+            #         loadingInformation[(dbfn, trajType)].append(p)
+            #     else:
+            #         loadingInformation[(dbfn, trajType)] = [p]
+            # for k, v in loadingInformation.iteritems():
+            #     objects += loadTrajectoriesFromSqlite(k[0], k[1], [p.getNum() for p in v])
     except sqlite3.OperationalError as error:
         printDBError(error)
     connection.close()
-    if len(set(trajectoryTypes)) > 1:
-        print('Different types of prototypes in database ({}).'.format(set(trajectoryTypes)))
-    return prototypeIndices, dbFilenames, trajectoryTypes, nMatchings, objects
+    if len(set([p.getTrajectoryType() for p in prototypes])) > 1:
+        print('Different types of prototypes in database ({}).'.format(set([p.getTrajectoryType() for p in prototypes])))
+    return prototypes
 
 def savePOIs(filename, gmm, gmmType, gmmId):
     '''Saves a Gaussian mixture model (of class sklearn.mixture.GaussianMixture)