extracting index from a 2D array Python

Question

I have an image (named gray_image ) of shape (1830, 1830) . After some image processing (I have created superpixels) I got a 2D array named segments (of shape (1830, 1830) ), which contains values from 0 to 72.

I need to take the index from where I find the value "0" in segments and use that index to save the value from gray_image in a new array (named: arr ).

I think an example will help you understand my problem better:

Let's say I have this image, an 2D array with 3x3 dimension:

gray_image = numpy.array([[1, 1, 1],
                          [2, 2, 2],
                          [3, 3, 3]])

And this is my segments array, an 2D array with 3x3 dimension:

segments = numpy.zeros([[0, 0, 1],
                        [0, 1, 2],
                        [1, 2, 2]])

I need to create an algorithm that outputs the array arr with shape (3, 3) .

 arr = np.array([[1, 1, 2]  # on the first line are the values from gray_image that correspond with value 0 from segments
                 [1, 2, 3]  # on the second line are the values from gray_image that correspond with value 1 from segments
                 [2, 3, 3]]) # on the third line are the values from gray_image that correspond with value 2 from segments

I just realized that the count of numbers from each segment are different, so I'm not sure if this is possible with and 2D array. I'm thinking of using a collection such as a dictionary to save all the information related to the index.

So, here is what I did so far:

i = 0
j = 0
k = 0
n = 0
m = 0
arr = np.empty([1830, 1830]) # empty array
for k in range(0, 72):
     for i in range(0,1829):
         for j in range(0,1829):                      
             if segments[i][j] == k:
                 arr[m][n] = gray_image[i][j]
                 n = n + 1
             if i == 1829 and j == 1829:
                 m = m + 1

But this isn't working at all, I've got this error:

arr[m][n] = gary_image[i][j]

IndexError: index 1830 is out of bounds for axis 0 with size 1830

I'm kinda stuck on this for a few days so any suggestion will be highly appreciated.

Answer 1

First, let's assume that your labels array is conducive to making a numpy array to begin with, ie, that number of elements per label N is constant and number of labels M x N is the same as the image size. If that is not the case, you can't construct a 2D numpy array as the result.

The trick is to identify the regions. For this, we will use np.argsort :

idxa = np.argsort(segments, axis=None).reshape(m, n)
arr = gray_image.ravel()[idxa]

If your labels are not conducive to an array output, the result above still useful. Instead of reshaping idx to the correct output shape, leave it as-is and figure out split indices to make a list of arrays:

idxl = np.argsort(segments, axis=None)
splits = np.flatnonzero(np.diff(segments.ravel()[idxl])) + 1
lst = np.split(gray_image.ravel()[idxl], splits)

The label corresponding to each segment does not have to be zero-based or in any way special, so you may want to get the values as an array of the same length as arr / lst :

labels = segments.ravel()[idxa[:, 0]]  # For arr-based solution
labels = segments.ravel()[idxl[np.r_[0, splits]]] # For lst-based solution

You can trivially turn the result into a dictionary by zipping labels and either arr or lst :

dct = dict(zip(labels, lst))  # Works with arr too