How to go through an array apply a threshold to each pixel

Viewed 579

I have a 3 channel numpy array and I would like to apply a function to each pixel. Specifically I want to process an image and return a greyscale image highlighting where specific colours appear in the image. If the red, green, blue channels are within 10 in L2 distance from the colour: (30,70,130) then set that pixel's value on the greyscale image to be 255, otherwise 0.

My current process to doing it is with:

def L2_dist(p1,p2):
    dist = ( (p1[0]-p2[0] )**2 + (p1[1]-p2[1] )**2 + (p1[2]-p2[2] )**2 ) **0.5
    if dist<10: return 255
    return 0

def colour_img(image):
    colour = my_colour
    img_dim = image.shape
    new_img = np.zeros((img_dim[0],img_dim[1])) # no alpha channel replica
    for c in range(img_dim[0]):
        for r in range(img_dim[1]):
            pixel = image[r,c,:3]
            new_img[r,c] = L2_dist(colour,pixel)
    return new_img

But it's very slow. How can I do this faster instead of using loops?

2 Answers

Simple one line solution

You can do what you want in a single line like this:

new_img = (((image - color)**2).sum(axis=2)**.5 <= 10) * 255

Optimized two line solution

The above line isn't the most efficient way possible to perform all of the operations that the OP wants. Here's a significantly faster way (credit to Paul Panzer for suggesting the optimizations in the comments, readability not guaranteed):

d = image - color
new_img = (np.einsum('...i, ...i', d, d) <= 100) * 255

Timings:

Given some test data with 100x100 pixels:

import numpy as np

color = np.array([30, 70, 130])
# random data within [20,60,120]-[40,80,140] for demo purposes
image = np.random.randint(10*2 + 1, size=[100,100,3]) + color - 10

Here's a comparison of the timings of the OP's method and the solutions from this answer. The one-line solution is about 100x faster than than the OP's, whereas the fully optimized version is about 300x faster:

%%timeit
# OP's code
img_dim = image.shape
new_img = np.zeros((img_dim[0],img_dim[1])) # no alpha channel replica
for c in range(img_dim[0]):
    for r in range(img_dim[1]):
        pixel = image[r,c,:3]
        new_img[r,c] = L2_dist(color,pixel)

43.8 ms ± 502 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
# one line solution
new_img = (((image - color)**2).sum(axis=2)**.5 <= 10) * 255

439 µs ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
# fully optimized solution
d = image - color
new_img = (np.einsum('...i, ...i', d, d) <= 100) * 255

145 µs ± 2.29 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Explanation of simple one line solution

The simple one-liner given as the first solution will:

  • Find the Euclidean distance between every pixel in image (which will be an array of shape (m, n, 3)) and color (which will be an array of shape (3)).

  • Check if any of those distances is within 10, and return a boolean array that is True wherever the condition is met and False otherwise.

  • A boolean array is really just an array of 0s and 1s, so we then multiply the boolean array by 255 to get the final result you wanted.

Explanation of optimized solution

Here's the list of optimizations used:

  • Uses einsum to calculate the sum of the squares required for the distance calculation. Under the hood, einsum makes use of the BLAS library that Numpy wraps to calculate the needed sum-product, so it should be faster.

  • Skips taking the square root by comparing the square of the distance to the square of the threshold.

  • I tried to find a way to minimize allocation/copying of arrays, but this actually made things slower. Here's a version of the optimized solution that allocates exactly two arrays (one for intermediate results and one for final result) and makes no other copies:

    %%timeit
    # fully optimized solution, makes as few array copies as possible
    scr = image.astype(np.double)
    new_img = np.zeros(image.shape[:2], dtype=np.uint8)
    np.multiply(np.less_equal(np.einsum('...i,...i', np.subtract(image, color, out=scr), scr, out=scr[:,:,0]), 100, out=new_img), 255, out=new_img)
    
    232 µs ± 7.72 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
    

You can do something like this

color = np.array([30, 70, 130])
L2 = np.sqrt(np.sum((image - color) ** 2, axis=2))  # L2 distance of each pixel from color

img_dim = image.shape
new_img = np.zeros((img_dim[0], img_dim[1]))
new_img[L2 < 10] = 255

But as you can see, we are iterating though the array twice, first to calculate L2 and then to do the thresholding in L2 < 10, We can improve it as is done in your code, through nested loops. But, loops in python are slow. So, JIT compile the function to get the fastest version. Below I use numba:

import numba as nb

@nb.njit(cache=True)
def L2_dist(p1,p2):
    dist = (p1[0]-p2[0] )**2 + (p1[1]-p2[1] )**2 + (p1[2]-p2[2] )**2
    if dist < 100: return 255
    return 0

@nb.njit(cache=True)
def color_img(image):
    n_rows, n_cols, _ = image.shape
    new_img = np.zeros((n_rows, n_cols), dtype=np.int32)
    for c in range(n_rows):
        for r in range(n_cols):
            pixel = image[r, c, :3]
            new_img[r,c] = L2_dist(color,pixel)
    return new_img

Timings:

# @tel's fully optimised solution(using einsum to short circuit np to get to BLAS directly, the no sqrt trick)
128 µs ± 6.94 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# JITed version without the sqrt trick
30.8 µs ± 10.2 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

# JITed version with the sqrt trick
24.8 µs ± 11.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

HTH.

Related