custom_undistort_v2

import numpy as np
import time

def _distort_z_1_impl(x, y, k1, k2, k3, k4, k5, k6, p1, p2):
    xv = x
    yv = y

    x2 = np.multiply(x, x)
    y2 = np.multiply(y, y)
    xy = np.multiply(x, y)

    r2 = np.add(x2, y2)
    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 = (2 * p1) * xy + p2 * (r2 + 2 * x2)
    tangential_y = p1 * (r2 + 2 * y2) + (2 * p2) * xy

    x_distorted = x * radial + tangential_x
    y_distorted = y * radial + tangential_y

    return x_distorted, y_distorted

def _distort_z_1(x, y, step_size, k1, k2, k3, k4, k5, k6, p1, p2):
    x_distorted = np.zeros_like(x)
    y_distorted = np.zeros_like(y)

    for r in range(0, x.shape[0], step_size):
        rend = min(x.shape[0], r + step_size)
        x_distorted[r:rend,:], y_distorted[r:rend,:] = _distort_z_1_impl(x[r:rend,:], y[r:rend,:], k1, k2, k3, k4, k5, k6, p1, p2)

    return x_distorted, y_distorted


# Change dimension from [2 x H x W] to [H x W x 3 x 1] to correctly multiply with [3 x 3] matrix
def _homogeneous_reshape(points_x, points_y):
    points_homogeneous_reshaped = (
        # Add extra axis to change from [H x W x 3] to [H x W x 3 x 1]
        np.expand_dims(
            # Change from [3 x H x W] to [H x W x 3]
            np.transpose(
                # Change from [2 x H x W] to [3 x H x W] (homogeneous coordinates)
                np.stack(
                    np.broadcast_arrays(points_x, points_y, np.float32(1))),
                (1, 2, 0)),
            -1))

    return points_homogeneous_reshaped


def _homogeneous_reshape_back(points_homogeneous_reshaped):
    points_homogeneous = (
        # Get back from [H x W x 3] to [3 x H x W]
        np.transpose(
            # Remove extra axis: [H x W x 3 x 1] to [H x W x 3]
            np.squeeze(
                points_homogeneous_reshaped),
            (2, 0, 1)))

    # Get back from homogeneous coordinates
    points_x, points_y, _ = points_homogeneous

    return points_x, points_y


def _get_undistort_rectify_maps(distortion_coefficients, camera_matrix, image_width, image_height, step_size = 20):
    image_points = np.meshgrid(np.arange(image_width, dtype=np.float32), np.arange(image_height, dtype=np.float32))

    # print("BEGIN: _homogeneous_reshape")
    start = time.time()
    image_points_homogeneous_reshaped = _homogeneous_reshape(*image_points)
    end = time.time()
    print("END: _homogeneous_reshape", end - start)

    camera_matrix_inv = np.linalg.inv(camera_matrix)

    # print("BEGIN: camera_matrix_inv @ image_points_homogeneous_reshaped")
    start = time.time()
    image_points_homogeneous_z_1_reshaped = np.matmul(camera_matrix_inv, image_points_homogeneous_reshaped)
    end = time.time()
    print("END: camera_matrix_inv @ image_points_homogeneous_reshaped", end - start)

    # print("BEGIN: _homogeneous_reshape_back")
    start = time.time()
    image_points_z_1 = _homogeneous_reshape_back(image_points_homogeneous_z_1_reshaped)
    end = time.time()
    print("END: _homogeneous_reshape_back", end - start)

    # print("BEGIN: _distort_z_1")
    start = time.time()
    x_distorted, y_distorted = _distort_z_1(
        *image_points_z_1
        , step_size=step_size
        , **distortion_coefficients)
    end = time.time()
    print("END: _distort_z_1", end - start)

    # print("BEGIN: _homogeneous_reshape")
    start = time.time()
    points_homogeneous_z_1_distorted_reshaped = _homogeneous_reshape(x_distorted, y_distorted)
    end = time.time()
    print("END: _homogeneous_reshape", end - start)

    # print("BEGIN: _homogeneous_reshape")
    start = time.time()
    points_homogeneous_distorted_reshaped = np.matmul(camera_matrix, points_homogeneous_z_1_distorted_reshaped)
    end = time.time()
    print("END: camera_matrix @ points_homogeneous_z_1_distorted_reshaped", end - start)

    # print("BEGIN: _homogeneous_reshape_back")
    start = time.time()
    points_homogeneous_distorted = _homogeneous_reshape_back(points_homogeneous_distorted_reshaped)
    end = time.time()
    print("END: _homogeneous_reshape_back", end - start)

    return (map.astype(np.float32) for map in points_homogeneous_distorted)


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

    distortion_coefficients = {
        "k1": 0, "k2": 0, "k3": 0, "k4": 0, "k5": 0, "k6": 0,
        "p1": 0, "p2": 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=",")