Mercurial > hg > nsaunier > traffic-intelligence

comparison python/storage.py @ 231:249d65ff6c35

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
merged modifications for windows
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Mon, 02 Jul 2012 23:49:39 -0400
parents b5772df11b37
children 584613399513
comparison
equal deleted inserted replaced
230:bc4ea09b1743 231:249d65ff6c35
1 #! /usr/bin/env python 1 #! /usr/bin/env python
2 # -*- coding: utf-8 -*-
2 '''Various utilities to save and load data''' 3 '''Various utilities to save and load data'''
3 4
4 import utils 5 import utils
5 import moving 6 import moving
6 7
8 9
9 10
10 ngsimUserTypes = {'twowheels':1, 11 ngsimUserTypes = {'twowheels':1,
11 'car':2, 12 'car':2,
12 'truck':3} 13 'truck':3}
14
15 def writeTrajectoriesToSqlite(objects, outFile, trajectoryType, objectNumbers = -1):
16 """
17 This function writers trajectories to a specified sqlite file
18 @param[in] objects -> a list of trajectories
19 @param[in] trajectoryType -
20 @param[out] outFile -> the .sqlite file containting the written objects
21 @param[in] objectNumber : number of objects loaded
22 """
23
24 import sqlite3
25 connection = sqlite3.connect(outFile)
26 cursor = connection.cursor()
27
28 schema = "CREATE TABLE \"positions\"(trajectory_id INTEGER,frame_number INTEGER, x_coordinate REAL, y_coordinate REAL, PRIMARY KEY(trajectory_id, frame_number))"
29 cursor.execute(schema)
30
31 trajectory_id = 0
32 frame_number = 0
33 if trajectoryType == 'feature':
34 if type(objectNumbers) == int and objectNumbers == -1:
35 for trajectory in objects:
36 trajectory_id += 1
37 frame_number = 0
38 for position in trajectory.getPositions():
39 frame_number += 1
40 query = "insert into positions (trajectory_id, frame_number, x_coordinate, y_coordinate) values (?,?,?,?)"
41 cursor.execute(query,(trajectory_id,frame_number,position.x,position.y))
42
43 connection.commit()
44 connection.close()
45
46 def loadPrototypeMatchIndexesFromSqlite(filename):
47 """
48 This function loads the prototypes table in the database of name <filename>.
49 It returns a list of tuples representing matching ids : [(prototype_id, matched_trajectory_id),...]
50 """
51 matched_indexes = []
52
53 import sqlite3
54 connection = sqlite3.connect(filename)
55 cursor = connection.cursor()
56
57 try:
58 cursor.execute('SELECT * from prototypes order by prototype_id, trajectory_id_matched')
59 except sqlite3.OperationalError as err:
60 print('DB Error: {0}'.format(err))
61 return []
62
63 for row in cursor:
64 matched_indexes.append((row[0],row[1]))
65
66 connection.close()
67 return matched_indexes
13 68
14 def loadTrajectoriesFromSqlite(filename, trajectoryType, objectNumbers = -1): 69 def loadTrajectoriesFromSqlite(filename, trajectoryType, objectNumbers = -1):
15 '''Loads nObjects or the indices in objectNumbers from the database 70 '''Loads nObjects or the indices in objectNumbers from the database
16 TODO: load velocities (replace table name 'positions' by 'velocities' 71 TODO: load velocities (replace table name 'positions' by 'velocities'
17 TODO: load features as well, other ways of averaging trajectories 72 TODO: load features as well, other ways of averaging trajectories
18 ''' 73 '''
19 import sqlite3 74 import sqlite3
20 75
21 connection = sqlite3.connect(filename) # add test if it open 76 connection = sqlite3.connect(filename) # add test if it open
22 cursor = connection.cursor() 77 cursor = connection.cursor()
23 78
24 if trajectoryType == 'feature': 79 try:
25 if type(objectNumbers) == int: 80 if trajectoryType == 'feature':
26 if objectNumbers == -1: 81 if type(objectNumbers) == int:
27 cursor.execute('SELECT * from positions order by trajectory_id, frame_number') 82 if objectNumbers == -1:
28 else: 83 cursor.execute('SELECT * from positions order by trajectory_id, frame_number')
29 cursor.execute('SELECT * from positions where trajectory_id between 0 and {0} order by trajectory_id, frame_number'.format(objectNumbers)) 84 else:
30 elif type(objectNumbers) == list: 85 cursor.execute('SELECT * from positions where trajectory_id between 0 and {0} order by trajectory_id, frame_number'.format(objectNumbers))
31 cursor.execute('SELECT * from positions where trajectory_id in ('+', '.join([str(n) for n in objectNumbers])+') order by trajectory_id, frame_number') 86 elif type(objectNumbers) == list:
32 elif trajectoryType == 'object': 87 cursor.execute('SELECT * from positions where trajectory_id in ('+', '.join([str(n) for n in objectNumbers])+') order by trajectory_id, frame_number')
33 if type(objectNumbers) == int: 88 elif trajectoryType == 'object':
34 if objectNumbers == -1: 89 if type(objectNumbers) == int:
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') 90 if objectNumbers == -1:
36 else: 91 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')
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)) 92 else:
38 elif type(objectNumbers) == list: 93 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 OF.object_id between 0 and {0} group by object_id, frame_number'.format(objectNumbers))
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') 94 elif type(objectNumbers) == list:
40 else: 95 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 OF.object_id in ('+', '.join([str(n) for n in objectNumbers])+') group by object_id, frame_number')
41 print('no trajectory type was chosen') 96 else:
42 97 print('no trajectory type was chosen')
98 except sqlite3.OperationalError as err:
99 print('DB Error: {0}'.format(err))
100 return []
101
43 objId = -1 102 objId = -1
44 obj = None 103 obj = None
45 objects = [] 104 objects = []
46 for row in cursor: 105 for row in cursor:
47 if row[0] != objId: 106 if row[0] != objId:
141 200
142 print nObjectsPerType 201 print nObjectsPerType
143 202
144 out.close() 203 out.close()
145 204
146 205 if __name__ == "__main__":
147 206 import doctest
148 207 import unittest
149 # if __name__ == "__main__": 208 suite = doctest.DocFileSuite('tests/storage.txt')
150 # import doctest 209 unittest.TextTestRunner().run(suite)
151 # import unittest
152 # suite = doctest.DocFileSuite('tests/ubc_utils.txt')
153 # unittest.TextTestRunner().run(suite)
154 # #doctest.testmod() 210 # #doctest.testmod()
155 # #doctest.testfile("example.txt") 211 # #doctest.testfile("example.txt")