Mercurial > hg > nsaunier > traffic-intelligence

comparison scripts/undistort-video.py @ 926:dbd81710d515

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
new feature tracking in image space with point undistortion
author Nicolas Saunier <nicolas.saunier@polymtl.ca>
date Mon, 10 Jul 2017 18:04:41 -0400
parents a71455bd8367
children c03d2c0a4c04
comparison
equal deleted inserted replaced
925:974077e23804 926:dbd81710d515
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)
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('-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)
20 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)
21 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')
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('--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('--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) 25 parser.add_argument('--bitrate', dest = 'bitrate', help = 'bitrate of the output video file if encoding', type = int, default = 5000)
25 26
26 args = parser.parse_args() 27 args = parser.parse_args()
27 28
28 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename) 29 intrinsicCameraMatrix = np.loadtxt(args.intrinsicCameraMatrixFilename)
29 #distortionCoefficients = args.distortionCoefficients
30 #undistortedImageMultiplication = args.undistortedImageMultiplication
31 #firstFrameNum = params.firstFrameNum
32 if args.destinationDirname is None: 30 if args.destinationDirname is None:
33 destinationDirname = '' 31 destinationDirname = ''
34 else: 32 else:
35 if not args.destinationDirname.endswith('/'): 33 if not args.destinationDirname.endswith('/'):
36 destinationDirname = args.destinationDirname+'/' 34 destinationDirname = args.destinationDirname+'/'
41 39
42 capture = cv2.VideoCapture(args.videoFilename) 40 capture = cv2.VideoCapture(args.videoFilename)
43 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) 41 width = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
44 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) 42 height = int(capture.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
45 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients) 43 [map1, map2] = cvutils.computeUndistortMaps(width, height, args.undistortedImageMultiplication, intrinsicCameraMatrix, args.distortionCoefficients)
44 if args.maskFilename is not None:
45 mask = cv2.imread(args.maskFilename)
46 undistortedMask = cv2.remap(mask, map1, map2, interpolation=cv2.INTER_LINEAR)/255
46 47
47 if capture.isOpened(): 48 if capture.isOpened():
48 ret = True 49 ret = True
49 frameNum = args.firstFrameNum 50 frameNum = args.firstFrameNum
50 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, args.firstFrameNum) 51 capture.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, args.firstFrameNum)
57 while ret and frameNum < lastFrameNum: 58 while ret and frameNum < lastFrameNum:
58 ret, img = capture.read() 59 ret, img = capture.read()
59 if ret: 60 if ret:
60 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR) 61 img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR)
61 cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img) 62 cv2.imwrite(destinationDirname+'undistorted-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), img)
63 if args.maskFilename is not None:
64 cv2.imwrite(destinationDirname+'undistorted+mask-{{:0{}}}.png'.format(nZerosFilename).format(frameNum), cv2.multiply(img, undistortedMask))
62 frameNum += 1 65 frameNum += 1
63 66
64 if args.encodeVideo: 67 if args.encodeVideo:
65 print('Encoding the images files in video') 68 print('Encoding the images files in video')
66 from subprocess import check_call 69 from subprocess import check_call