diff scripts/undistort-video.py @ 795:a34ec862371f

merged with dev branch
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Mon, 09 May 2016 15:33:11 -0400
parents 0f6b0f63eb07
children 52aa03260f03
line wrap: on
line diff
--- a/scripts/undistort-video.py	Tue Nov 03 13:48:56 2015 -0500
+++ b/scripts/undistort-video.py	Mon May 09 15:33:11 2016 -0400
@@ -7,6 +7,7 @@
 
 import cvutils
 from math import ceil, log10
+from os import path, mkdir
 
 parser = argparse.ArgumentParser(description='''The program converts a video into a series of images corrected for distortion. One can then use mencoder to generate a movie, eg
 $ mencoder 'mf://./*.png' -mf fps=[framerate]:type=png -ovc xvid -xvidencopts bitrate=[bitrate] -nosound -o [output.avi]''')
@@ -17,6 +18,10 @@
 parser.add_argument('--undistorted-multiplication', dest = 'undistortedImageMultiplication', help = 'undistorted image multiplication', type = float)
 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int)
 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int)
+parser.add_argument('-d', dest = 'destinationDirname', help = 'name of the directory where the undistorted frames are saved')
+parser.add_argument('--encode', dest = 'encodeVideo', help = 'indicate if video is generated at the end (default Xvid)', action = 'store_true')
+parser.add_argument('--fps', dest = 'fps', help = 'frame per second of the output video file if encoding', type = float, default = 30)
+parser.add_argument('--bitrate', dest = 'bitrate', help = 'bitrate of the output video file if encoding', type = int, default = 5000)
 
 args = parser.parse_args()
 
@@ -24,15 +29,24 @@
 #distortionCoefficients = args.distortionCoefficients
 #undistortedImageMultiplication = args.undistortedImageMultiplication
 #firstFrameNum = params.firstFrameNum
+if args.destinationDirname is None:
+    destinationDirname = ''
+else:
+    if not args.destinationDirname.endswith('/'):
+        destinationDirname = args.destinationDirname+'/'
+    else:
+        destinationDirname = args.destinationDirname
+    if not path.exists(destinationDirname):
+        mkdir(destinationDirname)
 
 capture = cv2.VideoCapture(args.videoFilename)
-width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
-height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
+width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients)
 if capture.isOpened():
     ret = True
     frameNum = args.firstFrameNum
-    capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, args.firstFrameNum)
+    capture.set(cv2.CAP_PROP_POS_FRAMES, args.firstFrameNum)
     if args.lastFrameNum is None:
         from sys import maxint
         lastFrameNum = maxint
@@ -43,5 +57,13 @@
             ret, img = capture.read()
             if ret:
                 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
-                cv2.imwrite('undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
+                cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
             frameNum += 1
+
+if args.encodeVideo:
+    print('Encoding the images files in video')
+    from subprocess import check_call
+    from storage import openCheck
+    out = openCheck("err.log", "w")
+    check_call("mencoder \'mf://"+destinationDirname+"*.png\' -mf fps={}:type=png -ovc xvid -xvidencopts bitrate={} -nosound -o ".format(args.fps, args.bitrate)+destinationDirname+"undistort.avi", stderr = out, shell = True)
+    out.close()