nsaunier/traffic-intelligence

update to add curvilinear positions

Commit 3d6ee243d5c0 · Nicolas Saunier · 2024-04-07 22:09 -0400

Changeset

3d6ee243d5c0097081294338ec4241d4e5ece7fc

View source at this commit

Comments

No comments yet.

Log in to comment

Diff

diff --git a/trafficintelligence/storage.py b/trafficintelligence/storage.py
--- a/trafficintelligence/storage.py
+++ b/trafficintelligence/storage.py
@@ -357,7 +357,15 @@
 
 def addCurvilinearTrajectoriesFromSqlite(filename, objects):
     '''Adds curvilinear positions (s_coordinate, y_coordinate, lane)
-    from a database to an existing MovingObject dict (indexed by each objects's num)'''
+    from a database to an existing MovingObject dict (indexed by each objects's num) or list'''
+    if type(objects) == list:
+        _objects = {obj.getNum(): obj for obj in objects}
+    elif type(objects) == dict:
+        _objects = objects
+    else:
+        print('objects should be a list of dictionary of objects indexed by their num')
+        return None
+    
     with sqlite3.connect(filename) as connection:
         cursor = connection.cursor()
 
@@ -372,11 +380,11 @@
         for row in cursor:
             if objNum != row[0]:
                 objNum = row[0]
-                if objNum in objects:
-                    objects[objNum].curvilinearPositions = moving.CurvilinearTrajectory()
+                if objNum in _objects:
+                    _objects[objNum].curvilinearPositions = moving.CurvilinearTrajectory()
                 else:
                     missingObjectNumbers.append(objNum)
-            if objNum in objects:
+            if objNum in _objects:
                 objects[objNum].curvilinearPositions.addPositionSYL(row[2],row[3],row[4])
         if len(missingObjectNumbers) > 0:
             print('List of missing objects to attach corresponding curvilinear trajectories: {}'.format(missingObjectNumbers))