Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/undistort-video.py @ 760:d72e4bcc1e36 dev

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
added functionality to generate avi from the undistorted videos
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Wed, 04 Nov 2015 17:44:07 -0500
parents 3058e00887bc
children 0f6b0f63eb07
comparison
equal deleted inserted replaced
759:a05b70f307dd 760:d72e4bcc1e36
5 import numpy as np 5 import numpy as np
6 import cv2 6 import cv2
7 7
8 import cvutils 8 import cvutils
9 from math import ceil, log10 9 from math import ceil, log10
10 from os import path, mkdir
10 11
11 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 12 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
12 $ mencoder 'mf://./*.png' -mf fps=[framerate]:type=png -ovc xvid -xvidencopts bitrate=[bitrate] -nosound -o [output.avi]''') 13 $ mencoder 'mf://./*.png' -mf fps=[framerate]:type=png -ovc xvid -xvidencopts bitrate=[bitrate] -nosound -o [output.avi]''')
13 14
14 parser.add_argument('-i', dest = 'videoFilename', help = 'filename of the video sequence') 15 parser.add_argument('-i', dest = 'videoFilename', help = 'filename of the video sequence')
15 parser.add_argument('--intrinsic', dest = 'intrinsicCameraMatrixFilename', help = 'name of the intrinsic camera file') 16 parser.add_argument('--intrinsic', dest = 'intrinsicCameraMatrixFilename', help = 'name of the intrinsic camera file')
16 parser.add_argument('--distortion-coefficients', dest = 'distortionCoefficients', help = 'distortion coefficients', nargs = '*', type = float) 17 parser.add_argument('--distortion-coefficients', dest = 'distortionCoefficients', help = 'distortion coefficients', nargs = '*', type = float)
17 parser.add_argument('--undistorted-multiplication', dest = 'undistortedImageMultiplication', help = 'undistorted image multiplication', type = float) 18 parser.add_argument('--undistorted-multiplication', dest = 'undistortedImageMultiplication', help = 'undistorted image multiplication', type = float)
18 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int) 19 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int)
19 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int) 20 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int)
21 parser.add_argument('-d', dest = 'destinationDirname', help = 'name of the directory where the undistorted frames are saved')
22 parser.add_argument('--encode', dest = 'encodeVideo', help = 'indicate if video is generated at the end (default Xvid)', action = 'store_true')
23 parser.add_argument('--fps', dest = 'fps', help = 'frame per second of the output video file if encoding', type = float, default = 30)
24 parser.add_argument('--bitrate', dest = 'bitrate', help = 'bitrate of the output video file if encoding', type = int, default = 5000)
20 25
21 args = parser.parse_args() 26 args = parser.parse_args()
22 27
23 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename) 28 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename)
24 #distortionCoefficients = args.distortionCoefficients 29 #distortionCoefficients = args.distortionCoefficients
25 #undistortedImageMultiplication = args.undistortedImageMultiplication 30 #undistortedImageMultiplication = args.undistortedImageMultiplication
26 #firstFrameNum = params.firstFrameNum 31 #firstFrameNum = params.firstFrameNum
32 if args.destinationDirname is None:
33 destinationDirname = ''
34 else:
35 if not args.destinationDirname.endswith('/'):
36 destinationDirname = args.destinationDirname+'/'
37 else:
38 destinationDirname = args.destinationDirname
39 if not path.exists(destinationDirname):
40 mkdir(destinationDirname)
27 41
28 capture = cv2.VideoCapture(args.videoFilename) 42 capture = cv2.VideoCapture(args.videoFilename)
29 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) 43 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
30 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) 44 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
31 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients) 45 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients)
41 nZerosFilename = int(ceil(log10(lastFrameNum))) 55 nZerosFilename = int(ceil(log10(lastFrameNum)))
42 while ret and frameNum < lastFrameNum: 56 while ret and frameNum < lastFrameNum:
43 ret, img = capture.read() 57 ret, img = capture.read()
44 if ret: 58 if ret:
45 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) 59 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
46 cv2.imwrite('undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img) 60 cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
47 frameNum += 1 61 frameNum += 1
62
63 if args.encodeVideo:
64 print('Encoding the images files in video')
65 from subprocess import check_call
66 from storage import openCheck
67 out = openCheck("err.log", "w")
68 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)
69 out.close()