Python Numpy array indexing and masking

Question

I need help indexing a 3 dimensional array (an RGB/BGR image) using a 2 dimensional array of indices. All values are either 0, 1, or 2 for different color channels. The result should be a 2D array of color values. If anyone can tell me the syntax for this in python that would be great!

For a context of what I am trying to do (also read TLDR below):

I am essentially trying to convert the following code from normal for loop syntax which is very very slow to a more efficient python/numpy syntax:

colorIndices = np.zeros((height,width));       # an array which has the index of the outstanding color
colorIndices -= 1;           # all -1's

for x in range(0,width):
    for y in range(0,height):

        pix = img[y,x];        # get the pixel, a 1D array of length 3
        colorID = np.argmax(pix);           #get which index has max value (candidate for outstanding color)

        if(pix[colorID]>np.average(pix)+np.std(pix)):      # if that value is more than one std dev away from the overall pixel's value, the pixel has an outstanding color
            colorIndices[y,x] = colorID;

I then would like to access the outstanding color channels in each pixel using something like:

img[:,:,:]=0;
img[colorIndices] = 255;

TLDR: I want to set a pixel to pure blue, green, or red if it is a shade of that color. The way I define if a pixel is shade red is if the R value of the pixel is more than one std above the average of the overall distribution of {R, G, B}.

The broken code I have so far:

  colorIDs = np.argmax(img, axis=2);

  averages = np.average(img, axis=2);
  stds = np.std(img, axis=2);
  cutoffs = averages + stds;

  print(img[colorIDs]);

Mad Physicist · Accepted Answer

You can convert a numerical index along a given axis into values, you can use np.take_along_axis or fancy indexing. When using a fancy index, you need to index along all axes with arrays whose shapes broadcast to the size of the final result. np.ogrid helps with this. For an MxNx3 array img (M, N, _ = img.shape), if you had ix = np.argmax(img, axis=2), the index would look like this:

r, c = np.ogrid[:M, :N]
maxes = img[r, c, ix]

Using take_along_axis saves you a step and some temp arrays:

maxes = np.take_along_axis(img, ix, 2)

Now create your mask:

significant = np.abs(maxes - img.mean(axis=2) > img.std(axis=2))

At this point you have a 2D boolean mask and an integer index in the third dimension. The simplest thing is probably to turn everything into a linear index:

r, c = np.where(significant)

Now you can construct the output:

color_index = np.zeros_like(img)
color_index[r, c, ix[significant]] = 255

While tempting, np.put_along_axis can not be used here in a straightforward manner. The issue is that masking ix with significant would invalidated its shape similarity. You could, however, create an intermediate 2D array containing 255 at the locations marked by significant, and use that with put_along_axis:

values = np.zeros(significant.shape, dtype=img.dtype)
values[significant] = 255
color_index = np.zeros_like(img)
np.put_along_axis(color_index, ix, values, 2)

All combined:

ix = np.argmax(img, axis=2)
significant = np.abs(np.take_along_axis(img, ix, 2) - img.mean(axis=2)) > img.std(axis=2)
color_index = np.zeros_like(img)
color_index[(*np.where(significant), ix[significant])] = 255

Python Numpy array indexing and masking

Answers (2)

Related Questions