Finding the (x,y) indexes of specific (R,G,B) color values from images stored in NumPy ndarrays

Question

I have an image stored in a numpy array, as yielded by imread():

>>> ndim
array([[[  0,   0,   0],
        [  4,   0,   0],
        [  8,   0,   0],
        ..., 
        [247,   0,  28],
        [251,   0,  28],
        [255,   0,  28]],

       [[  0, 255, 227],
        [  4, 255, 227],
        [  8, 255, 227],
        ..., 
        [247, 255, 255],
        [251, 255, 255],
        [255, 255, 255]]], dtype=uint8)
>>> ndim.shape
(512, 512, 3)

I want to efficiently find the (x, y) coordinate (or coordinates) of pixels with a specific color value, e.g.

>>> c
array([ 32,  32, 109], dtype=uint8)

>>> ndim[200,200]
array([ 32,  32, 109], dtype=uint8)

>>> ndim.T[0, 200, 200]
32
>>> ndim.T[1, 200, 200]
32
>>> ndim.T[2, 200, 200]
109

... in this case, I know the pixel at (200, 200) has the RGB value (32, 32, 109) -- I can test for this.

What I want to do is query the ndarray for a pixel value and get back the coordinates. In the above case, the putative function find_pixel(c) would return (200, 200).

Ideally this find_pixel() function would return a list of coordinate tuples and not just the first value it finds.

I've looked at numpy's "fancy indexing", which confused me greatly... Most of my attempts at figuring this out have been overwrought and unnecessarily baroque.

I am sure there is a very simple method that I am overlooking here. What is the best way to do this -- is there an altogether better mechanism to get these values than that which I have outlined?

Answer 1

For some array colour array a and a colour tuple c:

indices = numpy.where(numpy.all(a == c, axis=-1))

indices should now be a 2-tuple of arrays, the first of which contains the indices in the first dimensions and the second of which contains the indices in the second dimension corresponding to pixel values of c.

If you need this as a list of coordinate tuples, use zip:

coords = zip(indices[0], indices[1])

For example:

import numpy
a = numpy.zeros((4, 4, 3), 'int')    

for n in range(4):
    for m in range(4):
        a[n, m, :] = n + m
        if (n + m) == 4:
            print n, m

c = (4, 4, 4)
indices = numpy.where(numpy.all(a == c, axis=-1))
print indices
print zip(indices[0], indices[1])

will output:

1 3
2 2
3 1
(array([1, 2, 3]), array([3, 2, 1]))
[(1, 3), (2, 2), (3, 1)]

which corresponds to all the pixels of value (4, 4, 4).