diff python/cvutils.py @ 509:935430b1d408

corrected mask bug in feature tracking, updated display-trajectories to display on undistorted image
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Fri, 23 May 2014 16:27:26 -0400
parents 343cfd185ca6
children b0dac840c24f
line wrap: on
line diff
--- a/python/cvutils.py	Fri May 23 10:35:51 2014 -0400
+++ b/python/cvutils.py	Fri May 23 16:27:26 2014 -0400
@@ -202,13 +202,24 @@
         return imgcrop, yCropMin, yCropMax, xCropMin, xCropMax
 
 
-    def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1, saveAllImages = False):
+    def displayTrajectories(videoFilename, objects, boundingBoxes = {}, homography = None, firstFrameNum = 0, lastFrameNumArg = None, printFrames = True, rescale = 1., nFramesStep = 1, saveAllImages = False, undistort = False, intrinsicCameraMatrix = None, distortionCoefficients = None, undistortedImageMultiplication = 1.):
         '''Displays the objects overlaid frame by frame over the video '''
         from moving import userTypeNames
         from math import ceil, log10
+        from numpy import identity, array
+
         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))
+
+        if undistort: # setup undistortion
+            newImgSize = (int(round(width*undistortedImageMultiplication)), int(round(height*undistortedImageMultiplication)))
+            from copy import deepcopy
+            newCameraMatrix = deepcopy(intrinsicCameraMatrix)
+            newCameraMatrix[0,2] = newImgSize[0]/2.
+            newCameraMatrix[1,2] = newImgSize[1]/2.
+            [map1, map2] = cv2.initUndistortRectifyMap(intrinsicCameraMatrix, array(distortionCoefficients), identity(3), newCameraMatrix, newImgSize, cv2.CV_32FC1)
+
         if capture.isOpened():
             key = -1
             ret = True
@@ -223,6 +234,8 @@
             while ret and not quitKey(key) and frameNum < lastFrameNum:
                 ret, img = capture.read()
                 if ret:
+                    if undistort:
+                        img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
                     if printFrames:
                         print('frame {0}'.format(frameNum))
                     for obj in objects: