Mercurial Hosting > traffic-intelligence

comparison scripts/polytracktopdtv.py @ 403:f3938bb6da7f

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
added commande line options to script to convert polytrack annotation data
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Tue, 30 Jul 2013 19:45:03 -0400
parents f29204e68aab
children 4dd68b0765b1
comparison
equal deleted inserted replaced
402:f29204e68aab 403:f3938bb6da7f
1 #! /usr/bin/env python
2
3 from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State 1 from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State
4 import sys, os, datetime 2 import sys, os, datetime, argparse
5 import shutil 3 import shutil
6 import sqlite3 4 import sqlite3
7 import zipfile 5 import zipfile
8 import utils 6 import utils
9 import cvutils 7 import cvutils
49 for row in cursor: 47 for row in cursor:
50 typeDict[row[0]]= row[1] 48 typeDict[row[0]]= row[1]
51 49
52 return typeDict 50 return typeDict
53 51
54 def extractFrames(videoFile, framePath, fps, time, startFrame = 0, endFrame = None): 52 def extractFrames(videoFile, framePath, fps, time, firstFrameNum = 0, lastFrameNum = None):
55 '''Method to extract all the frames of the video''' 53 '''Method to extract all the frames of the video'''
56 print('Extracting frame') 54 print('Extracting frame')
57 deltaTimestamp = 1000.0/float(fps); 55 deltaTimestamp = 1000.0/float(fps);
58 time+=datetime.timedelta(microseconds=startFrame*deltaTimestamp*1000) 56 time+=datetime.timedelta(microseconds=firstFrameNum*deltaTimestamp*1000)
59 57
60 inc = 1000 #How many frame we fetch in the video at a time 58 inc = 1000 #How many frame we fetch in the video at a time
61 59
62 if endFrame != None: 60 if lastFrameNum != None:
63 delta = endFrame-startFrame 61 delta = lastFrameNum-firstFrameNum
64 if delta < inc: 62 if delta < inc:
65 inc = delta 63 inc = delta
66 64
67 currentIdx = startFrame 65 currentIdx = firstFrameNum
68 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc) 66 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
69 67
70 while len(frameList) == inc and inc > 0: 68 while len(frameList) == inc and inc > 0:
71 69
72 for f in frameList: 70 for f in frameList:
73 cv2.imwrite(os.path.join(framePath,time.strftime("%Y%m%d-%H%M%S.%f")[:-3]+'.jpg'), f) 71 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) 72 time += datetime.timedelta(microseconds=deltaTimestamp*1000)
75 currentIdx = currentIdx + inc 73 currentIdx = currentIdx + inc
76 74
77 if endFrame != None: 75 if lastFrameNum != None:
78 delta = endFrame-currentIdx 76 delta = lastFrameNum-currentIdx
79 if delta < inc: 77 if delta < inc:
80 inc = delta 78 inc = delta
81 if inc: 79 if inc:
82 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc) 80 frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
83 print('Extracting frame ' + str(currentIdx)) 81 print('Extracting frame ' + str(currentIdx))
84 return len(frameList) > 0 82 return len(frameList) > 0
85 83
86 84
87 85
88 def convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = None, videoFolderExist = False, startFrame = 0, endFrame = None): 86 def convertDatabase(workDirname, sectionName, sceneFilename = None, databaseFilename = None, videoFilename = None, videoFolderExist = False, firstFrameNum = 0, lastFrameNum = None, cameraCalibrationFilename = None, outputFileName = 'roaduser.json'):
89 ''' 87 '''
90 Method to convert database from polytrack to PDTV: 88 Method to convert database from polytrack to PDTV:
91 workDirname is the current working directory 89 workDirname is the current working directory
92 sceneFilename is the path to the .cfg file 90 sceneFilename is the path to the .cfg file
93 sectionName is the name of the section we want to process in this file 91 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/ 92 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 93 firstFrameNum 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) 94 lastFrameNum is the last frame we want to extract (or None if we want to extract everything)
97 ''' 95 '''
98 96 error = False
99 scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename)) 97 if sceneFilename != None:
98 scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename))
99 time = scene[sectionName].date
100 inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename)
101 videoFile = os.path.join(workDirname, scene[sectionName].videoFilename)
102
100 if databaseFilename != None: 103 if databaseFilename != None:
101 inputDb = os.path.join(workDirname, databaseFilename) 104 inputDb = os.path.join(workDirname, databaseFilename)
102 else: 105 if videoFilename != None:
103 inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename) 106 videoFile = os.path.join(workDirname, videoFilename)
104 videoFile = os.path.join(workDirname, scene[sectionName].videoFilename) 107 # elif videoFolderExist == False:
105 108 # print('No video path specified')
106 error = False 109 # error = True
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 110
118 videoFolderPath = os.path.join(workDirname, "videoframes/") 111 videoFolderPath = os.path.join(workDirname, "videoframes/")
119 fileName = scene[sectionName].sitename 112 fileName = sectionName
120 113
121 fps = cvutils.getFPS(videoFile) 114 if videoFile != None:
122 print('Video should run at ' + str(fps) + ' fps') 115 fps = cvutils.getFPS(videoFile)
123 deltaTimestamp = 1000.0/float(fps); 116 print('Video should run at ' + str(fps) + ' fps')
124 if videoFolderExist == False: 117 deltaTimestamp = 1000.0/float(fps);
125 if os.path.exists(videoFolderPath): 118 if videoFolderExist == False:
126 shutil.rmtree(videoFolderPath) 119 if os.path.exists(videoFolderPath):
127 utils.mkdir(videoFolderPath) 120 shutil.rmtree(videoFolderPath)
128 if extractFrames(videoFile, videoFolderPath, fps, time, startFrame, endFrame) == 0: 121 utils.mkdir(videoFolderPath)
129 print("Error. Frame were not extracted") 122 if extractFrames(videoFile, videoFolderPath, fps, time, firstFrameNum, lastFrameNum) == 0:
130 error = True 123 print("Error. Frame were not extracted")
124 error = True
125
131 126
132 if not error: 127 if not error:
133 masterTimestamp = 0.0 128 masterTimestamp = 0.0
134 masterTimestampList = list() 129 masterTimestampList = list()
135 video_timestringsList = list() 130 video_timestringsList = list()
136 frameNumberToMasterTimestamp = dict() 131 frameNumberToMasterTimestamp = dict()
137 for r,d,f in os.walk(videoFolderPath): 132 for r,d,f in os.walk(videoFolderPath):
138 i = startFrame 133 i = firstFrameNum
139 for files in f: 134 for files in f:
140 name, ext = os.path.splitext(files) 135 name, ext = os.path.splitext(files)
141 video_timestringsList.append(name) 136 video_timestringsList.append(name)
142 masterTimestampList.append(masterTimestamp) 137 masterTimestampList.append(masterTimestamp)
143 frameNumberToMasterTimestamp[i] = masterTimestamp 138 frameNumberToMasterTimestamp[i] = masterTimestamp
148 print('Zipping files...') 143 print('Zipping files...')
149 if not os.path.exists(inputZipVideoName) or not videoFolderExist: 144 if not os.path.exists(inputZipVideoName) or not videoFolderExist:
150 zipFolder(videoFolderPath, inputZipVideoName) 145 zipFolder(videoFolderPath, inputZipVideoName)
151 print('Zipping files...Done.') 146 print('Zipping files...Done.')
152 #We generate the structure for ZipVideo 147 #We generate the structure for ZipVideo
148 if cameraCalibrationFilename != None:
149 calibrationFile = cameraCalibrationFilename
150 else:
151 calibrationFile = 'calib.json'
153 zipVideo = ZipVideo(video_zip_file=inputZipVideoName, 152 zipVideo = ZipVideo(video_zip_file=inputZipVideoName,
154 time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file='calib.json') 153 time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file=calibrationFile)
155 zipVideo.save(os.path.join(workDirname, 'video.json')) 154 zipVideo.save(os.path.join(workDirname, 'video.json'))
156 155
157 print('ZipVideo saved') 156 print('ZipVideo saved')
158 obj = SyncedVideos(master_timestamps = [masterTimestamp], 157 obj = SyncedVideos(master_timestamps = [masterTimestamp],
159 video_timestrings = [video_timestringsList], 158 video_timestrings = [video_timestringsList],
202 201
203 idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.], 202 idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.],
204 master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)], 203 master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)],
205 bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]])) 204 bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]]))
206 print('Saving db in json') 205 print('Saving db in json')
207 trackset.save(os.path.join(workDirname, 'roaduser.json')) 206 trackset.save(os.path.join(workDirname, outputFileName))
208 connection.close() 207 connection.close()
209 print('Done.') 208 print('Done.')
210 209
211 210
212 211
213 212
214 if __name__ == "__main__": 213 if __name__ == "__main__":
215 workDirname = "./" 214 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.')
216 sceneFilename = "scene.cfg" 215 parser.add_argument('-w', dest = 'workDirname', help = 'use a different work directory', default = "./",type = str)
217 sectionName = "amherst-11" 216 parser.add_argument('--scene', dest = 'sceneFilename', help = 'name of the configuration file', type = str, default = None)
218 # videoFrameAlreadyExtracted Set this to true if you have the frame already extracted in the folder workDirname/videoframes/ 217 parser.add_argument('--section', dest = 'sectionName', help = 'name of the section', type = str, default = None)
219 convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = 'amherst-11_gt.sqlite', videoFolderExist = True, startFrame = 0, endFrame = 4000) 218 parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', type = str, default = None)
220 219 parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file', type = str, default = None)
221 220 parser.add_argument('-c', dest = 'cameraCalibrationFilename', help = 'name of the camera json file', type = str, default = None)
221 parser.add_argument('-o', dest = 'outputFilename', help = 'name of the output json file', type = str, default = 'roaduser.json')
222 parser.add_argument('-s', dest = 'firstFrameNum', help = 'forced start frame', type = int, default = 0)
223 parser.add_argument('-e', dest = 'lastFrameNum', help = 'forced end frame', type = int, default = None)
224 #parser.add_argument('-t', dest = 'useDatabaseTimestamp', help = 'use the timestamp of the database', default= False, type = bool)
225 parser.add_argument('-u', dest = 'useCurrentVideoFile', help = 'use the previously generated video file', action = 'store_true')
226 args = parser.parse_args()
227 #convertDatabase(args)
228
229 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 )