CODEWITHSUNDEEP
pythonnumpy
kainev
kainev

Reputation: 265

Numpy: Efficiently set all pixels to black that meet a condition

I need to change all pixels that meet some conditions to black. My quick solution is to just iterate over each pixel. My application is performance sensitive and is slower than required. The image sizes that it is looping over are quite small, ranging from roughly 25x25 to 50x50.

I've provided an example of what I'm doing. Apologies if there are any mistakes in the image loading or width/height part.

image = cv2.imread("image.png")
width, height = image.shape

for x in range(width):
    for y in range(height):
        blue = image[x, y, 0]
        green = image[x, y, 1]
        red = image[x, y, 2]

        if (abs(green - blue) > 20 and green > 30) or abs(red - green) < 40:
            output[x, y] = [0, 0, 0]

Upvotes: 4

Views: 850

Answers (2)

hilberts_drinking_problem
hilberts_drinking_problem

Reputation: 11602

Incidentally, your code lends itself well to numba.jit, e.g.

from numba import jit

@jit
def change_color(image):
    width, height = image.shape[:2]
    image = image.copy()
    for x in range(width):
        for y in range(height):
            blue = image[x, y, 0]
            green = image[x, y, 1]
            red = image[x, y, 2]
            if (abs(green - blue) > 20 and green > 30) or abs(red - green) < 40:
                image[x, y, 0] = 0
                image[x, y, 1] = 0
                image[x, y, 2] = 0
    return image

This is 10-15 times faster than numpy on a 512x512 image I tested.

Upvotes: 2

Divakar
Divakar

Reputation: 221564

Use vectorized operations alongwith masking -

import numpy as np

# Convert to int datatype, as we want to avoid overflow later on
image_int = image.astype(int)

# Split up the channels
B,G,R = image_int.transpose(2,0,1)
# Or image_int[...,0],image_int[...,1],image_int[...,2]

# Use numpy.abs and 2D sliced data to get 2D mask
mask = mask = ((np.abs(G - B) > 20) & (G>30)) | (np.abs(R - G) < 40)

# Use the mask to assign zeros into o/p with boolean-indexing
output[mask] = [0,0,0]

Upvotes: 4

Related Questions