Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/polytracktopdtv.py @ 402:f29204e68aab

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