changeset 403:f3938bb6da7f

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
files scripts/polytracktopdtv.py
diffstat 1 files changed, 58 insertions(+), 50 deletions(-) [+]
line wrap: on
line diff
--- a/scripts/polytracktopdtv.py	Mon Jul 29 19:45:43 2013 -0400
+++ b/scripts/polytracktopdtv.py	Tue Jul 30 19:45:03 2013 -0400
@@ -1,7 +1,5 @@
-#! /usr/bin/env python
-
 from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State
-import sys, os, datetime
+import sys, os, datetime, argparse
 import shutil
 import sqlite3
 import zipfile
@@ -51,20 +49,20 @@
 
     return typeDict
 
-def extractFrames(videoFile, framePath, fps, time, startFrame = 0, endFrame = None):
+def extractFrames(videoFile, framePath, fps, time, firstFrameNum = 0, lastFrameNum = None):
     '''Method to extract all the frames of the video'''
     print('Extracting frame')
     deltaTimestamp = 1000.0/float(fps);
-    time+=datetime.timedelta(microseconds=startFrame*deltaTimestamp*1000)
+    time+=datetime.timedelta(microseconds=firstFrameNum*deltaTimestamp*1000)
     
     inc = 1000 #How many frame we fetch in the video at a time
 
-    if endFrame != None:
-        delta = endFrame-startFrame
+    if lastFrameNum != None:
+        delta = lastFrameNum-firstFrameNum
         if delta < inc:
             inc = delta
     
-    currentIdx = startFrame
+    currentIdx = firstFrameNum
     frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
     
     while len(frameList) == inc and inc > 0:
@@ -74,8 +72,8 @@
             time += datetime.timedelta(microseconds=deltaTimestamp*1000)
         currentIdx = currentIdx + inc
 
-        if endFrame != None:            
-            delta = endFrame-currentIdx
+        if lastFrameNum != None:            
+            delta = lastFrameNum-currentIdx
             if delta < inc:
                 inc = delta        
         if inc:
@@ -85,49 +83,46 @@
 
     
 
-def convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = None, videoFolderExist = False, startFrame = 0, endFrame = None):
+def convertDatabase(workDirname, sectionName, sceneFilename = None, databaseFilename = None, videoFilename = None, videoFolderExist = False, firstFrameNum = 0, lastFrameNum = None, cameraCalibrationFilename = None, outputFileName = 'roaduser.json'):
     '''
     Method to convert database from polytrack to PDTV:
     workDirname  is the current working directory
     sceneFilename is the path to the .cfg file
     sectionName is the name of the section we want to process in this file
     videoFolderExist specifiy if we want to reextract the video frame or if they already exist at workdir/videoframes/
-    startFrame is the first frame we want to extract
-    endFrame is the last frame we want to extract (or None if we want to extract everything)
+    firstFrameNum is the first frame we want to extract
+    lastFrameNum is the last frame we want to extract (or None if we want to extract everything)
     '''
+    error = False
+    if sceneFilename != None:
+        scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename))
+        time = scene[sectionName].date
+        inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename)
+        videoFile = os.path.join(workDirname, scene[sectionName].videoFilename)
 
-    scene = utils.SceneParameters.loadConfigFile(os.path.join(workDirname, sceneFilename))
     if databaseFilename != None:
         inputDb = os.path.join(workDirname, databaseFilename)
-    else:
-        inputDb = os.path.join(workDirname, scene[sectionName].databaseFilename)
-    videoFile = os.path.join(workDirname, scene[sectionName].videoFilename)
-
-    error = False
-
-    if not os.path.exists(inputDb):
-        print('Database path is not valid: ' + inputDb)
-        error = True
-    if not os.path.exists(videoFile):
-        print('Video path is not valid: ' + videoFile)
-        error = True
-
-    time = scene[sectionName].date    
-
+    if videoFilename != None:
+        videoFile = os.path.join(workDirname, videoFilename)
+    # elif videoFolderExist == False:            
+    #     print('No video path specified')
+    #     error = True
     
     videoFolderPath = os.path.join(workDirname, "videoframes/")
-    fileName = scene[sectionName].sitename
+    fileName = sectionName
     
