Undistort Script

undistort.py

#!/usr/bin/env python

import sys
import argparse
import os

import cv
import cv2
import numpy

if __name__ == "__main__":
    distort = {'nexus5':(
                [[  2.04885461e+03,   0.00000000e+00,   1.03469869e+03],
                [  0.00000000e+00,   2.03764082e+03,   7.20988436e+02],
                [  0.00000000e+00,   0.00000000e+00,   1.00000000e+00]],
                [  7.66240319e-02,  1.82510928e-01, -8.76278643e-04,  6.98555155e-04, -1.48187088e+00]
            ),
               'gopro':(
                [[  1.84645694e+03,  0.00000000e+00,  1.98328678e+03],
                 [  0.00000000e+00,  1.84390571e+03,  1.59156265e+03],
                 [  0.00000000e+00,  0.00000000e+00,  1.00000000e+00]],
                [-0.29320943, 0.1426103, -0.00153012, 0.00053567,-0.05368968]
            )
         }

    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--model','-m')
    parser.add_argument('--images', '-i', nargs='+', help='List of images to undistort')
    parser.add_argument('--out', '-o', help='Ouput directory')
    
    args = parser.parse_args()
    
        
    for idx, fn in enumerate(args.images):
        img = cv2.imread(fn)
        A, d = distort[args.model]
        flat_img = cv2.undistort(img,numpy.array(A),numpy.array(d))
        out_path = os.path.join(args.out,os.path.basename(fn))
        if not os.path.exists(out_path):
            cv2.imwrite(out_path, flat_img)
        else:
            print("Skipping %s, file already exists." % out_path)