Modify numpy 3d array (shape x,y,3) based on value from another 2d array shape (x,y)

Question

Starting with a 3d array (like a 2d image with RGB). I'd like to change the color based on the value of another 2d matrix.

import numpy as np

img=np.zeros((2,2,3)) # a black image
print('\nimg=',list(img))

b=np.array([[1,2],[3,4]]) # some 2d array of values

#img=np.where(b==1,[9,9,9],img) # ValueError: operands could not be broadcast together with shapes (2,2) (3,) (2,2,3) 
#print(img)

# Trying to color the coordinate where b==1 with the RGB color 9,9,9
whatIwant=np.array([[9,9,9],[0,0,0],[0,0,0],[0,0,0]])
print('\nwhatIwant=\n',list(whatIwant))

expected output:

 img=[array([[0., 0., 0.],
   [0., 0., 0.]]), array([[0., 0., 0.],
   [0., 0., 0.]])]

 whatIwant=
 [array([9, 9, 9]), array([0, 0, 0]), array([0, 0, 0]), array([0, 0, 0])]

Answer 1

In [13]: img=np.zeros((2,2,3)) # a black image
In [14]: b=np.array([[1,2],[3,4]]) # some 2d array of values
    

boolean test array:

In [15]: b==1
Out[15]: 
array([[ True, False],
       [False, False]])

A boolean mask has to match all dimensions, or just one:

In [16]: img[b]
Traceback (most recent call last):
  File "<ipython-input-16-228af24ace6b>", line 1, in <module>
    img[b]
IndexError: index 2 is out of bounds for axis 0 with size 2

But if we get the indices of the True value(s):

In [17]: idx = np.nonzero(b==1)
In [18]: idx
Out[18]: (array([0]), array([0]))

we can use that to index the 3d array, for get or for set:

In [19]: img[idx]
Out[19]: array([[0., 0., 0.]])
In [20]: img[idx]=[9,8,7]
In [21]: img
Out[21]: 
array([[[9., 8., 7.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 0., 0.]]])

Sometimes it's easier to unpack the nonzero tuple:

In [22]: I,J = np.nonzero(b==1)
In [23]: I,J
Out[23]: (array([0]), array([0]))
In [24]: img[I,J,:]
Out[24]: array([[9., 8., 7.]])