如何在3通道RGB图像上应用2D条件来优化numpy操作？

Question

我正在尝试将2D alpha图像的计算应用于3通道RGB图像。我需要基于2D alpha图像中的相应像素值来更新每个通道中的像素强度。以下是我创建的一个MWE来说明这一概念。

MWE：

# test alpha 2D image
test_a1 = np.array([
    [0, 0, 50], 
    [0, 0, 150],
    [0, 0, 225]
    ])

# test 3 channel RGB image
test_ir1 = np.ones((3,3,3))

# getting indices of alpha where cond is satisfied
idx = np.unravel_index(np.where(test_a1.ravel()>0),test_a1.shape)
test_output = np.zeros_like(test_ir1)
n_idx = len(idx[0][0])

# applying computation on 3 channel RGB image only where cond is satisfied.
for i in range(n_idx):

    # multiply only where test_a1 > 0
    r_idx, c_idx = idx[0][0][i], idx[1][0][i]
    test_output[r_idx,c_idx,0] = test_a1[r_idx, c_idx] * test_ir1[r_idx, c_idx, 0]
    test_output[r_idx,c_idx,1] = test_a1[r_idx, c_idx] * test_ir1[r_idx, c_idx, 1]
    test_output[r_idx,c_idx,2] = test_a1[r_idx, c_idx] * test_ir1[r_idx, c_idx, 2]

test_output = test_output.astype('uint8')
plt.imshow(test_output, vmin=0, vmax=3)

输出：

我基本上尝试在满足条件的2D alpha图像中查找索引，并尝试将这些索引应用于图像的所有通道。

是否有一种方法可以优化上述操作（不适用于通道循环）？我特别希望避免代码中的for循环，对每个索引执行numpy。对于常规图像来说速度很慢。

Answer 1

您可能会发现：
test_output = np.einsum('ij,ijk->ijk', test_a1, test_ir1)
不确定这是否恰好不是您想要的目的，是否可以帮助重新安排稍有不同的MWE

===已编辑===

我仍然会使用einsum，因为它可以使您对矢量化的多维线性代数有很多控制。

只要您可以将操作简化为一些数学描述，例如：

  

如果test_a1大于零，则在每个通道上测得的像素强度加倍。

然后您可以通过以下方式进行操作：

mask = test_a1 > 0
output = np.einsum('ij,ijk->ijk', mask, test_ir1) + test_ir1