[SOLVED] Is there a way to modify an element in a Numpy array based on the value of other elements?

Is there a way to modify an element in a Numpy array based on the value of other elements?

Say I opened an image as a Numpy array and wanted to change the transparency of each pixel based on whether or not any of the RGB values for that pixel fall below a certain threshold.

For example, given the following array:

[[[255 255 255 255] 
  [247 248 251 255] 
  [243 247 247 255]]
 [[ 29  26  73 255]
  [ 36  26  68 255]
  [ 88  78 118 255]]]

With a threshold of 245, the expected output would be:

[[[255 255 255 255] 
  [247 248 251 255] 
  [243 247 247   0]]
 [[ 29  26  73   0]
  [ 36  26  68   0]
  [ 88  78 118   0]]]

I have very little experience (knowledge) with Numpy. I've been digging through the docs and here on SO, but haven't found anything that works like that. All of the examples I've seen after searching all day are about changing the value being checked rather than a different value based on that check.

So, as the question states, is there a way to do this efficiently in Numpy or should I just look for another way to do it?

Solution

Yes, you could achieve it with following code:

import numpy as np

image_array = np.array([[[255, 255, 255, 255], 
                         [247, 248, 251, 255], 
                         [243, 247, 247, 255]],
                        [[ 29,  26,  73, 255],
                         [ 36,  26,  68, 255],
                         [ 88,  78, 118, 255]]])

threshold = 245

mask = np.any(image_array[:, :, :3] < threshold,axis=2)
image_array[mask, 3] = 0
print(image_array)

np.any(image_array[:, :, :3] < threshold,axis=2) creates mask that check each pixel RGB values.

image_array[mask, 3] = 0 uses the mask to set alpha channel of pixels that meet condition to 0