diff python/cvutils.py @ 513:dbf4b83afbb9

pulled in and merged the new functionalities to deal with camera distortion (eg GoPro cameras)
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 04 Jun 2014 10:57:09 -0400
parents ad518f0c3218
children 727e3c529519
line wrap: on
line diff
--- a/python/cvutils.py	Wed Jun 04 10:56:12 2014 -0400
+++ b/python/cvutils.py	Wed Jun 04 10:57:09 2014 -0400
@@ -110,6 +110,15 @@
         else:
             cv2.imshow(windowName, img)
 
+    def computeUndistortMaps(width, height, undistortedImageMultiplication, intrinsicCameraMatrix, distortionCoefficients):
+        from copy import deepcopy
+        from numpy import identity, array
+        newImgSize = (int(round(width*undistortedImageMultiplication)), int(round(height*undistortedImageMultiplication)))
+        newCameraMatrix = deepcopy(intrinsicCameraMatrix)
+        newCameraMatrix[0,2] = newImgSize[0]/2.
+        newCameraMatrix[1,2] = newImgSize[1]/2.
+        return cv2.initUndistortRectifyMap(intrinsicCameraMatrix, array(distortionCoefficients), identity(3), newCameraMatrix, newImgSize, cv2.CV_32FC1)
+
     def playVideo(filename, firstFrameNum = 0, frameRate = -1, interactive = False, printFrames = True, text = None, rescale = 1.):
         '''Plays the video'''
         wait = 5
@@ -202,13 +211,17 @@
         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
+
         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
+            [map1, map2] = computeUndistortMaps(width, height, undistortedImageMultiplication, intrinsicCameraMatrix, distortionCoefficients)
         if capture.isOpened():
             key = -1
             ret = True
@@ -223,6 +236,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: