I have an array with shape (100, 80, 3) which is an rgb image.
I have a boolean mask with shape (100, 80).
I want each pixel where the mask is True to have value of pix_val = np.array([0.1, 0.2, 0.3]).
cols = 100
rows = 80
img = np.random.rand(rows, cols, 3)
mask = np.random.randint(2, size=(rows, cols), dtype=np.bool_)
px = np.array([0.1, 0.2, 0.3])
for ch in range(3):
img[:, :, ch][mask] = px[ch]
I thought broadcasting:
img[mask[:, :, None]] = px
would work. But it did not.
I am looking for a vectorized (efficient) way to implement it.
I'll attempt to explain why your indexing attempt didn't work.
Make a smaller 3d array, and 2d mask:
In [1]: import numpy as np
In [2]: img = np.arange(24).reshape(2,3,4)
In [3]: mask = np.array([[1,0,1],[0,1,1]],bool);mask
Out[3]:
array([[ True, False, True],
[False, True, True]])
Using @mozway's indexing, produces a (4,4) array. The first first 4 is the number of True values in the mask, the second is the trailing dimension:
In [4]: img[mask]
Out[4]:
array([[ 0, 1, 2, 3],
[ 8, 9, 10, 11],
[16, 17, 18, 19],
[20, 21, 22, 23]])
With your indexing attempt, we get an error. (You really should have shown the error message):
In [5]: img[mask[:,:,None]]
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[5], line 1
----> 1 img[mask[:,:,None]]
IndexError: boolean index did not match indexed array
along dimension 2; dimension is 4 but corresponding
boolean dimension is 1
With this None, mask dimension is (2,3,1). That last 1 doesn't match the 4 of img. broadcasting doesn't apply in this context.
Now if we attempt to use mask in a multiplication, the (2,3,4) and (2,3) don't broadcast together:
In [6]: img*mask
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[6], line 1
----> 1 img*mask
ValueError: operands could not be broadcast together
with shapes (2,3,4) (2,3)
But (2,3,1) does broadcast with (2,3,4), producing a select number of 0 rows:
In [7]: img*mask[:,:,None]
Out[7]:
array([[[ 0, 1, 2, 3],
[ 0, 0, 0, 0],
[ 8, 9, 10, 11]],
[[ 0, 0, 0, 0],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
As I commented, using a boolean mask is equivalent to indexing with nonzero arrays:
In [13]: I,J = np.nonzero(mask); I,J
Out[13]: (array([0, 0, 1, 1], dtype=int64), array([0, 2, 1, 2], dtype=int64))
In [14]: img[I,J,:]
Out[14]:
array([[ 0, 1, 2, 3],
[ 8, 9, 10, 11],
[16, 17, 18, 19],
[20, 21, 22, 23]])
In the assignment expresion, a size (4,) value can broadcast to the (n,4) indexed img[mask]. Now if we were attempting to mask other dimensions we might need to make a px[:,None,:] or something like that.