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

Question

I am trying to apply a computation from 2D alpha image to 3 channeled RGB image. I need to update pixel intensity in each channel, based on respective pixel value in 2D alpha image. Below is one MWE I created to illustrate the concept.

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)

output:

I basically tried to find the indices in 2D alpha image, where condition is met, and tried to apply those indices to all channels of the image.

Is there a way to optimize above operation (not for channel looping)? I am specifically looking to avoid the for loop in the code, doing numpy for each index. It is ver slow for regular images.

Answer 1

You might observe that:
test_output = np.einsum('ij,ijk->ijk', test_a1, test_ir1)
Not sure if this helps reformulate a slightly different MWE if this isnt precisely what you intend

=== edited ===

I would still make use of einsum since it allows you a lot of control of vectorised multidimensional linear algebra.

Provided you can reduce your operations to some mathematical description, for example:

Where test_a1 is greater than zero, double the intensity of the pixels measured over each channel.

Then you do that in the following way:

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