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

Question:

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.

Asked By: TheMaster

||

Answers:

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]]])
Answered By: PaulS

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]]])
Answered By: mathfux

As your image contains only two colours, I would suggest you consider saving it as a palette image, a.k.a. an indexed image.

Rather than needlessly inflating your image by a factor of 3 to enable it to store 16.7 million colours, you can just store one byte per pixel which will still enable you to have 256 colours which seems plenty when you only have 2 "colours", namely black and white.

That looks like this:

import numpy as np
from PIL import Image

# Make Numpy array "na" from your list
na = np.array(mask, dtype=np.uint8)

# Make PIL Image from Numpy array - this image will be 'L' mode
im = Image.fromarray(na)

# Now push a palette into the image that says:
#   index 0 => black, i.e. [0,0,0]
#   index 1 => white, i.e. [255,255,255]
#   all other 254 indices are black
# Afterwards the image will be 'P' mode
im.putpalette([0,0,0, 255,255,255] + [0,0,0]*254)

# Save
im.save('result.png')

enter image description here

Answered By: Mark Setchell