Vectorized approach to masking red and blue channels in an image

I am trying to change all the red and blue pixels in an image to black to have only green pixels (based on certain conditions).

To do this right now I am using multiple for loops and the process, although it works, is extremely slow.

The code I have so far –

### Test image is my original RGB image
mat = np.asarray(testim)

for elemento in mat: 

    for pixel in element:

        if ((pixel[0] + pixel[2]) >= (2*pixel[1])): #Non-vegetation background changing to black according to the rule (R + B >= 2G)
          pixel[0] = 0
          pixel[1] = 0
          pixel[2] = 0
        
        elif pixel[1] > pixel[0] and pixel[1] > pixel[2] and pixel[1] > 25: # Treat these as green vegetation and do not change anything
          continue
        
        else: # Change background to black
          pixel[0] = 0
          pixel[1] = 0
          pixel[2] = 0
cv2_imshow(testim)
print(testim.shape)

Is there any way I can vectorize this using Numpy without using nested for loops to make the same process work faster? I am slightly new to NumPy operations and am confused as to how to get it done. I appreciate the help!

For example : My input image – [1]: https://i.stack.imgur.com/zWGtA.jpg

The output image I have now with the above logic – [2]: https://i.stack.imgur.com/VkmLC.jpg

I would like the same output with a faster code preferably using NumPy instead of the nested for loops that I currently have

Answer

You should be able to access your red matrix with mat[:,:,0], and similar with green and blue.

Your mask of green pixels is then:

mask = (mat[:,:,1]>mat[:,:,0]) & (mat[:,:,1]>mat[:,:,2]) & (mat[:,:,1]>25)

where you are doing element-wise comparisons and element-wise AND to combine comparisons in a boolean 2D matrix.

You can then zero each other pixel by doing:

# Repeat mask along the three channels
mask = np.repeat(mask[:,:,np.newaxis], 3, axis=2)

# Zero non-green elements
mat[~mask] = 0