Converting a numpy image array based on a boolean mask

Question

I have 2 numpy arrays. One is a 3D integer array (image RGB values) with dimensions (988, 790, 3) and the other is a mask boolean array with the same shape. I want to use the mask to convert False values in the image array to black and leave true values as is.

I tried (image & mask) which appears to convert the entire image to black (or white) instead of just the False locations. I want to avoid loops for efficiency so looking for a numpy solution.

Answer 1

Since True and False values are treated as 1 and 0 respectively, you can simply use element-wise multiplication to get your desired result:

image * mask

Answer 2

You could use np.where. If your mask is also 3D, then np.where(mask, image, 0) should do the trick. But suspecting your mask is 2D...

import numpy as np

H, W = 2, 4

image = np.arange(1, 25).reshape(H, W, 3)
mask = np.random.choice([True, False], size=(H, W))

out = np.where(mask[..., None], image, 0)

Note that in the above, I'm assuming mask only has 2 dimensions; it does not have a trailing channel dimension. In order for np.where to work, its values need to be broadcastable, which I facilitated above by insert a trailing axis with shape 1 to the mask.

PS: if you calculate your mask based on your image, you may be able to use keyword argument keepdims=True somewhere in your calculations to get a mask where there is already such a trailing axis; which could clean up the code a bit.

results:

>>> image
array([[[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9],
        [10, 11, 12]],

       [[13, 14, 15],
        [16, 17, 18],
        [19, 20, 21],
        [22, 23, 24]]])
>>> mask
array([[False,  True, False,  True],
       [ True, False,  True, False]])
>>> out
array([[[ 0,  0,  0],
        [ 4,  5,  6],
        [ 0,  0,  0],
        [10, 11, 12]],

       [[13, 14, 15],
        [ 0,  0,  0],
        [19, 20, 21],
        [ 0,  0,  0]]])