numpy get 2d array, где последний размер индексируется в соответствии с массивом 2d

Question

Я читал о индексировании numpy, но я не нашел то, что искал.

У меня есть изображение 288 * 384, где каждый пиксель может иметь маркировку в [0,15]. Он хранится в 3d (288,384,16) -образном массиве numpy im.

С im[:,:,1], можно, например, получить изображение, в котором все пиксели имеют метку 1.

У меня есть другой 2d массив labelling, (288 * 384) -образный, содержащий метку для каждого пикселя.

Как получить изображение, в котором каждый пиксель имеет соответствующий пиксель, используя некоторые умные нарезки?

Использование петель, которые бы:

result = np.zeros((288,384)) 
for x,y in zip(range(288), range(384)): 
    result[x,y] = im[x,y,labelling[x,y]] 

Но это, конечно, неэффективное.

Answer 1

Новый результат

Короткий результат

np.choose(labelling,im.transpose(2,0,1)) 

Старый Результат

Попробуйте

im[np.arange(288)[:,None],np.arange(384)[None,:],labelling] 

Он работает для следующих ситуа на:

import numpy 
import numpy.random 
import itertools 

a = numpy.random.randint(5,size=(2,3,4)) 
array([[[4, 4, 0, 0], 
     [0, 4, 1, 1], 
     [3, 4, 4, 2]], 

     [[4, 0, 0, 2], 
     [1, 4, 2, 2], 
     [4, 2, 4, 4]]]) 

b = numpy.random.randint(4,size=(2,3)) 
array([[1, 1, 0], 
     [1, 2, 2]]) 

res = a[np.arange(2)[:,None],np.arange(3)[None,:],b] 
array([[4, 4, 3], 
     [0, 2, 4]]) 

# note that zip is not doing what you expect it to do 
result = np.zeros((2,3)) 
for x,y in itertools.product(range(2),range(3)): 
    result[x,y] = a[x,y,b[x,y]] 

array([[4., 4., 3.], 
     [0., 2., 4.]]) 

Обратите внимание, что zip не делать то, что вы ожидаете

zip(range(2),range(3)) 
[(0, 0), (1, 1)] 

Возможно, вы имели в виду что-то вроде itertools.product

list(itertools.product(range(2),range(3))) 
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)] 

Ужасно глядя [:,None] и т.д. можно избежать с помощью numpy.ix_

xx,yy = np.ix_(np.arange(2), np.arange(3)) 

res = a[xx,yy,b]