getting indices when comparing multidimensional ar

问题:

I have two numpy arrays, one an RGB image, one a lookup table of pixel values, for example:

img = np.random.randint(0, 9 , (3, 3, 3))
lut = np.random.randint(0, 9, (1,3,3))

What I'd like is to know the x,y coordinate in lut of pixels whose values are common to img and lut, so I tried:

for x in xrange(img.shape[0]):
    for y in xrange(img.shape[1]):
            print np.transpose(np.concatenate(np.where(lut == img[x,y])))

At this point, the problem is that img[x,y], which will be in the form of [int_r, int_g, int_b] does not get evaluated as a single element, so the three components get sought for separately in img...

I would like the output to be something like:

(x_coord, y_coord)

But I only get output in the form of:

[0 0 0]
[0 2 1]
[0 0 2]
[0 0 0]
[0 0 0]
[0 0 2]
[0 0 1]
[0 2 2]
[0 1 2]

Can anyone please help? Thanks!

回答1:

img = np.random.randint(0, 9 , (3, 3, 3))
lut2 = img[1,2,:] # so that we know exactly the answer

# compare two matrices
img == lut2

array([[[False, False, False],
        [False, False, False],
        [False,  True, False]],

       [[False, False, False],
        [False, False, False],
        [ True,  True,  True]],

       [[ True, False, False],
        [ True, False, False],
        [False, False, False]]], dtype=bool)

# rows with all true are the matching ones
np.where( (img == lut2).sum(axis=2) == 3 )

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

I don't really know why lut is filled with random numbers. But, I assume that you want to look for the pixels that have the exactly same color. If so, this seems to work. Is this what you need to do?

回答2:

@otterb 's answer works if lut is defined as a single [r,g,b] pixel slice, but it needs to be tweaked a little if you want to generalize this process to a multi-pixel lut:

img = np.random.randint(0, 9 , (3, 3, 3))
lut2 = img[0:1,0:2,:]

for x in xrange(lut2.shape[0]):
    for y in xrange(lut2.shape[1]):
        print lut2[x,y]
        print np.concatenate(np.where( (img == lut2[x,y]).sum(axis=2) == 3 ))

yields:

[1 1 7]
[0 0]
[8 7 4]
[0 1]

where triplets are pixel values, and doublets are their coordinates in the lut.

Cheers, and thanks to @otterb!

PS: iteration over numpy arrays is bad. The above is not production code.