undistortion.py

import cv2
import os
import numpy as np


mtx = [ 7.8612894299052164e+02, 0., 7.3830147560671890e+02, 0.,
       7.8839089780230699e+02, 6.5433443157976046e+02, 0., 0., 1. ]
mtx = np.array( mtx ).reshape((3, 3))

dist =  [ -3.3554454300236025e-01, 1.3238788177666494e-01,
       -4.2860590649250341e-05, 4.3244374088120965e-04,
       -2.5863792676102677e-02 ]
dist = np.array(dist)


folder = 'drive3_trafficlights'
img_name = '00235.jpg'
img = cv2.imread(os.path.join(folder, img_name))
h,  w = img.shape[:2]
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1, (w,h))

images = os.listdir(folder)
for image in images:
    if image.endswith('.jpg'):
        image_path = os.path.join(folder, image)
        img = cv2.imread(image_path)

        # Undistorting
        dst = cv2.undistort(img, mtx, dist, None, newcameramtx)

        # Cropping the image
        x,y,w,h = roi
        dst = dst[y:+y+h, x:x+w]
        cv2.imwrite(image, dst)
        break