custom_undistort_v3

import numpy as np
import time


def _distort_z_1(x, y, step_size, k1, k2, k3, k4, k5, k6, p1, p2):
    col_count = x.shape[0]
    row_count = y.shape[0]

    xx = np.multiply(x, x)
    yy = np.multiply(y, y)

    two_p1 = 2 * p1
    two_p2 = 2 * p2
    two_xx = 2 * xx

    x_distorted = np.zeros((row_count, col_count), np.float32)
    y_distorted = np.zeros((row_count, col_count), np.float32)

    for i in range(0, row_count, step_size):
        i_e = min(row_count, 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,:] = x * radial + tangential_x
        y_distorted[i:i_e,:] = yv * radial + tangential_y

    return x_distorted, y_distorted


def _get_undistort_rectify_maps(distortion_coefficients, camera_matrix, image_width, image_height, step_size = 20):
    start = time.time()
    coords_x = np.arange(image_width, dtype=np.float32)
    coords_y = np.arange(image_height, dtype=np.float32)

    camera_matrix_inv = np.linalg.inv(camera_matrix)
    fx_inv = camera_matrix_inv[0,0]
    cx_inv = camera_matrix_inv[0,2]
    fy_inv = camera_matrix_inv[1,1]
    cy_inv = camera_matrix_inv[1,2]

    coords_x[:] = fx_inv * coords_x + cx_inv
    coords_y[:] = fy_inv * coords_y + cy_inv

    end = time.time()
    print("END: camera_matrix_inv * (x,y) equivalent", end - start)


    start = time.time()
    x_distorted, y_distorted = _distort_z_1(
        coords_x
        , coords_y
        , step_size
        , **distortion_coefficients)
    end = time.time()
    print("END: _distort_z_1", end - start)


    start = time.time()
    fx = camera_matrix[0,0]
    cx = camera_matrix[0,2]
    fy = camera_matrix[1,1]
    cy = camera_matrix[1,2]

    result_x = fx * x_distorted + cx
    result_y = fy * y_distorted + cy
    end = time.time()
    print("END: camera_matrix * (x_d,y_d) equivalent", end - start)

    return (result_x, result_y)


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
        , 20)
    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=",")