diff 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
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/scripts/polytracktopdtv.py	Fri Dec 05 12:13:53 2014 -0500
@@ -0,0 +1,231 @@
+#! /usr/bin/env python
+
+from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State
+import sys, os, datetime, argparse
+import shutil
+import sqlite3
+import zipfile
+import utils
+import cvutils
+import cv2
+
+
+def zipFolder(inputFolder, outputFile):
+    '''Method to compress the content of the inputFolder in the outputFile'''
+    zip = zipfile.ZipFile(outputFile, 'w')
+    for root, dirs, files in os.walk(inputFolder):
+        for file in files:
+            zip.write(root+file, file)
+    zip.close()
+
+
+
+
+def getTypeDict(cursor):
+    '''Return a dictionnary with integer key and associated type string
+    i.e.: "0"  -> "unknown"
+          "1"  -> "car"
+          "2"  -> "pedestrians"
+          "3"  -> "motorcycle"
+          "4"  -> "bicycle"
+          "5"  -> "bus"
+          "6"  -> "truck"
+          ... and other type if the objects_type table is defined in SQLite'''
+    typeDict = dict()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='objects_type'")
+    data = cursor.fetchone()
+
+    if(data == None):
+        typeDict["0"] = "unknown"
+        typeDict["1"] = "car"
+        typeDict["2"] = "pedestrians"
+        typeDict["3"] = "motorcycle"
+        typeDict["4"] = "bicycle"
+        typeDict["5"] = "bus"
+        typeDict["6"] = "truck"
+        
+    else:
+        cursor.execute("SELECT road_user_type, type_string FROM objects_type")
+        for row in cursor:
+            typeDict[row[0]]= row[1]
+
+    return typeDict
+
+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=firstFrameNum*deltaTimestamp*1000)
+    
+    inc = 1000 #How many frame we fetch in the video at a time
+
+    if lastFrameNum != None:
+        delta = lastFrameNum-firstFrameNum
+        if delta < inc:
+            inc = delta
+    
+    currentIdx = firstFrameNum
+    frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
+    
+    while len(frameList) == inc and inc > 0:
+        
+        for f in frameList:
+            cv2.imwrite(os.path.join(framePath,time.strftime("%Y%m%d-%H%M%S.%f")[:-3]+'.jpg'), f)
+            time += datetime.timedelta(microseconds=deltaTimestamp*1000)
+        currentIdx = currentIdx + inc
+
+        if lastFrameNum != None:            
+            delta = lastFrameNum-currentIdx
+            if delta < inc:
+                inc = delta        
+        if inc:
+            frameList = cvutils.getImagesFromVideo(videoFile, firstFrameNum = currentIdx, nFrames = inc)
+        print('Extracting frame ' + str(currentIdx))
+    return len(frameList) > 0
+
+    
+
+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/
+    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)
+
+    if databaseFilename != None:
+        inputDb = os.path.join(workDirname, databaseFilename)
+    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 = sectionName
+    
+    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
+        masterTimestampList = list()
+        video_timestringsList = list()
+        frameNumberToMasterTimestamp = dict()
+        for r,d,f in os.walk(videoFolderPath):
+            i = firstFrameNum
+            for files in f:
+                name, ext = os.path.splitext(files)           
+                video_timestringsList.append(name)
+                masterTimestampList.append(masterTimestamp)            
+                frameNumberToMasterTimestamp[i] = masterTimestamp
+                i = i +1
+                masterTimestamp = masterTimestamp+deltaTimestamp
+        
+        inputZipVideoName = fileName+".zip"
+        print('Zipping files...')
+        if not os.path.exists(inputZipVideoName) or not videoFolderExist:
+            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=calibrationFile)
+        zipVideo.save(os.path.join(workDirname, 'video.json'))
+
+        print('ZipVideo saved')
+        obj = SyncedVideos(master_timestamps = [masterTimestamp],
+                               video_timestrings = [video_timestringsList],
+                               video_files = ['video.json'],
+                               fps=fps)
+        obj.save(os.path.join(workDirname, 'syncedvideo.json'))
+
+        print('SyncedVideos saved')
+
+        print('Opening db')
+        connection = sqlite3.connect(inputDb)
+        cursor = connection.cursor()
+      
+
+        #Tracket database
+        trackset = TrackSet(synced_file = ['syncedvideo.json'])
+
+        
+        #1) We build the type index dictionnary
+        typeDict = getTypeDict(cursor)
+        
+        idToTrackDict = dict()
+        #2) We read the object database    
+        cursor.execute("SELECT object_id, road_user_type FROM objects")
+        for row in cursor:
+            objectId = row[0]
+            objectType = row[1]
+            t = Track(type=typeDict.get(objectType, "unknown"))
+            idToTrackDict[objectId] = t;
+            trackset.append(t)
+        print('Reading boundingbox table')
+        #3) We read the bounding box table    
+        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='bounding_boxes'")
+        data = cursor.fetchone()
+        if data == None:
+            print('No bounding box table. Maybe it was not generated ?')
+        else:
+            cursor.execute("SELECT object_id, frame_number, x_top_left, y_top_left, x_bottom_right, y_bottom_right FROM bounding_boxes")
+            for row in cursor:
+                objectId = row[0]
+                frameNumber = row[1]
+                x_top_left = row[2]
+                y_top_left = row[3]
+                x_bottom_right = row[4]
+                y_bottom_right = row[5]
+                
+                idToTrackDict[objectId].append(State(frame=int(frameNumber), world_position = [0.,0.], 
+                                                 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, outputFileName))
+        connection.close()
+        print('Done.')
+
+    
+    
+
+if __name__ == "__main__":
+    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 )