Thresholded pixel indices of a NumPy array

Question

I'm sure this question is Googleable, but I don't know what keywords to use. I'm curious about a specific case, but also about how to do it in general. Lets say I have a RGB image as an array of shape (width, height, 3) and I want to find all the pixels where the red channel is greater than 100. I feel like image > [100, 0, 0] should give me an array of indices (and would if I was comparing a scalar and using a greyscale image) but this compares each element with the list. How do I compare over the first two dimensions where each "element" is the last dimension?

Answer 1

To detect for red-channel only, you can do something like this -

np.argwhere(image[:,:,0] > threshold)

Explanation :

1. Compare the red-channel with the threshold to give us a boolean array of same shape as the input image without the third axis (color channel).
2. Use np.argwhere to get the indices of successful matches.

For a case when you want to see if any channel is above some threshold, use .any(-1) (any elements that satisfy the condition along the last axis/color channel).

np.argwhere((image > threshold).any(-1))

Sample run

Input image :

In [76]: image
Out[76]: 
array([[[118,  94, 109],
        [ 36, 122,   6],
        [ 85,  91,  58],
        [ 30,   2,  23]],

       [[ 32,  47,  50],
        [  1, 105, 141],
        [ 91, 120,  58],
        [129, 127, 111]]], dtype=uint8)

In [77]: threshold
Out[77]: 100

Case #1: Red-channel only

In [69]: np.argwhere(image[:,:,0] > threshold)
Out[69]: 
array([[0, 0],
       [1, 3]])

In [70]: image[0,0]
Out[70]: array([118,  94, 109], dtype=uint8)

In [71]: image[1,3]
Out[71]: array([129, 127, 111], dtype=uint8)

Case #2: Any-channel

In [72]: np.argwhere((image > threshold).any(-1))
Out[72]: 
array([[0, 0],
       [0, 1],
       [1, 1],
       [1, 2],
       [1, 3]])

In [73]: image[0,1]
Out[73]: array([ 36, 122,   6], dtype=uint8)

In [74]: image[1,1]
Out[74]: array([  1, 105, 141], dtype=uint8)

In [75]: image[1,2]
Out[75]: array([ 91, 120,  58], dtype=uint8)

---

Faster alternative to np.any in np.einsum

np.einsum could be tricked to perform np.any's work and as it turns out is a tad faster.

Thus, boolean_arr.any(-1) would be equivalent to np.einsum('ijk->ij',boolean_arr).

Here are the associated runtimes across various datasizes -

In [105]: image = np.random.randint(0,255,(30,30,3)).astype('uint8')
     ...: %timeit np.argwhere((image > threshold).any(-1))
     ...: %timeit np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: out1 = np.argwhere((image > threshold).any(-1))
     ...: out2 = np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: print np.allclose(out1,out2)
     ...: 
10000 loops, best of 3: 79.2 µs per loop
10000 loops, best of 3: 56.5 µs per loop
True

In [106]: image = np.random.randint(0,255,(300,300,3)).astype('uint8')
     ...: %timeit np.argwhere((image > threshold).any(-1))
     ...: %timeit np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: out1 = np.argwhere((image > threshold).any(-1))
     ...: out2 = np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: print np.allclose(out1,out2)
     ...: 
100 loops, best of 3: 5.47 ms per loop
100 loops, best of 3: 3.69 ms per loop
True

In [107]: image = np.random.randint(0,255,(3000,3000,3)).astype('uint8')
     ...: %timeit np.argwhere((image > threshold).any(-1))
     ...: %timeit np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: out1 = np.argwhere((image > threshold).any(-1))
     ...: out2 = np.argwhere(np.einsum('ijk->ij',image>threshold))
     ...: print np.allclose(out1,out2)
     ...: 
1 loops, best of 3: 833 ms per loop
1 loops, best of 3: 640 ms per loop
True