diff python/cvutils.py @ 482:f6415f012640

adding functionalities (save images directly to display trajectories to create movies
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 02 Apr 2014 16:12:24 -0400
parents d337bffd7283
children 343cfd185ca6
line wrap: on
line diff
--- a/python/cvutils.py	Wed Apr 02 01:45:53 2014 -0400
+++ b/python/cvutils.py	Wed Apr 02 16:12:24 2014 -0400
@@ -202,10 +202,10 @@
         return imgcrop, yCropMin, yCropMax, xCropMin, xCropMax
 
 
-    def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1):
+    def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1, saveAllImages = False):
         '''Displays the objects overlaid frame by frame over the video '''
         from moving import userTypeNames
-
+        from math import ceil, log10
         capture = cv2.VideoCapture(videoFilename)
         width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
         height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
@@ -214,11 +214,12 @@
             ret = True
             frameNum = firstFrameNum
             capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, firstFrameNum)
-            if not lastFrameNumArg:
+            if lastFrameNumArg == None:
                 from sys import maxint
                 lastFrameNum = maxint
             else:
                 lastFrameNum = lastFrameNumArg
+            nZerosFilename = int(ceil(log10(lastFrameNum)))
             while ret and not quitKey(key) and frameNum < lastFrameNum:
                 ret, img = capture.read()
                 if ret:
@@ -242,10 +243,11 @@
                             if userTypeNames[obj.userType] != 'unknown':
                                 objDescription += userTypeNames[obj.userType][0].upper()
                             cv2.putText(img, objDescription, obj.projectedPositions[frameNum-obj.getFirstInstant()].asint().astuple(), cv2.cv.CV_FONT_HERSHEY_PLAIN, 1, cvRed)
-                    cvImshow('frame', img, rescale)
-                    key = cv2.waitKey()
-                    if saveKey(key):
-                        cv2.imwrite('image.png', img)
+                    if not saveAllImages:
+                        cvImshow('frame', img, rescale)
+                        key = cv2.waitKey()
+                    if saveAllImages or saveKey(key):
+                        cv2.imwrite('image-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
                     frameNum += nFramesStep
                     if nFramesStep > 1:
                         capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frameNum)