Fast way to apply custom function to every pixel in image

Question

I'm looking for a faster way to apply a custom function to an image which I use to remove a blue background. I have a function that calculates the distance each pixel is from approximately the blue colour in the background. The original code with a loop looked like this:

def dist_to_blue(pix):
    rdist = 76 - pix[0]
    gdist = 150 - pix[1]
    bdist = 240 - pix[2]
    return rdist*rdist + gdist*gdist + bdist*bdist

imgage.shape #outputs (576, 720, 3)
for i, row in enumerate(image):
    for j, pix in enumerate(row):
        if dist_to_blue(pix) < 12000: image[i,j] = [255,255,255]

However this code takes around 8 seconds to run for this relatively small image. I've been trying to use numpy's "vectorize" function but that applies the function to every value individually. However I want to do it to every pixel aka not expand the z/rgb dimension

the only improvements I've come up with is replacing the for loops with the following:

m = np.apply_along_axis(lambda pix: (255,255,255) if dist_to_blue(pix) < 12000 else pix, 2, image)

Which runs in about 7 seconds which is still painfully slow. Is there something I'm missing that could speed this up to a reasonable execution time

Answer 1

This is just a shot in the dark but maybe precomputing some data would help? I don't know for sure but the table lookup may be faster than the add and multiply?

def square(x): # maybe there's a library function for this?
    return x*x

RDIST = [square(76 - r) for r in range(256)]
GDIST = [square(150 - g) for g in range(256)]
BDIST = [square(240 - b) for b in range(256)]

def dist_to_blue(pix):
    return RDIST[pix[0]] + GDIST[pix[1]] + BDIST[pix[2]]

I suspect too if you have a way to just get an array of pixels per row that might be faster, instead of indexing each individual pixel, but I don't know the libraries in play.

Answer 2

This should be a lil bit faster ... ;)

import numpy as np 
blue = np.full_like(image, [76,150,250])
mask = np.sum((image-blue)**2,axis=-1) < 12000
image[mask] = [255,0,255]

Here you're generating the ideal blue image, squaring the difference of the images pixel by pixel, then summing over the last axis (the rgb vectors) before generating a mask and using it to modify values in the original image.