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.

`a_image[mask, :] = b_image[mask, :]`? – DYZ Aug 25 '20 at 23:26 — DYZ, Aug 25 '20 at 23:26
Are you sure your code does what you think it does...? – gmds Aug 25 '20 at 23:36 — gmds, Aug 25 '20 at 23:36

score 2 · Accepted Answer · answered Aug 26 '20 at 01:25

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]]]

Numpy better ways of iterating over colour channels using binary mask

1 Answers1