Ускорьте вычисления в массиве в Numpy

Question

Я хочу реализовать пользовательскую функцию undistortion, как в OpenCV, используя модуль numpy на Python.

Из документации известно, что функция undistort является простокомбинация initUndistortRectifyMap() и remap().

Поскольку remap() - довольно простая операция, основной проблемой является реализация карт для remap().

Я написал код для построениякарты, но мне кажется, что он работает довольно медленно.

Код состоит из трех основных частей:

1. Изменение формы исходного изображения указывает на массив правильной формы, чтобы умножить его наинвертировать матрицу камеры и выполнить умножение.
2. Исказить точки в плоскости z = 1.
3. Снова изменить форму, чтобы выполнить другое умножение, чтобы вернуться к точкам изображения.

Я сделал снимок размером (4032 x 3024).

Умножение одной матрицы работает на моем компьютере около 1 секунды. И функция искажения работает около 2,4 сек.

Я пытался умножить матрицы одинаковой формы с OpenCV Матс на C++, и я взял 0,0002 сек.

Вопрос в том, какускорить вычисления, потому что мне кажется, что я делаю что-то не так из-за такой большой разницы.

Я нашел здесь совет, чтобы сделать все массивы смежными, но это не помогло

код:

import numpy
import time


def _distort_z_1(x, y, k1, k2, k3, k4, k5, k6, p1, p2):
    x2 = x * x
    y2 = y * y
    xy = x * y

    r2 = x2 + y2
    r4 = r2 * r2
    r6 = r4 * r2

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

    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


# 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]
        numpy.expand_dims(
            # Change from [3 x H x W] to [H x W x 3]
            numpy.transpose(
                # Change from [2 x H x W] to [3 x H x W] (homogeneous coordinates)
                numpy.stack(
                    numpy.broadcast_arrays(points_x, points_y, 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]
        numpy.transpose(
            # Remove extra axis: [H x W x 3 x 1] to [H x W x 3]
            numpy.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):
    image_points = numpy.meshgrid(range(image_width), range(image_height))

    # 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 = numpy.linalg.inv(camera_matrix)

    # print("BEGIN: camera_matrix_inv @ image_points_homogeneous_reshaped")
    start = time.time()
    image_points_homogeneous_z_1_reshaped = 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,
        **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 = 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(numpy.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 = numpy.array([
        [1000, 0, 2016],
        [0, 1000, 1512],
        [0, 0, 1]])

    map_x, map_y = _get_undistort_rectify_maps(
        distortion_coefficients,
        camera_matrix,
        image_width,
        image_height)