Question

我想创建具有多个通道的零图像批处理，并且每个图像的某个像素给定的值为1。

如果仅按通道数索引图像，则以下代码可以正常工作：

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]

使用沿最后两个轴的索引数组对4D数组进行索引-NumPy / Python

1 个答案: