Replacing ones and zeros in a 2D numpy array with another array?-CodePudding

I have a simple problem that I am trying to solve using numpy in an efficient manner. The jist of it is that I have a simple 2D array containing ones and zeros representing an image mask.

What I want to do is convert these ones and zeros into their RGB equivalent where one is a white pixel [255, 255, 255] and zero is a black pixel [0, 0, 0].

How would I go about doing this using NumPy?

mask = [[0, 0, 1],
[1, 0, 0]]

# something

result = [
[[0, 0, 0], [0, 0, 0], [255, 255, 255]],
[[255, 255, 255], [0, 0, 0], [0, 0, 0]]
]

The intent is to take the result and feed it into PIL to save into a PNG.

I've tried using numpy.where but can't seem to coax it into broadcasting another array out.

CodePudding user response：

A possible solution:

np.stack([255 * mask, 255 * mask, 255 * mask], axis=2)

Output:

array([[[  0,   0,   0],
        [  0,   0,   0],
        [255, 255, 255]],

       [[255, 255, 255],
        [  0,   0,   0],
        [  0,   0,   0]]])

CodePudding user response：

Since you need to repeat each item three times, np.repeat in conjunction with reshape could be used:

mask = np.array([[0, 0, 1], [1, 0, 0]])
255 * np.repeat(mask, 3, axis=1).reshape(*mask.shape, -1)
>>> array([[[  0,   0,   0],
            [  0,   0,   0],
            [255, 255, 255]],
           [[255, 255, 255],
            [  0,   0,   0],
            [  0,   0,   0]]])