用 3D 替换 2D numpy 数组（元素到向量）

Question

This question has probably been asked before so I'm happy to be pointed to the answer but I couldn't find it.

I have a 2D numpy array of True and False . Now I need to convert it into a black and white image (a 3D numpy array), that is, I need [0,0,0] in place of every False and [1,1,1] in place of every True . What's the best way to do this? For example,

Input:

[[False, True],
 [True, False]]

Output:
[[[0, 0, 0], [1, 1, 1]],
 [[1, 1, 1], [0, 0, 0]]]

(As you probably know, 3D images are arrays of shape (height, width, 3) where 3 is the depth dimension ie number of channels.)

Bonus points if someone can tell me how to also convert it back, ie, if I have a pure black and white image (purely [0,0,0] and [0,0,1] pixels), how do I get a 2D matrix of the same height-width dimensions but with True in place of white pixels ([1,1,1]) and False in place of black pixels ([0,0,0]).

Answer 1

The cheapest way is to view your bool data as np.uint8 , and add a fake dimension:

img = np.lib.stride_tricks.as_strided(mask.view(np.uint8),
                                      strides=mask.strides + (0,),
                                      shape=mask.shape + (3,))

Unlike mask.astype(np.uint8) , mask.view(np.uint8) does not copy the data, instead harnesses the fact that bool_ is stored in a single byte. Similarly, the new dimension created by np.lib.stride_tricks.as_strided is a view which does not copy any data.

You can bypass as_strided and view by creating a new array object manually:

img = np.ndarray(shape=mask.shape + (3,), dtype=np.uint8,
                 strides=mask.strides + (0,), buffer=mask)

Answer 2

I think the clearest way is this:

a = np.array([[False, True], [True, False]])
out = np.zeros((*a.shape, 3), dtype=np.uint8)
out[a.nonzero()] = 1

>>> out
array([[[0, 0, 0],
        [1, 1, 1]],

       [[1, 1, 1],
        [0, 0, 0]]], dtype=uint8)