如何在 python 中使用循環來提高處理 numpy 數組的計算效率？

Question

我正在嘗試用 python 中的 numpy 做一些圖像轉換任務。 這個想法是，假設我已經將一個圖像文件加載到一個 numpy 數組img ，然后我創建了一個新數組new_img ，並且還定義了像素坐標之間的映射：新圖像中的[x,y]對應於[old_x,old_y] 。 然后我使用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)

基本上，這實現了極坐標變換。 很抱歉沒有時間在我的代碼中添加注釋。 但是在您的任何圖像上嘗試它應該很有趣。 另外，我實際上在這個例子中使用了火炬張量，但它們與 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個參數）返回(n,m)形狀的值. 這將 3 個矩陣的每個 i,j 坐標映射在一起。

np.mgrid是獲取這些行和列索引數組的快速方法。 看文檔