如何加速图像数组/numpy数组的迭代

Question

我想在图像上应用 Arnold 猫图。我从Wikipedia 管理此代码，但是当我将600 的迭代次数增加600 图像和83 迭代时它变得很慢，它花了time=567.8921346664429 秒。可以对此代码进行任何改进。我猜循环的东西让这段代码很慢，O(n^2) 每次迭代。

from PIL.Image import open as load_pic, new as new_pic


def main(path, iterations, keep_all=False, name="arnold_cat-{name}-{index}.png"):
    """
    Params
        path:str
            path to photograph
        iterations:int
            number of iterations to compute
        name:str
            formattable string to use as template for file names
    """
    title = os.path.splitext(os.path.split(path)[1])[0]
    counter = 0
    while counter < iterations:
        with load_pic(path) as image:
            dim = width, height = image.size
            with new_pic(image.mode, dim) as canvas:
                for x in range(width):
                    for y in range(height):
                        nx = (2 * x + y) % width
                        ny = (x + y) % height

                        canvas.putpixel((nx, height-ny-1), image.getpixel((x, height-y-1)))

        if counter > 0 and not keep_all:
            os.remove(path)
        counter += 1
        print(counter, end="\r")
        path = name.format(name=title, index=counter)
        canvas.save(path)

    return canvas

我为我的用例修改代码，如下所示：

def cat_map(image_matrix,MAX):
    dim = width, height = image_matrix.shape
    transformed_matrix = np.zeros((width,height)) 
    # Apply Arnold cat map on image_matrix
    index = []
    counter = 0 
    iterations = MAX
    # collect initial index of the image_matrix
    for i in range(len(image_matrix)):
        for j in range(len(image_matrix[0])):
            index.append([i,j])
    forward_index = []
    while counter < iterations:
        for coordinate in index:
            x = coordinate[0]
            y = coordinate[1]
            new_x = (2 * x + y) % width
            new_y = (x + y) % height

            transformed_matrix[new_x][new_y] = image_matrix[x][y]
            forward_index.append([new_x,new_y])
        index = forward_index
        # apply recursive transformation on image_matrix
        image_matrix = transformed_matrix
        # only store the last index matrix
        if counter != iterations - 1:
            forward_index = []
            # re-initialize transformation matrix
            transformed_matrix = np.zeros((width,height)) 
        counter += 1
    return transformed_matrix,forward_index

我还需要最后一次迭代的index_array。我很欣赏@dankal444 的矢量化思想，但是如何存储索引数组？

Answer 1

矢量化版本大约快 40 到 80 倍（取决于迭代次数以及您是否愿意使用 opencv）。

多次迭代运行的最大加速来自转换 idx 矩阵（x_image、y_image、nx_image、ny_image）的存储和重用。

现在保存图像是一个瓶颈，因此如果您不需要中间图像，您可以将其注释掉并获得另一个显着的 (~x2-3) 加速。

def main(path, iterations, keep_all=False, name="arnold_cat-{name}-{index}.png"):
    """
    Params
        path:str
            path to photograph
        iterations:int
            number of iterations to compute
        name:str
            formattable string to use as template for file names
    """
    title = os.path.splitext(os.path.split(path)[1])[0]
    counter = 0

    with load_pic(path) as image:

        width, height = image.size
        current_image = np.array(image).copy()
        n_channels = current_image.shape[-1]
        x_image = np.repeat(np.arange(width).reshape(-1, 1), height, axis=-1).T
        y_image = np.repeat(np.arange(height).reshape(-1, 1), width, axis=-1)
        nx_image = (2 * x_image + y_image) % width
        ny_image = (x_image + y_image) % height
        transformed_image = np.zeros((width, height, n_channels)).astype(np.uint8)
        ny_image = height - ny_image - 1
        y_image = height - y_image - 1
        while counter < iterations:
            transformed_image[ny_image, nx_image] = current_image[y_image, x_image]

            if counter > 0 and not keep_all:
                os.remove(path)
            counter += 1
            print(counter, end="\r")
            path = name.format(name=title, index=counter)

            # slower saving image:
            # image = fromarray(transformed_image)
            # image.save(path)

            # this is faster alternative of saving image, without it it is still very fast
            import cv2
            cv2.imwrite(path, transformed_image)

            current_image = transformed_image.copy()

        # use canvas at the end, for me this is unnecessary:
        image = fromarray(current_image)

    return image

编辑：OP 也希望返回转换索引，因此重构代码以具有单独的函数来计算转换索引

def main(path, iterations, keep_all=False, name="arnold_cat-{name}-{index}.png"):
    """
    Params
        path:str
            path to photograph
        iterations:int
            number of iterations to compute
        name:str
            formattable string to use as template for file names
    """
    title = os.path.splitext(os.path.split(path)[1])[0]
    counter = 0

    with load_pic(path) as image:

        width, height = image.size
        current_image = np.array(image).copy()
        n_channels = current_image.shape[-1]
        transformation_indices = get_transformation_indices(counter, height, iterations, width)
        current_image = np.array(image)[transformation_indices[:, :, 0], transformation_indices[:, :, 1]]
        # use canvas at the end, for me this is unnecessary:
        image = fromarray(current_image)

    return image, transformation_indices


def get_transformation_indices(counter, height, iterations, width):
    x_image = np.repeat(np.arange(width).reshape(-1, 1), height, axis=-1).T
    y_image = np.repeat(np.arange(height).reshape(-1, 1), width, axis=-1)
    nx_image = (2 * x_image + y_image) % width
    ny_image = (x_image + y_image) % height
    transformation_indices = np.stack((y_image, x_image), axis=-1)
    ny_image = height - ny_image - 1
    y_image = height - y_image - 1
    while counter < iterations:
        transformation_indices[ny_image, nx_image] = transformation_indices[y_image, x_image]
        counter += 1
    return transformation_indices