-    fps = cvutils.getFPS(videoFile)
-    print('Video should run at ' + str(fps) + ' fps')
-    deltaTimestamp = 1000.0/float(fps);
-    if videoFolderExist == False:
-        if os.path.exists(videoFolderPath):
-            shutil.rmtree(videoFolderPath)
-        utils.mkdir(videoFolderPath)
-        if extractFrames(videoFile, videoFolderPath, fps, time, startFrame, endFrame) == 0:
-            print("Error. Frame were not extracted")
-            error = True
+    if videoFile != None:
+        fps = cvutils.getFPS(videoFile)
+        print('Video should run at ' + str(fps) + ' fps')
+        deltaTimestamp = 1000.0/float(fps);
+        if videoFolderExist == False:
+            if os.path.exists(videoFolderPath):
+                shutil.rmtree(videoFolderPath)
+            utils.mkdir(videoFolderPath)
+            if extractFrames(videoFile, videoFolderPath, fps, time, firstFrameNum, lastFrameNum) == 0:
+                print("Error. Frame were not extracted")
+                error = True
+
     
     if not error:
         masterTimestamp = 0.0
@@ -135,7 +130,7 @@
         video_timestringsList = list()
         frameNumberToMasterTimestamp = dict()
         for r,d,f in os.walk(videoFolderPath):
-            i = startFrame
+            i = firstFrameNum
             for files in f:
                 name, ext = os.path.splitext(files)           
                 video_timestringsList.append(name)
@@ -150,8 +145,12 @@
             zipFolder(videoFolderPath, inputZipVideoName)
         print('Zipping files...Done.')
         #We generate the structure for ZipVideo
+        if cameraCalibrationFilename != None:
+            calibrationFile = cameraCalibrationFilename
+        else:
+            calibrationFile = 'calib.json'
         zipVideo = ZipVideo(video_zip_file=inputZipVideoName,
-                        time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file='calib.json')
+                        time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file=calibrationFile)
         zipVideo.save(os.path.join(workDirname, 'video.json'))
 
         print('ZipVideo saved')
@@ -204,7 +203,7 @@
                                                  master_timestamp=frameNumberToMasterTimestamp[int(frameNumber)],
                                                  bounding_boxes=[[(x_top_left, x_bottom_right), (y_top_left, y_bottom_right)]]))
         print('Saving db in json')
-        trackset.save(os.path.join(workDirname, 'roaduser.json'))
+        trackset.save(os.path.join(workDirname, outputFileName))
         connection.close()
         print('Done.')
 
@@ -212,10 +211,19 @@
     
 
 if __name__ == "__main__":
-    workDirname = "./"
-    sceneFilename = "scene.cfg"
-    sectionName = "amherst-11"
-    # videoFrameAlreadyExtracted Set this to true if you have the frame already extracted in the folder workDirname/videoframes/
-    convertDatabase(workDirname, sceneFilename, sectionName, databaseFilename = 'amherst-11_gt.sqlite', videoFolderExist = True, startFrame = 0, endFrame = 4000) 
-    
-    
+    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.')
+    parser.add_argument('-w', dest = 'workDirname', help = 'use a different work directory', default = "./",type = str)
+    parser.add_argument('--scene', dest = 'sceneFilename', help = 'name of the configuration file', type = str, default = None)
+    parser.add_argument('--section', dest = 'sectionName', help = 'name of the section', type = str, default = None)
+    parser.add_argument('-d', dest = 'databaseFilename', help = 'name of the Sqlite database file', type = str, default = None)
+    parser.add_argument('-i', dest = 'videoFilename', help = 'name of the video file', type = str, default = None) 
+    parser.add_argument('-c', dest = 'cameraCalibrationFilename', help = 'name of the camera json file', type = str, default = None)
+    parser.add_argument('-o', dest = 'outputFilename', help = 'name of the output json file', type = str, default = 'roaduser.json')
+    parser.add_argument('-s', dest = 'firstFrameNum', help = 'forced start frame', type = int, default = 0)
+    parser.add_argument('-e', dest = 'lastFrameNum', help = 'forced end frame', type = int, default = None)
+    #parser.add_argument('-t', dest = 'useDatabaseTimestamp', help = 'use the timestamp of the database', default= False, type = bool) 
+    parser.add_argument('-u', dest = 'useCurrentVideoFile', help = 'use the previously generated video file', action = 'store_true')
+    args = parser.parse_args()
+    #convertDatabase(args)
+
+    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 )