Mercurial Hosting > traffic-intelligence

annotate scripts/undistort-video.py @ 947:053484e08947

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
found a more elegant solution, making a copy of the list to iterate
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Fri, 21 Jul 2017 11:31:42 -0400
parents dbd81710d515
children c03d2c0a4c04
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
1 #! /usr/bin/env python
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
2
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
3 import sys, argparse
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
4
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
5 import numpy as np
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
6 import cv2
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
7
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
8 import cvutils
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
9 from math import ceil, log10
760
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
10 from os import path, mkdir
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
11
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
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
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
13 $ mencoder 'mf://./*.png' -mf fps=[framerate]:type=png -ovc xvid -xvidencopts bitrate=[bitrate] -nosound -o [output.avi]''')
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
14
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
15 parser.add_argument('-i', dest = 'videoFilename', help = 'filename of the video sequence')
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
16 parser.add_argument('--intrinsic', dest = 'intrinsicCameraMatrixFilename', help = 'name of the intrinsic camera file')
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
17 parser.add_argument('--distortion-coefficients', dest = 'distortionCoefficients', help = 'distortion coefficients', nargs = '*', type = float)
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
18 parser.add_argument('--undistorted-multiplication', dest = 'undistortedImageMultiplication', help = 'undistorted image multiplication', type = float)
926
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
19 parser.add_argument('--mask', dest = 'maskFilename', help = 'name of the mask file, to undistort to see how it covers the undistortion errors')
807
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
20 parser.add_argument('-f', dest = 'firstFrameNum', help = 'number of first frame number to display', type = int, default = 0)
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
21 parser.add_argument('-l', dest = 'lastFrameNum', help = 'number of last frame number to save', type = int)
760
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
22 parser.add_argument('-d', dest = 'destinationDirname', help = 'name of the directory where the undistorted frames are saved')
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
23 parser.add_argument('--encode', dest = 'encodeVideo', help = 'indicate if video is generated at the end (default Xvid)', action = 'store_true')
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
24 parser.add_argument('--fps', dest = 'fps', help = 'frame per second of the output video file if encoding', type = float, default = 30)
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
25 parser.add_argument('--bitrate', dest = 'bitrate', help = 'bitrate of the output video file if encoding', type = int, default = 5000)
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
26
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
27 args = parser.parse_args()
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
28
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
29 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename)
760
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
30 if args.destinationDirname is None:
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
31 destinationDirname = ''
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
32 else:
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
33 if not args.destinationDirname.endswith('/'):
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
34 destinationDirname = args.destinationDirname+'/'
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
35 else:
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
36 destinationDirname = args.destinationDirname
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
37 if not path.exists(destinationDirname):
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
38 mkdir(destinationDirname)
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
39
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
40 capture = cv2.VideoCapture(args.videoFilename)
807
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
41 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
42 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
43 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients)
926
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
44 if args.maskFilename is not None:
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
45 mask = cv2.imread(args.maskFilename)
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
46 undistortedMask = cv2.remap(mask, map1, map2, interpolation=cv2.INTER_LINEAR)/255
924
a71455bd8367 work in progress on undistortion acceleration
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 807
diff changeset
47
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
48 if capture.isOpened():
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
49 ret = True
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
50 frameNum = args.firstFrameNum
807
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
51 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, args.firstFrameNum)
636
3058e00887bc removed all issues because of tests with None, using is instead of == or !=
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 572
diff changeset
52 if args.lastFrameNum is None:
572
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
53 from sys import maxint
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
54 lastFrameNum = maxint
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
55 else:
9c429c7efe89 added script to generate undistorted images from video
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents:
diff changeset
56 lastFrameNum = args.lastFrameNum
807
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
57 nZerosFilename = int(ceil(log10(lastFrameNum)))
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
58 while ret and frameNum < lastFrameNum:
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
59 ret, img = capture.read()
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
60 if ret:
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
61 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
62 cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
926
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
63 if args.maskFilename is not None:
dbd81710d515 new feature tracking in image space with point undistortion
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 924
diff changeset
64 cv2.imwrite(destinationDirname+'undistorted+mask-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), cv2.multiply(img, undistortedMask))
807
52aa03260f03 reversed all code to OpenCV 2.4.13
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 770
diff changeset
65 frameNum += 1
760
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
66
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
67 if args.encodeVideo:
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
68 print('Encoding the images files in video')
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
69 from subprocess import check_call
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
70 from storage import openCheck
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
71 out = openCheck("err.log", "w")
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
72 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)
d72e4bcc1e36 added functionality to generate avi from the undistorted videos
Nicolas Saunier <nicolas.saunier@polymtl.ca>
parents: 636
diff changeset
73 out.close()