pythonnumpyarray-broadcastingnumpy-indexing

Advanced Indexing in 3 Dimensional Numpy ndarray In Python


I have a ndarray of shape (68, 64, 64) called 'prediction'. These dimensions correspond to image_number, height, width. For each image, I have a tuple of length two that contains coordinates that corresponds to a particular location in each 64x64 image, for example (12, 45). I can stack these coordinates into another Numpy ndarray of shape (68,2) called 'locations'.

How can I construct a slice object or construct the necessary advanced indexing indices to access these locations without using a loop? Looking for help on the syntax. Using pure Numpy matrixes without loops is the goal.

Working loop structure

Import numpy as np
# example code with just ones...The real arrays have 'real' data.
prediction = np.ones((68,64,64), dtype='float32')
locations = np.ones((68,2), dtype='uint32')

selected_location_values = np.empty(prediction.shape[0], dtype='float32')
for index, (image, coordinates) in enumerate(zip(prediction, locations)):
    selected_locations_values[index] = image[coordinates]

Desired approach

selected_location_values = np.empty(prediction.shape[0], dtype='float32')

correct_indexing = some_function_here(locations). # ?????
selected_locations_values = predictions[correct_indexing]

Solution

  • A straightforward indexing should work:

    img = np.arange(locations.shape[0])
    r = locations[:, 0]
    c = locations[:, 1]
    selected_locations_values = predictions[img, r, c]
    

    Fancy indexing works by selecting elements of the indexed array that correspond to the shape of the broadcasted indices. In this case, the indices are quite straightforward. You just need the range to tell you what image each location corresponds to.