diff scripts/polytracktopdtv.py @ 402:f29204e68aab

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
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/scripts/polytracktopdtv.py	Mon Jul 29 19:45:43 2013 -0400
@@ -0,0 +1,221 @@
+#! /usr/bin/env python
+
+from pdtv import TsaiCamera, ZipVideo, SyncedVideos, TrackSet, Track, State
+import sys, os, datetime
+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, startFrame = 0, endFrame = 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)
+    
+    inc = 1000 #How many frame we fetch in the video at a time
+
+    if endFrame != None:
+        delta = endFrame-startFrame
+        if delta < inc:
+            inc = delta
+    
+    currentIdx = startFrame
+    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 endFrame != None:            
+            delta = endFrame-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, sceneFilename, sectionName, databaseFilename = None, videoFolderExist = False, startFrame = 0, endFrame = None):
+    '''
+    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)
+    '''
+
+    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    
+
+    
+    videoFolderPath = os.path.join(workDirname, "videoframes/")
+    fileName = scene[sectionName].sitename
+    
+    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 not error:
+        masterTimestamp = 0.0
+        masterTimestampList = list()
+        video_timestringsList = list()
+        frameNumberToMasterTimestamp = dict()
+        for r,d,f in os.walk(videoFolderPath):
+            i = startFrame
+            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
+        zipVideo = ZipVideo(video_zip_file=inputZipVideoName,
+                        time_offset=0.0, time_scale=1.0, master_timestamps=masterTimestampList, calibration_file='calib.json')
+        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, 'roaduser.json'))
+        connection.close()
+        print('Done.')
+
+    
+    
+
+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) 
+    
+