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?

Solution

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).

Comments

[fish2000]

That worked flawlessly, thanks!... Incedentally, do you know if this method is vectorizable -- like, will it work on a list of RGB triples, or can it be adapted to do so?

[Henry Gomersall]

As it stands, it should work on any array in which the last axis contains the colour information. So you can create, for example, a 4-D array in which the first axis indexes the image - something like a = numpy.empty((5, 512, 512, 3), dtype='int'); a[0,:,:,:] = image0; a[1,:,:,:] = image1 .... indices will now contain an extra entry (at the beginning) which are the image indices. If you want to make the colour axis an axis other than the last, you will probably need to play around with the structure of c (raise its dimensionality or something).

[Rishit Bansal]

Instead of zip(...), numpy.transpose(indices) can be used as mentioned in the docs. This might slightly benefit performance....

[mLstudent33]

does this work as alternative syntax for the same thing?np.where((images[i] == (0, 255, 0)).all(axis=-1))