CODEWITHSUNDEEP
pythonnumpy
Clamari
Clamari

Reputation: 353

Replace elements in array if a condition meet

I'm trying to optimize my function with numpy. I have an array[x, y, 4] which is a RGBA image (obtained from bgra = cv2.cvtColor(bgra, cv2.COLOR_RGB2RGBA))

I want to do something like this:

        for pixel_row in bgra:
            for pixel in pixel_row:
                if pixel[0] == 0:
                    pixel[3] = 0

or this:

        for pixel_row in bgra:
            for pixel in pixel_row:
                if np.array_equal([0,0,0,255], pixel):
                    pixel[3] = 0

but using numpy which is wide faster than manual iterations. I tried this

bgra = np.where(bgra == [0,0,0,255], [0,0,0,0], bgra)

but the result is not what I expected because each value of the pixel is compared and replaced individually. I want to make black pixels fully transparent.

Upvotes: 1

Views: 141

Answers (1)

a_guest
a_guest

Reputation: 36329

You can use .all(axis=2) to enforce the condition on the color channel axis. Then you can use np.where to update the color channel axis:

import numpy as np
  
img = np.zeros((3, 3, 4))
img[..., 3] = 255
img[1, 1] = [1, 2, 3, 255]

new = img.copy()
new[...,3] = np.where((img == [0,0,0,255]).all(axis=2), 0, img[...,3])

print(img, new, sep='\n\n')

The output is:

[[[  0.   0.   0. 255.]
  [  0.   0.   0. 255.]
  [  0.   0.   0. 255.]]

 [[  0.   0.   0. 255.]
  [  1.   2.   3. 255.]
  [  0.   0.   0. 255.]]

 [[  0.   0.   0. 255.]
  [  0.   0.   0. 255.]
  [  0.   0.   0. 255.]]]

[[[  0.   0.   0.   0.]
  [  0.   0.   0.   0.]
  [  0.   0.   0.   0.]]

 [[  0.   0.   0.   0.]
  [  1.   2.   3. 255.]
  [  0.   0.   0.   0.]]

 [[  0.   0.   0.   0.]
  [  0.   0.   0.   0.]
  [  0.   0.   0.   0.]]]

Upvotes: 2

Related Questions