Mercurial Hosting > traffic-intelligence
comparison scripts/polytracktopdtv.py @ 614:5e09583275a4
Merged Nicolas/trafficintelligence into default
author | Mohamed Gomaa <eng.m.gom3a@gmail.com> |
---|---|
date | Fri, 05 Dec 2014 12:13:53 -0500 |
parents | 4dd68b0765b1 |
children | 3058e00887bc |
comparison
equal
deleted
inserted
replaced
598:11f96bd08552 | 614:5e09583275a4 |
---|---|
1 #! /usr/bin/env python | |
2 | |
3 from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State | |
4 import sys, os, datetime, argparse | |
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, firstFrameNum = 0, lastFrameNum = 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=firstFrameNum*deltaTimestamp*1000) | |
59 | |
60 inc = 1000 #How many frame we fetch in the video at a time | |
61 | |
62 if lastFrameNum != None: | |
63 delta = lastFrameNum-firstFrameNum | |
64 if delta < inc: | |
65 inc = delta | |
66 | |
67 currentIdx = firstFrameNum | |
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 lastFrameNum != None: | |
78 delta = lastFrameNum-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, sectionName, sceneFilename = None, databaseFilename = None, videoFilename = None, videoFolderExist = False, firstFrameNum = 0, lastFrameNum = None, cameraCalibrationFilename = None, outputFileName = 'roaduser.json'): | |
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 firstFrameNum is the first frame we want to extract | |
96 lastFrameNum is the last frame we want to extract (or None if we want to extract everything) | |
97 ''' | |
98 error = False | |
99 if sceneFilename != None: | |
100 scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename)) | |
101 time = scene[sectionName].date | |
102 inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename) | |
103 videoFile = os.path.join(workDirname, scene[sectionName].videoFilename) | |
104 | |
105 if databaseFilename != None: | |
106 inputDb = os.path.join(workDirname, databaseFilename) | |
107 if videoFilename != None: | |
108 videoFile = os.path.join(workDirname, videoFilename) | |
109 # elif videoFolderExist == False: | |
110 # print('No video path specified') | |
111 # error = True | |
112 | |
113 videoFolderPath = os.path.join(workDirname, "videoframes/") | |
114 fileName = sectionName | |
115 | |
116 if videoFile != None: | |
117 fps = cvutils.getFPS(videoFile) | |
118 print('Video should run at ' + str(fps) + ' fps') | |
119 deltaTimestamp = 1000.0/float(fps); | |
120 if videoFolderExist == False: | |
121 if os.path.exists(videoFolderPath): | |
122 shutil.rmtree(videoFolderPath) | |
123 utils.mkdir(videoFolderPath) | |
124 if extractFrames(videoFile, videoFolderPath, fps, time, firstFrameNum, lastFrameNum) == 0: | |
125 print("Error. Frame were not extracted") | |
126 error = True | |
127 | |
128 | |
129 if not error: | |
130 masterTimestamp = 0.0 | |
131 masterTimestampList = list() | |
132 video_timestringsList = list() | |
133 frameNumberToMasterTimestamp = dict() | |
134 for r,d,f in os.walk(videoFolderPath): | |
135 i = firstFrameNum | |
136 for files in f: | |
137 name, ext = os.path.splitext(files) | |
138 video_timestringsList.append(name) | |
139 masterTimestampList.append(masterTimestamp) | |
140 frameNumberToMasterTimestamp[i] = masterTimestamp | |
141 i = i +1 | |
142 masterTimestamp = masterTimestamp+deltaTimestamp | |
143 | |
144 inputZipVideoName = fileName+".zip" | |
145 print('Zipping files...') | |
146 if not os.path.exists(inputZipVideoName) or not videoFolderExist: | |
147 zipFolder(videoFolderPath, inputZipVideoName) | |
148 print('Zipping files...Done.') | |
149 #We generate the structure for ZipVideo | |
150 if cameraCalibrationFilename != None: | |
151 calibrationFile = cameraCalibrationFilename | |
152 else: | |
153 calibrationFile = 'calib.json' | |
154 zipVideo = ZipVideo(video_zip_file=inputZipVideoName, | |
155 time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file=calibrationFile) | |
156 zipVideo.save(os.path.join(workDirname, 'video.json')) | |
157 | |
158 print('ZipVideo saved') | |
159 obj = SyncedVideos(master_timestamps = [masterTimestamp], | |
160 video_timestrings = [video_timestringsList], | |
161 video_files = ['video.json'], | |
162 fps=fps) | |
163 obj.save(os.path.join(workDirname, 'syncedvideo.json')) | |
164 | |
165 print('SyncedVideos saved') | |
166 | |
167 print('Opening db') | |
168 connection = sqlite3.connect(inputDb) | |
169 cursor = connection.cursor() | |
170 | |
171 | |
172 #Tracket database | |
173 trackset = TrackSet(synced_file = ['syncedvideo.json']) | |
174 | |
175 | |
176 #1) We build the type index dictionnary | |
177 typeDict = getTypeDict(cursor) | |
178 | |
179 idToTrackDict = dict() | |
180 #2) We read the object database | |
181 cursor.execute("SELECT object_id, road_user_type FROM objects") | |
182 for row in cursor: | |
183 objectId = row[0] | |
184 objectType = row[1] | |
185 t = Track(type=typeDict.get(objectType, "unknown")) | |
186 idToTrackDict[objectId] = t; | |
187 trackset.append(t) | |
188 print('Reading boundingbox table') | |
189 #3) We read the bounding box table | |
190 cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='bounding_boxes'") | |
191 data = cursor.fetchone() | |
192 if data == None: | |
193 print('No bounding box table. Maybe it was not generated ?') | |
194 else: | |
195 cursor.execute("SELECT object_id, frame_number, x_top_left, y_top_left, x_bottom_right, y_bottom_right FROM bounding_boxes") | |
196 for row in cursor: | |
197 objectId = row[0] | |
198 frameNumber = row[1] | |
199 x_top_left = row[2] | |
200 y_top_left = row[3] | |
201 x_bottom_right = row[4] | |
202 y_bottom_right = row[5] | |
203 | |
204 idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.], | |
205 master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)], | |
206 bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]])) | |
207 print('Saving db in json') | |
208 trackset.save(os.path.join(workDirname, outputFileName)) | |
209 connection.close() | |
210 print('Done.') | |
211 | |
212 | |
213 | |
214 | |
215 if __name__ == "__main__": | |
216 parser = argparse.ArgumentParser(description='The program convert polytrack.sqlite database to pdtv bounding boxes', epilog = 'Either the configuration filename or the other parameters (at least video and database filenames) need to be provided.') | |
217 parser.add_argument('-w', dest = 'workDirname', help = 'use a different work directory', default = "./",type = str) | |
218 parser.add_argument('--scene', dest = 'sceneFilename', help = 'name of the configuration file', type = str, default = None) | |
219 parser.add_argument('--section', dest = 'sectionName', help = 'name of the section', type = str, default = None) | |
220 parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', type = str, default = None) | |
221 parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file', type = str, default = None) | |
222 parser.add_argument('-c', dest = 'cameraCalibrationFilename', help = 'name of the camera json file', type = str, default = None) | |
223 parser.add_argument('-o', dest = 'outputFilename', help = 'name of the output json file', type = str, default = 'roaduser.json') | |
224 parser.add_argument('-s', dest = 'firstFrameNum', help = 'forced start frame', type = int, default = 0) | |
225 parser.add_argument('-e', dest = 'lastFrameNum', help = 'forced end frame', type = int, default = None) | |
226 #parser.add_argument('-t', dest = 'useDatabaseTimestamp', help = 'use the timestamp of the database', default= False, type = bool) | |
227 parser.add_argument('-u', dest = 'useCurrentVideoFile', help = 'use the previously generated video file', action = 'store_true') | |
228 args = parser.parse_args() | |
229 #convertDatabase(args) | |
230 | |
231 convertDatabase(args.workDirname, args.sectionName, args.sceneFilename, videoFilename = args.videoFilename, databaseFilename = args.databaseFilename, videoFolderExist = args.useCurrentVideoFile, firstFrameNum = args.firstFrameNum, lastFrameNum = args.lastFrameNum, cameraCalibrationFilename = args.cameraCalibrationFilename, outputFileName=args.outputFilename ) |