custom_undistort_v4

import numpy as np
import time

def _extract_camera_matrix(camera_matrix):
    fx = camera_matrix[0,0]
    cx = camera_matrix[0,2]
    fy = camera_matrix[1,1]
    cy = camera_matrix[1,2]
    return fx, cx, fy, cy

def _extract_distortion_coefficients(k1, k2, k3, k4, k5, k6, p1, p2):
    return k1, k2, k3, k4, k5, k6, p1, p2

def _get_undistort_rectify_maps(distortion_coefficients, camera_matrix, image_width, image_height, step_size = 20):
    start = time.time()

    # Prepare coordinates
    coords_x = np.arange(image_width, dtype=np.float32)
    coords_y = np.arange(image_height, dtype=np.float32)

    # Apply inverse camera matrix
    camera_matrix_inv = np.linalg.inv(camera_matrix)
    fx_inv, cx_inv, fy_inv, cy_inv = _extract_camera_matrix(camera_matrix_inv)

    x = fx_inv * coords_x + cx_inv
    y = fy_inv * coords_y + cy_inv

    # Precalculate some common things
    xx = np.multiply(x, x)
    two_xx = 2 * xx

    yy = np.multiply(y, y)

    # Preallocate result arrays
    x_distorted = np.zeros((image_height, image_width), np.float32)
    y_distorted = np.zeros((image_height, image_width), np.float32)

    # Extract the various coefficients
    k1, k2, k3, k4, k5, k6, p1, p2 = _extract_distortion_coefficients(**distortion_coefficients)
    fx, cx, fy, cy = _extract_camera_matrix(camera_matrix)

    # Precalculate (minor, but still useful)
    two_p1, two_p2 = (2 * p1), (2 * p2)

    for i in range(0, image_height, step_size):
        # Handle case when last chunk is shorter then step_size
        i_e = min(image_height, i + step_size)

        yv = y[i:i_e,np.newaxis]
        yyv = yy[i:i_e,np.newaxis]

        xy = x * yv

        r2 = np.add(xx, yyv)
        r4 = np.multiply(r2, r2)
        r6 = np.multiply(r4, r2)

        radial_a = (1 + k1 * r2 + k2 * r4 + k3 * r6)
        radial_b = (1 + k4 * r2 + k5 * r4 + k6 * r6)
        radial = radial_a / radial_b

        tangential_x = two_p1 * xy + p2 * (r2 + two_xx)
        tangential_y = p1 * (r2 + 2 * yyv) + two_p2 * xy

        x_distorted[i:i_e,:] = fx * (x * radial + tangential_x) + cx
        y_distorted[i:i_e,:] = fy * (yv * radial + tangential_y) + cy

    end = time.time()
    print("END: process", end - start)

    return (x_distorted, y_distorted)


if __name__ == "__main__":
    image_width = 4032
    image_height = 3024

    distortion_coefficients = {
        "k1": np.float32(0)
        , "k2": np.float32(0)
        , "k3": np.float32(0)
        , "k4": np.float32(0)
        , "k5": np.float32(0)
        , "k6": np.float32(0)
        , "p1": np.float32(0)
        , "p2": np.float32(0)
        }

    camera_matrix = np.array([
        [1000, 0, 2016],
        [0, 1000, 1512],
        [0, 0, 1]], dtype=np.float32)

    start = time.time()
    map_x, map_y = _get_undistort_rectify_maps(
        distortion_coefficients
        , camera_matrix
        , image_width
        , image_height
        , 1)
    end = time.time()
    print("END: _get_undistort_rectify_maps", end - start)

    #np.savetxt("custund_mapx.csv", map_x, delimiter=",")
    #np.savetxt("custund_mapy.csv", map_y, delimiter=",")