--- a/python/storage.py	Wed Jul 25 18:22:18 2012 -0400
+++ b/python/storage.py	Wed Jul 25 22:06:51 2012 -0400
@@ -76,6 +76,23 @@
     connection.close()
     return matched_indexes
 
+def getTrajectoryIdQuery(objectNumbers, trajectoryType):
+    if trajectoryType == 'feature':
+        statementBeginning = ' where trajectory_id'
+    elif trajectoryType == 'object':
+        statementBeginning =  ' and OF.object_id'
+    else:
+        print('no trajectory type was chosen')
+
+    if type(objectNumbers) == int:
+        if objectNumbers == -1:
+            query = ''
+        else:
+            query = statementBeginning+' between 0 and {0}'.format(objectNumbers)
+    elif type(objectNumbers) == list:
+        query = statementBeginning+' in ('+', '.join([str(n) for n in objectNumbers])+')'
+    return query
+
 def loadTrajectoriesFromTable(connection, tableName, trajectoryType, objectNumbers = -1):
     '''Loads trajectories (in the general sense) from the given table
     can be positions or velocities
@@ -87,21 +104,11 @@
 
     try:
         if trajectoryType == 'feature':
-            if type(objectNumbers) == int:
-                if objectNumbers == -1:
-                    cursor.execute('SELECT * from '+tableName+' order by trajectory_id, frame_number')
-                else:
-                    cursor.execute('SELECT * from {0} where trajectory_id between 0 and {1} order by trajectory_id, frame_number'.format(tableName, objectNumbers))
-            elif type(objectNumbers) == list:
-                cursor.execute('SELECT * from '+tableName+' where trajectory_id in ('+', '.join([str(n) for n in objectNumbers])+') order by trajectory_id, frame_number')
+            trajectoryIdQuery = getTrajectoryIdQuery(objectNumbers, trajectoryType)
+            cursor.execute('SELECT * from '+tableName+trajectoryIdQuery+' order by trajectory_id, frame_number')
         elif trajectoryType == 'object':
-            if type(objectNumbers) == int:
-                if objectNumbers == -1:
-                    cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from '+tableName+' P, objects_features OF where P.trajectory_id = OF.trajectory_id group by object_id, frame_number')
-                else:
-                    cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from {0} P, objects_features OF where P.trajectory_id = OF.trajectory_id and OF.object_id between 0 and {1} group by object_id, frame_number'.format(tableName, objectNumbers))
-            elif type(objectNumbers) == list:
-                cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from '+tableName+' P, objects_features OF where P.trajectory_id = OF.trajectory_id and OF.object_id in ('+', '.join([str(n) for n in objectNumbers])+') group by object_id, frame_number')
+            objectIdQuery = getTrajectoryIdQuery(objectNumbers, trajectoryType)
+            cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from '+tableName+' P, objects_features OF where P.trajectory_id = OF.trajectory_id '+objectIdQuery+' group by object_id, frame_number')
         else:
             print('no trajectory type was chosen')
     except sqlite3.OperationalError as err:
@@ -145,6 +152,28 @@
             else:
                 print('Could not match positions {0} with velocities {1}'.format(o.num, v.num))
 
+    if trajectoryType == 'object':
+        cursor = connection.cursor()
+        try:
+            objectIdQuery = getTrajectoryIdQuery(objectNumbers, trajectoryType)
+            cursor.execute('SELECT P.trajectory_id, OF.object_id from positions P, objects_features OF where P.trajectory_id = OF.trajectory_id '+objectIdQuery+' group by P.trajectory_id order by OF.object_id')
+
+            # attribute feature numbers to objects
+            objId = -1
+            featureNumbers = {}
+            for row in cursor:
+                if row[1] != objId:
+                    objId = row[1]
+                    featureNumbers[objId] = [row[0]]
+                else:
+                    featureNumbers[objId].append(row[0])
+                    
+            for obj in objects:
+                obj.featureNumbers = featureNumbers[obj.num]
+        except sqlite3.OperationalError as err:
+            print('DB Error: {0}'.format(err))
+            return []
+
     connection.close()
     return objects