Numpy better ways of iterating over colour channels using binary mask

Question

Is there a better way to apply a binary mask to colour channels in numpy? I end up having to do this all the time and it feels like there should be.

for c in range(3):
    a_image[mask, c] = b_image[mask, c]

shapes are (x, y, c) for a_image and b_image, and (x, y) for mask.

Answer 1

You can simply use the 2-D mask on a 3-D array without a loop. Numpy will broadcast it to the third dimension for you.

a_image[mask] = b_image[mask]

simple example:

a_image = np.arange(6).reshape(1,2,3)
#[[[0 1 2]
#  [3 4 5]]]

b_image = np.ones((1,2,3))
#[[[1. 1. 1.]
#  [1. 1. 1.]]]

mask = np.array([[False,True]])
#[[False  True]]

output:

[[[0 1 2]
  [1 1 1]]]