Индексирование массива 4D с индексированием массива по двум последним осям - NumPy / Python

Question

Я хочу создать серию нулевых изображений с несколькими каналами и с некоторым заданным пикселем на изображение со значением один.

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

num_channels = 3
im_size = 2
images = np.zeros((num_channels, im_size, im_size))

# random locations for the ones
pixels = np.random.randint(low=0, high=im_size,
                           size=(num_channels, 2))
images[np.arange(num_channels), pixels[:, 0], pixels[:, 1]] = 1

Однако аналогичный код не работает, если мы хотим рассмотреть пакет также:

batch_size = 4
num_channels = 3
im_size = 2
images = np.zeros((batch_size, num_channels, im_size, im_size))

# random locations for the ones
pixels = np.random.randint(low=0, high=im_size,
                           size=(batch_size, num_channels, 2))
images[np.arange(batch_size), np.arange(num_channels), pixels[:, :, 0], pixels[:, :, 1]] = 1

, который дает ошибку

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (4,) (3,) (4,3) (4,3) 

Следующий код выполнит работу, используя неэффективный цикл:

batch_size = 4
num_channels = 3
im_size = 2
images = np.zeros((batch_size, num_channels, im_size, im_size))

# random locations for the ones
pixels = np.random.randint(low=0, high=im_size,
                       size=(batch_size, num_channels, 2))
for k in range(batch_size):
    images[k, np.arange(num_channels), pixels[k, :, 0], pixels[k, :, 1]] = 1

Как бы вы получили векторизованное решение?

Answer 1

Простая векторизация с использованием advanced-indexing будет -

I,J = np.arange(batch_size)[:,None],np.arange(num_channels)
images[I, J, pixels[...,0], pixels[...,1]] = 1

Альтернативный более простой способ получить эти индексаторы I, J был бы с np.ogrid -

I,J = np.ogrid[:batch_size,:num_channels]