Mercurial Hosting > traffic-intelligence
comparison scripts/polytracktopdtv.py @ 402:f29204e68aab
function to generate homography from PDTV Tsai format and script to generate trajectories from sqlite bounding boxes
author | Nicolas Saunier <nicolas.saunier@polymtl.ca> |
---|---|
date | Mon, 29 Jul 2013 19:45:43 -0400 |
parents | |
children | f3938bb6da7f |
comparison
equal
deleted
inserted
replaced
401:b829ebdc18e6 | 402:f29204e68aab |
---|---|
1 #! /usr/bin/env python | |
2 | |
3 from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State | |
4 import sys, os, datetime | |
5 import shutil | |
6 import sqlite3 | |
7 import zipfile | |
8 import utils | |
9 import cvutils | |
10 import cv2 | |
11 | |
12 | |
13 def zipFolder(inputFolder, outputFile): | |
14 '''Method to compress the content of the inputFolder in the outputFile''' | |
15 zip = zipfile.ZipFile(outputFile, 'w') | |
16 for root, dirs, files in os.walk(inputFolder): | |
17 for file in files: | |
18 zip.write(root+file, file) | |
19 zip.close() | |
20 | |
21 | |
22 | |
23 | |
24 def getTypeDict(cursor): | |
25 '''Return a dictionnary with integer key and associated type string | |
26 i.e.: "0" -> "unknown" | |
27 "1" -> "car" | |
28 "2" -> "pedestrians" | |
29 "3" -> "motorcycle" | |
30 "4" -> "bicycle" | |
31 "5" -> "bus" | |
32 "6" -> "truck" | |
33 ... and other type if the objects_type table is defined in SQLite''' | |
34 typeDict = dict() | |
35 cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='objects_type'") | |
36 data = cursor.fetchone() | |
37 | |
38 if(data == None): | |
39 typeDict["0"] = "unknown" | |
40 typeDict["1"] = "car" | |
41 typeDict["2"] = "pedestrians" | |
42 typeDict["3"] = "motorcycle" | |
43 typeDict["4"] = "bicycle" | |
44 typeDict["5"] = "bus" | |
45 typeDict["6"] = "truck" | |
46 | |
47 else: | |
48 cursor.execute("SELECT road_user_type, type_string FROM objects_type") | |
49 for row in cursor: | |
50 typeDict[row[0]]= row[1] | |
51 | |
52 return typeDict | |
53 | |
54 def extractFrames(videoFile, framePath, fps, time, startFrame = 0, endFrame = None): | |
55 '''Method to extract all the frames of the video''' | |
56 print('Extracting frame') | |
57 deltaTimestamp = 1000.0/float(fps); | |
58 time+=datetime.timedelta(microseconds=startFrame*deltaTimestamp*1000) | |
59 | |
60 inc = 1000 #How many frame we fetch in the video at a time | |
61 | |
62 if endFrame != None: | |
63 delta = endFrame-startFrame | |
64 if delta < inc: | |
65 inc = delta | |
66 | |
67 currentIdx = startFrame | |
68 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc) | |
69 | |
70 while len(frameList) == inc and inc > 0: | |
71 | |
72 for f in frameList: | |
73 cv2.imwrite(os.path.join(framePath,time.strftime("%Y%m%d-%H%M%S.%f")[:-3]+'.jpg'), f) | |
74 time += datetime.timedelta(microseconds=deltaTimestamp*1000) | |
75 currentIdx = currentIdx + inc | |
76 | |
77 if endFrame != None: | |
78 delta = endFrame-currentIdx | |
79 if delta < inc: | |
80 inc = delta | |
81 if inc: | |
82 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc) | |
83 print('Extracting frame ' + str(currentIdx)) | |
84 return len(frameList) > 0 | |
85 | |
86 | |
87 | |
88 def convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = None, videoFolderExist = False, startFrame = 0, endFrame = None): | |
89 ''' | |
90 Method to convert database from polytrack to PDTV: | |
91 workDirname is the current working directory | |
92 sceneFilename is the path to the .cfg file | |
93 sectionName is the name of the section we want to process in this file | |
94 videoFolderExist specifiy if we want to reextract the video frame or if they already exist at workdir/videoframes/ | |
95 startFrame is the first frame we want to extract | |
96 endFrame is the last frame we want to extract (or None if we want to extract everything) | |
97 ''' | |
98 | |
99 scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename)) | |
100 if databaseFilename != None: | |
101 inputDb = os.path.join(workDirname, databaseFilename) | |
102 else: | |
103 inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename) | |
104 videoFile = os.path.join(workDirname, scene[sectionName].videoFilename) | |
105 | |
106 error = False | |
107 | |
108 if not os.path.exists(inputDb): | |
109 print('Database path is not valid: ' + inputDb) | |
110 error = True | |
111 if not os.path.exists(videoFile): | |
112 print('Video path is not valid: ' + videoFile) | |
113 error = True | |
114 | |
115 time = scene[sectionName].date | |
116 | |
117 | |
118 videoFolderPath = os.path.join(workDirname, "videoframes/") | |
119 fileName = scene[sectionName].sitename | |
120 | |
121 fps = cvutils.getFPS(videoFile) | |
122 print('Video should run at ' + str(fps) + ' fps') | |
123 deltaTimestamp = 1000.0/float(fps); | |
124 if videoFolderExist == False: | |
125 if os.path.exists(videoFolderPath): | |
126 shutil.rmtree(videoFolderPath) | |
127 utils.mkdir(videoFolderPath) | |
128 if extractFrames(videoFile, videoFolderPath, fps, time, startFrame, endFrame) == 0: | |
129 print("Error. Frame were not extracted") | |
130 error = True | |
131 | |
132 if not error: | |
133 masterTimestamp = 0.0 | |
134 masterTimestampList = list() | |
135 video_timestringsList = list() | |
136 frameNumberToMasterTimestamp = dict() | |
137 for r,d,f in os.walk(videoFolderPath): | |
138 i = startFrame | |
139 for files in f: | |
140 name, ext = os.path.splitext(files) | |
141 video_timestringsList.append(name) | |
142 masterTimestampList.append(masterTimestamp) | |
143 frameNumberToMasterTimestamp[i] = masterTimestamp | |
144 i = i +1 | |
145 masterTimestamp = masterTimestamp+deltaTimestamp | |
146 | |
147 inputZipVideoName = fileName+".zip" | |
148 print('Zipping files...') | |
149 if not os.path.exists(inputZipVideoName) or not videoFolderExist: | |
150 zipFolder(videoFolderPath, inputZipVideoName) | |
151 print('Zipping files...Done.') | |
152 #We generate the structure for ZipVideo | |
153 zipVideo = ZipVideo(video_zip_file=inputZipVideoName, | |
154 time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file='calib.json') | |
155 zipVideo.save(os.path.join(workDirname, 'video.json')) | |
156 | |
157 print('ZipVideo saved') | |
158 obj = SyncedVideos(master_timestamps = [masterTimestamp], | |
159 video_timestrings = [video_timestringsList], | |
160 video_files = ['video.json'], | |
161 fps=fps) | |
162 obj.save(os.path.join(workDirname, 'syncedvideo.json')) | |
163 | |
164 print('SyncedVideos saved') | |
165 | |
166 print('Opening db') | |
167 connection = sqlite3.connect(inputDb) | |
168 cursor = connection.cursor() | |
169 | |
170 | |
171 #Tracket database | |
172 trackset = TrackSet(synced_file = ['syncedvideo.json']) | |
173 | |
174 | |
175 #1) We build the type index dictionnary | |
176 typeDict = getTypeDict(cursor) | |
177 | |
178 idToTrackDict = dict() | |
179 #2) We read the object database | |
180 cursor.execute("SELECT object_id, road_user_type FROM objects") | |
181 for row in cursor: | |
182 objectId = row[0] | |
183 objectType = row[1] | |
184 t = Track(type=typeDict.get(objectType, "unknown")) | |
185 idToTrackDict[objectId] = t; | |
186 trackset.append(t) | |
187 print('Reading boundingbox table') | |
188 #3) We read the bounding box table | |
189 cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='bounding_boxes'") | |
190 data = cursor.fetchone() | |
191 if data == None: | |
192 print('No bounding box table. Maybe it was not generated ?') | |
193 else: | |
194 cursor.execute("SELECT object_id, frame_number, x_top_left, y_top_left, x_bottom_right, y_bottom_right FROM bounding_boxes") | |
195 for row in cursor: | |
196 objectId = row[0] | |
197 frameNumber = row[1] | |
198 x_top_left = row[2] | |
199 y_top_left = row[3] | |
200 x_bottom_right = row[4] | |
201 y_bottom_right = row[5] | |
202 | |
203 idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.], | |
204 master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)], | |
205 bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]])) | |
206 print('Saving db in json') | |
207 trackset.save(os.path.join(workDirname, 'roaduser.json')) | |
208 connection.close() | |
209 print('Done.') | |
210 | |
211 | |
212 | |
213 | |
214 if __name__ == "__main__": | |
215 workDirname = "./" | |
216 sceneFilename = "scene.cfg" | |
217 sectionName = "amherst-11" | |
218 # videoFrameAlreadyExtracted Set this to true if you have the frame already extracted in the folder workDirname/videoframes/ | |
219 convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = 'amherst-11_gt.sqlite', videoFolderExist = True, startFrame = 0, endFrame = 4000) | |
220 | |
221 |