diff trafficintelligence/storage.py @ 1041:fc7c0f38e8a6

added nObjects to MovingObject, with loading/saving
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 04 Jul 2018 16:06:23 -0400
parents 20799ac9524e
children 75a6ad604cc5 1748c02f9ac3
line wrap: on
line diff
--- a/trafficintelligence/storage.py	Wed Jul 04 15:37:04 2018 -0400
+++ b/trafficintelligence/storage.py	Wed Jul 04 16:06:23 2018 -0400
@@ -229,16 +229,23 @@
 
     return objects
 
-def loadUserTypesFromTable(cursor, objectNumbers):
+def loadObjectAttributesFromTable(cursor, objectNumbers, loadNObjects = False):
     objectCriteria = getObjectCriteria(objectNumbers)
-    queryStatement = 'SELECT object_id, road_user_type FROM objects'
+    queryStatement = 'SELECT object_id, road_user_type'
+    if loadNObjects:
+        queryStatement += ', n_objects'
+    queryStatement += ' FROM objects'
     if objectNumbers is not None:
         queryStatement += ' WHERE object_id '+objectCriteria
     cursor.execute(queryStatement)
-    userTypes = {}
-    for row in cursor:
-        userTypes[row[0]] = row[1]
-    return userTypes
+    attributes = {}
+    if loadNObjects:
+        for row in cursor:
+            attributes[row[0]] = row[1:]
+    else:
+        for row in cursor:
+            attributes[row[0]] = row[1]
+    return attributes
 
 def loadTrajectoriesFromSqlite(filename, trajectoryType, objectNumbers = None, withFeatures = False, timeStep = None, maxNObjectFeatures = 1):
     '''Loads the trajectories (in the general sense, 
@@ -290,9 +297,11 @@
                     obj.featureNumbers = featureNumbers[obj.getNum()]
 
                 # load userType
-                userTypes = loadUserTypesFromTable(cursor, objectNumbers)
+                attributes = loadObjectAttributesFromTable(cursor, objectNumbers, True)
                 for obj in objects:
-                    obj.userType = userTypes[obj.getNum()]
+                    userType, nObjects = attributes[obj.getNum()]
+                    obj.setUserType(userType)
+                    obj.setNObjects(nObjects)
 
                 if withFeatures:
                     nFeatures = 0
@@ -452,7 +461,7 @@
         if objectType == 'bb':
             topCorners = loadTrajectoriesFromTable(connection, 'bounding_boxes', 'bbtop', objectNumbers, timeStep)
             bottomCorners = loadTrajectoriesFromTable(connection, 'bounding_boxes', 'bbbottom', objectNumbers, timeStep)
-            userTypes = loadUserTypesFromTable(connection.cursor(), objectNumbers) # string format is same as object
+            userTypes = loadObjectAttributesFromTable(connection.cursor(), objectNumbers) # string format is same as object
 
             for t, b in zip(topCorners, bottomCorners):
                 num = t.getNum()
@@ -1033,7 +1042,7 @@
         for row_index, row in instants.iterrows():
             objNum = int(row['NO'])
             tmp = data[data['NO'] == objNum]
-            objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(row['first'], row['last']))
+            objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(row['first'], row['last']), nObjects = 1)
             # positions should be rounded to nDecimals decimals only
             objects[objNum].curvilinearPositions = moving.CurvilinearTrajectory(S = npround(tmp['POS'].tolist(), nDecimals), Y = npround(tmp['POSLAT'].tolist(), nDecimals), lanes = tmp['LANE'].tolist())
             if objectNumbers is not None and objectNumbers > 0 and len(objects) >= objectNumbers:
@@ -1052,7 +1061,7 @@
                 if objNum not in objects:
                     if warmUpLastInstant is None or instant >= warmUpLastInstant:
                         if objectNumbers is None or len(objects) < objectNumbers:
-                            objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(instant, instant))
+                            objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(instant, instant), nObjects = 1)
                             objects[objNum].curvilinearPositions = moving.CurvilinearTrajectory()
                 if (warmUpLastInstant is None or instant >= warmUpLastInstant) and objNum in objects:
                     objects[objNum].timeInterval.last = instant
@@ -1076,7 +1085,7 @@
                         if objNum not in objects:
                             if warmUpLastInstant is None or instant >= warmUpLastInstant:
                                 if objectNumbers is None or len(objects) < objectNumbers:
-                                    objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(instant, instant))
+                                    objects[objNum] = moving.MovingObject(num = objNum, timeInterval = moving.TimeInterval(instant, instant), nObjects = 1)
                                     objects[objNum].curvilinearPositions = moving.CurvilinearTrajectory()
                         if (warmUpLastInstant is None or instant >= warmUpLastInstant) and objNum in objects:
                             objects[objNum].timeInterval.last = instant
@@ -1181,7 +1190,7 @@
         obj = moving.MovingObject(num = int(numbers[0]), 
                                   timeInterval = moving.TimeInterval(firstFrameNum, lastFrameNum), 
                                   positions = moving.Trajectory([[float(numbers[6])],[float(numbers[7])]]), 
-                                  userType = int(numbers[10]))
+                                  userType = int(numbers[10]), nObjects = 1)
         obj.userType = int(numbers[10])
         obj.laneNums = [int(numbers[13])]
         obj.precedingVehicles = [int(numbers[14])] # lead vehicle (before)