python - Numpy 向量化

Question

如何对以下 python 代码(for 循环)进行 numpy 向量化？任何帮助将不胜感激

arr1 = np.ndarray(shape = (184,184))
arr2 = np.ndarray(shape = (184,184))
arr3 = np.full((184, 184), 0.0, dtype=float)
for i in range(arr1.size[0]):
    for j in range(arr1.size[1]):
        if arr2[i,j] == 0 or arr1[i,j] == 0:
            arr3[i,j]=0
        elif arr2[i,j] == 255 and arr1[i,j] == 255:
            arr3[i, j] = 255

Answer 1

利用 masks and boolean-indexing 的向量化运算-

mask1 = (arr1==0) | (arr2==0)
mask2 = (arr1==255) & (arr2==255)

arr3[mask1] = 0
arr3[mask2] = 255

如果 arr3 已经用 zeros 初始化，我们可以跳过 arr3[mask1] 部分，因为那是分配 zeros 无论如何，因为没有其他条件语句，我们可以直接使用 mask2 获取 arr3，就像这样 -

arr3 = 255.0*mask2

用于验证的 sample 运行 -

In [23]: # Setup input
    ...: np.random.seed(0)
    ...: arr1 = (np.random.rand(184,184)>0.5)*255
    ...: arr2 = (np.random.rand(184,184)>0.5)*255

In [24]: # Run original code
    ...: arr3 = np.full((184, 184), 0.0, dtype=float)
    ...: for i in range(arr1.shape[0]):
    ...:     for j in range(arr1.shape[1]):
    ...:         if arr2[i,j] == 0 or arr1[i,j] == 0:
    ...:             arr3[i,j]=0
    ...:         elif arr2[i,j] == 255 and arr1[i,j] == 255:
    ...:             arr3[i, j] = 255

In [25]: # Run proposed code#1
    ...: out = np.full((184, 184), 0.0, dtype=float)
    ...: mask1 = (arr1==0) | (arr2==0)
    ...: mask2 = (arr1==255) & (arr2==255)
    ...: 
    ...: out[mask1] = 0
    ...: out[mask2] = 255

In [26]: np.allclose(arr3, out) #verify code#1
Out[26]: True

In [27]: np.allclose(arr3, 255.0*mask2) #verify code#2
Out[27]: True