comparison python/storage.py @ 204:966c2cd2bd9f

added code to load object trajectories (average of features)
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Tue, 06 Mar 2012 18:44:03 -0500
parents e2f31813ade6
children d9855499fc88
comparison
equal deleted inserted replaced
203:e2f31813ade6 204:966c2cd2bd9f
6 6
7 __metaclass__ = type 7 __metaclass__ = type
8 8
9 9
10 ngsimUserTypes = {'twowheels':1, 10 ngsimUserTypes = {'twowheels':1,
11 'car':2, 11 'car':2,
12 'truck':3} 12 'truck':3}
13 13
14 def loadTrajectoriesFromSqlite(filename, objectNumbers = -1): 14 def loadTrajectoriesFromSqlite(filename, trajectoryType, objectNumbers = -1):
15 '''Loads nObjects or the indices in objectNumbers from the database 15 '''Loads nObjects or the indices in objectNumbers from the database
16 TODO: load velocities''' 16 TODO: load velocities (replace table name 'positions' by 'velocities'
17 TODO: load features as well, other ways of averaging trajectories
18 '''
17 import sqlite3 19 import sqlite3
18 20
19 connection = sqlite3.connect(filename) # add test if it open 21 connection = sqlite3.connect(filename) # add test if it open
20 cursor = connection.cursor() 22 cursor = connection.cursor()
21 23
22 if type(objectNumbers) == int: 24 if trajectoryType == 'feature':
23 if objectNumbers == -1: 25 if type(objectNumbers) == int:
24 cursor.execute('SELECT * from positions order by trajectory_id, frame_number') 26 if objectNumbers == -1:
25 else: 27 cursor.execute('SELECT * from positions order by trajectory_id, frame_number')
26 cursor.execute('SELECT * from positions where trajectory_id between 0 and {0} order by trajectory_id, frame_number'.format(objectNumbers)) 28 else:
27 elif type(objectNumbers) == list: 29 cursor.execute('SELECT * from positions where trajectory_id between 0 and {0} order by trajectory_id, frame_number'.format(objectNumbers))
28 cursor.execute('SELECT * from positions where trajectory_id in (' 30 elif type(objectNumbers) == list:
29 +', '.join([str(n) for n in objectNumbers]) 31 cursor.execute('SELECT * from positions where trajectory_id in ('+', '.join([str(n) for n in objectNumbers])+') order by trajectory_id, frame_number')
30 +') order by trajectory_id, frame_number') 32 elif trajectoryType == 'object':
33 if type(objectNumbers) == int:
34 if objectNumbers == -1:
35 cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from positions P, objects_features OF where P.trajectory_id = OF.trajectory_id group by object_id, frame_number')
36 else:
37 cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from positions P, objects_features OF where P.trajectory_id = OF.trajectory_id and trajectory_id between 0 and {0} group by object_id, frame_number'.format(objectNumbers))
38 elif type(objectNumbers) == list:
39 cursor.execute('SELECT OF.object_id, P.frame_number, avg(P.x_coordinate), avg(P.y_coordinate) from positions P, objects_features OF where P.trajectory_id = OF.trajectory_id and trajectory_id in ('+', '.join([str(n) for n in objectNumbers])+') group by object_id, frame_number')
40 else:
41 print('no trajectory type was chosen')
31 42
32 objId = -1 43 objId = -1
33 obj = None 44 obj = None
34 objects = [] 45 objects = []
35 for row in cursor: 46 for row in cursor:
45 if obj: 56 if obj:
46 objects.append(obj) 57 objects.append(obj)
47 58
48 connection.close() 59 connection.close()
49 return objects 60 return objects
50
51 def loadObjectsFromSqlite(filename, objectNumbers = -1):
52 '''Loads objects as averages of feature trajectories
53 TODO: load features as well, other ways of averaging trajectories
54 need to provide table name(s) ?'''
55 # elect frame_number, avg(x_coordinate), avg(y_coordinate) from positions where trajectory_id in (select trajectory_id from objects_features where object_id=12) group by frame_number;
56
57
58
59 61
60 def loadTrajectoriesFromNgsimFile(filename, nObjects = -1, sequenceNum = -1): 62 def loadTrajectoriesFromNgsimFile(filename, nObjects = -1, sequenceNum = -1):
61 '''Reads data from the trajectory data provided by NGSIM project 63 '''Reads data from the trajectory data provided by NGSIM project
62 and returns the list of Feature objects''' 64 and returns the list of Feature objects'''
63 objects = [] 65 objects = []