如何提高在 python 中使用循环处理 numpy 数组的计算效率？

Question

我正在尝试在 python 中使用 numpy 执行一些图像转换任务。想法是，假设我已经将一个图像文件加载到一个numpy数组img中，然后我创建一个新数组new_img，并且还定义了像素坐标之间的映射：新图像中的[x,y]对应于@ 987654325@。然后我使用new_img[x,y] = img[old_x,old_y] 来计算转换。实际计算这个转换的循环看起来像这样（不是真正可运行的，因为我省略了很多关于转换规则、图像的宽度和长度、边界检查等的细节，但你明白了）强>.

def get_old_coord(y,x):
  # this is the function to compute the corresponding pixel coordinates
  # some computation here yields old_x and old_y
  # ...
  return old_x,old_y

for x in range(height_of_new_img):
  for y in range(width_of_new_img):
    new_img[y,x] = img[get_old_coord(x,y)]

# new_img is then as desired.

我遇到的问题是双循环非常耗时。 1000x1000 的图像需要一两分钟。另一方面，由于转换规则get_old_coord 可以是很多东西，我不认为我可以通过使用数组算术的一些内置函数来改进这一点。我怎样才能让这个过程更有效率？

---

更新：对于那些想要一个完整示例的人，这里有一个

import math
import torch
import torch.nn
import torchvision
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

def get_old_coord(x,y,old_h,old_w,new_h,new_w):
    cent_x, cent_y = new_h/2, new_w/2
    if ((x+1)%200==0 and (y+1)%new_w==0) or (x+1==new_h and (y+1)%new_w==0):
        print('Progress = %.2f %%'%((x+1)/new_h*100))
    rel_x, rel_y = x-cent_x, y-cent_y
    #print('Doint %d,%d, center = %d,%d, rel_x, rel_y = %d,%d'%(x,y,cent_x,cent_y,rel_x,rel_y))
    rho = math.sqrt(rel_x**2+rel_y**2)
    rho /= (min(new_h,new_w)/2)
    if rel_x==0 and rel_y>=0:
        theta = math.pi/2
    if rel_x==0 and rel_y<0:
        theta = -1*math.pi/2
    if rel_x>0:
        theta = math.atan(rel_y/rel_x)
    if rel_x<0:
        theta = math.atan(rel_y/rel_x)+math.pi
    theta = 2*math.pi-theta
    rho = 1-rho
    old_x = (int)(rho*old_h)
    old_y = (int)(theta*old_w/(2*math.pi))%old_w
    old_x = min(old_h-1,max(0,old_x))
    #old_y = min(old_w-1,max(0,old_y))
    #print('rho = %f, theta = %f, old_x, old_y = %d,%d'%(rho,theta,old_x,old_y))
    return old_x, old_y

def transform(new_h,new_w,old_im):
    old_h, old_w, _ = old_im.size()
    new_im = torch.zeros(new_h,new_w,3).int()

    for i in range(new_h):
        for j in range(new_w):
            new_im[i,j] = old_im[get_old_coord(i,j,old_h,old_w,new_h,new_w)]

    return new_im

input_file_name = 'in.jpg'
output_file_name = 'out.jpg'

new_h = 800
new_w = 800

old_im = torch.tensor(plt.imread(input_file_name))
new_im = transform_radial(new_h,new_w,old_im)

new_im = Image.fromarray(np.uint8(new_im))
new_im.save(output_file_name)

基本上，这实现了极坐标变换。很抱歉还没有时间将 cmets 添加到我的代码中。但是在您的任何图像上尝试它应该很有趣。另外，我在这个例子中实际上使用了 Torch 张量，但它们与 numpy 数组并没有太大区别。

Answer 1

您基本上需要的是矢量化，但可以访问正在对其执行元素明智操作的数组的索引。

对于 NxM 数组，我找到了一种方法

N,M = 4,6 # random number of dimensions

foo_mat = np.random.rand(N,M) #<-- random array representing data

def bar_func(elem, i, j): # <-- i,j are the indices of the matrix 
    return pass #<- Do whatever you want here

bar_func = np.vectorize(bar_func, signature="(n,m),(n,m),(n,m)->(n,m)")

baz_solved_mat = bar_func(foo_mat, *np.mgrid[:N,:M])

它是如何工作的？

bar_func = np.vectorize(bar_func, signature="(n,m),(n,m),(n,m)->(n,m)")

这里，签名说，对于 (n,m) 形状的元素值矩阵、(n,m) 形状的 i 值和 (n,m) 形状的 j 值（bar_func 的 3 个参数）返回 @987654325 @整形值。这将 3 个矩阵的每个 i,j 坐标映射在一起。

np.mgrid 是获取这些行和列索引数组的快速方法。 see doc