Mercurial Hosting > traffic-intelligence

comparison scripts/undistort-video.py @ 1024:acb4f6f6545d

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
corrected issues with update to OpenCV 3
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Sun, 10 Jun 2018 23:26:14 -0400
parents 933670761a57
children 97247e44b827
comparison
equal deleted inserted replaced
1023:a13f47c8931d 1024:acb4f6f6545d
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 from pathlib import Path
11 11
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 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
13 $ 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]''')
14 14
15 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')
16 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')
17 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)
18 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, default = 1.)
19 parser.add_argument('--mask', dest = 'maskFilename', help = 'name of the mask file, to undistort to see how it covers the undistortion errors') 19 parser.add_argument('--mask', dest = 'maskFilename', help = 'name of the mask file, to undistort to see how it covers the undistortion errors')
20 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int, default = 0) 20 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int, default = 0)
21 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int) 21 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int)
22 parser.add_argument('-d', dest = 'destinationDirname', help = 'name of the directory where the undistorted frames are saved') 22 parser.add_argument('-d', dest = 'destinationDirname', help = 'name of the directory where the undistorted frames are saved')
23 parser.add_argument('--encode', dest = 'encodeVideo', help = 'indicate if video is generated at the end (default Xvid)', action = 'store_true') 23 parser.add_argument('--encode', dest = 'encodeVideo', help = 'indicate if video is generated at the end (default Xvid)', action = 'store_true')
26 26
27 args = parser.parse_args() 27 args = parser.parse_args()
28 28
29 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename) 29 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename)
30 if args.destinationDirname is None: 30 if args.destinationDirname is None:
31 destinationDirname = '' 31 destinationPath = Path('.')
32 else: 32 else:
33 if not args.destinationDirname.endswith('/'): 33 destinationPath = Path(destinationPath)
34 destinationDirname = args.destinationDirname+'/' 34 if not destinationPath.exists():
35 else: 35 destinationPath.mkdir()
36 destinationDirname = args.destinationDirname
37 if not path.exists(destinationDirname):
38 mkdir(destinationDirname)
39 36
40 capture = cv2.VideoCapture(args.videoFilename) 37 capture = cv2.VideoCapture(args.videoFilename)
41 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) 38 width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
42 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) 39 height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
43 [map1, map2], newCameraMatrix = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients) 40 [map1, map2], newCameraMatrix = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients)
44 if args.maskFilename is not None: 41 if args.maskFilename is not None:
45 mask = cv2.imread(args.maskFilename) 42 mask = cv2.imread(args.maskFilename)
46 undistortedMask = cv2.remap(mask, map1, map2, interpolation=cv2.INTER_LINEAR)/255 43 undistortedMask = cv2.remap(mask, map1, map2, interpolation=cv2.INTER_LINEAR)/255
47 44
48 if capture.isOpened(): 45 if capture.isOpened():
49 ret = True 46 ret = True
50 frameNum = args.firstFrameNum 47 frameNum = args.firstFrameNum
51 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, args.firstFrameNum) 48 capture.set(cv2.CAP_PROP_POS_FRAMES, args.firstFrameNum)
52 if args.lastFrameNum is None: 49 if args.lastFrameNum is None:
53 lastFrameNum = float('inf') 50 lastFrameNum = float('inf')
54 else: 51 else:
55 lastFrameNum = args.lastFrameNum 52 lastFrameNum = args.lastFrameNum
56 nZerosFilename = int(ceil(log10(lastFrameNum))) 53 nZerosFilename = int(ceil(log10(lastFrameNum)))
57 while ret and frameNum < lastFrameNum: 54 while ret and frameNum < lastFrameNum:
58 ret, img = capture.read() 55 ret, img = capture.read()
59 if ret: 56 if ret:
60 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) 57 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
61 cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img) 58 cv2.imwrite(str(destinationPath/Path('undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum))), img)
62 if args.maskFilename is not None: 59 if args.maskFilename is not None:
63 cv2.imwrite(destinationDirname+'undistorted+mask-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), cv2.multiply(img, undistortedMask)) 60 cv2.imwrite(str(destinationPath/Path('undistorted+mask-{{:0{}}}.png'.format(nZerosFilename).format(frameNum))), cv2.multiply(img, undistortedMask, dtype = 16))
64 frameNum += 1 61 frameNum += 1
65 62
66 if args.encodeVideo: 63 if args.encodeVideo:
67 print('Encoding the images files in video') 64 print('Encoding the images files in video')
68 from subprocess import check_call 65 from subprocess import check_call
69 from storage import openCheck 66 from storage import openCheck
70 out = openCheck("err.log", "w") 67 out = openCheck("err.log", "w")
71 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) 68 check_call("mencoder \'mf://"+destinationPath+"*.png\' -mf fps={}:type=png -ovc xvid -xvidencopts bitrate={} -nosound -o ".format(args.fps, args.bitrate)+destinationDirname+"undistort.avi", stderr = out, shell = True)
72 out.close() 69 out.close()