Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/polytracktopdtv.py @ 614:5e09583275a4

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
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 )