在Numpy中加快nd数组计算

时间:2019-11-07 12:23:18

标签: python arrays numpy opencv optimization

我想在Python上使用undistortion模块来实现自定义OpenCV函数,就像在numpy中一样。

documentation可以知道undistort函数只是initUndistortRectifyMap()remap()的组合。

由于remap()是非常简单的操作,所以主要问题是为remap()实现映射。

我写了一个代码来构造地图,但是在我看来它的运行速度很慢。

代码由三个主要部分组成:

  1. 将原始图像重塑为一个形状良好的阵列,以将其与相机矩阵的逆相乘并相乘。
  2. z = 1平面中的扭曲点。
  3. 再次重塑点以执行另一次乘法以返回到图像点。

我拍了一张(4032 x 3024)大小的图像。

一次矩阵乘法在我的PC上工作大约1秒钟。失真功能的工作时间约为2.4秒。

我尝试将相同形状的矩阵与OpenCV上的C++ Mats相乘,并花费了0.0002秒。

问题是如何加快计算速度,因为在我看来,由于如此大的差异,我做错了事。

I found here an advice to make all arrays contiguous, but this did not help

代码:

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)

0 个答案:

没有答案