Skip to content

Instantly share code, notes, and snippets.

@N-McA

N-McA/numpy_index_extractor.py

Created April 9, 2018 11:32
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save N-McA/64b729a545e31f543f1d08b99c1afaf5 to your computer and use it in GitHub Desktop.
Save N-McA/64b729a545e31f543f1d08b99c1afaf5 to your computer and use it in GitHub Desktop.
Download ZIP
Pull out an array of indexes in numpy, eg for lookup in image
Raw
numpy_index_extractor.py
def extract(idxs, target):
trailing_shape = target.shape[idxs.shape[1]:]
flat_idxs = np.ravel_multi_index(idxs.T, target.shape[:idxs.shape[1]])
result_shape = (idxs.shape[0], ) + trailing_shape
return target.reshape([-1, *trailing_shape])[flat_idxs].reshape(result_shape)
@N-McA
Copy link
Author

N-McA commented Apr 13, 2018

The ever-wise Eric Wieser points out that this is
target[tuple(idxs.T)]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment