Numpy masking in 3 channel array-CodePudding

The following Snippet will create a test image

# Create 3x3x3 image
test_image = []
for i in range(9):
    if i < 6:
        image.append([255, 22, 96])
    else:
        image.append([255, 0, 0])

Out:

array([[[255,  22,  96],
        [255,  22,  96],
        [255,  22,  96]],

       [[255,  22,  96],
        [255,  22,  96],
        [255,  22,  96]],

       [[255,   0,   0],
        [255,   0,   0],
        [255,   0,   0]]], dtype=int32)

My goal is to create a single-channel image of zeros BUT for every [255, 22, 96] in test_image, I want to set the number 100 in the new single_channel image:

I tried the following:

test_image = np.array(test_image)
height, width, channels = test_image.shape
single_channel_img = np.zeros(test_image.shape, dtype=int)

msk = test_image ==  [255, 22, 96] # DOES NOT WORK AS EXPECTED
single_channel_img[msk] = 100

Which does not work because the result of the masking:

msk = test_image ==  [255, 22, 96]

returns:

array([[[ True,  True,  True],
        [ True,  True,  True],
        [ True,  True,  True]],

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

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

Why does the masking return True for the last 3 Pixel and how can I make sure that that comparison returns True only if all 3 Values are the same. My assumption was that the way I mask should just return True when all 3 RGB values are equal to [255, 22, 96].

CodePudding user response：

You can convert msk to a 3-D array using array broadcasting: https://numpy.org/doc/stable/user/basics.broadcasting.html

The command .reshape can be used to change the dimensions of an array. Numpy will automatically fill out the "thin" dimension. So for example,comparing arrays with shapes (n,n,3) and(1,1,3) is the same as comparing each sub-array test_image[i,j,:] with the target (1,1,3).

import  numpy as np

# Create 3x3x3 image
test_image = []
for i in range(9):
    if i < 6:
        test_image.append([255, 22, 96])
    else:
        test_image.append([255, 0, 0])

test_image         = np.array(test_image).reshape((3,3,3)) # test image shape needed to be fixed
single_channel_img = np.zeros(test_image.shape, dtype=int)

msk = test_image == np.array([255,22,96]).reshape((1,1,3)) # now it works
single_channel_img[msk] = 100

print(single_channel_img)
# [[[100 100 100]
#   [100 100 100]
#   [100 100 100]]
# 
#  [[100 100 100]
#   [100 100 100]
#   [100 100 100]]
# 
#  [[100   0   0]
#   [100   0   0]
#   [100   0   0]]]

PS. PyTorch also has array broadcasting, it is really useful in deep learning.