Question

我有两个输入数组：尺寸为（i，j，k）的data_arr和尺寸为（i，j）的index_arr。 index_arr中的条目是[0，k-1]范围内的整数。我想创建一个尺寸为（i，j）的输出数组（output_arr），其中对于output_arr的每个元素，index_arr告诉我要选择哪个元素。

换句话说，output_arr [i，j] = data_arr [i，j，index_arr [i，j]]

很明显，我可以以两倍于for循环的速度完成此操作。我更喜欢使用智能索引的快照工具。目前，我能设计的最好的方法是创建两个额外的大小为（i，j）的2D矩阵。

下面是一个简单的MWE，其构图是使用标准的Bayer图案从RGB图像创建马赛克图像。我希望能够摆脱X_ind和Y_ind

import numpy as np
import time


if __name__ == '__main__':
    img_width = 1920
    img_height = 1080
    img_num_colours = 3

    red_arr = np.ones([img_height, img_width], dtype=np.uint16) * 10
    green_arr = np.ones([img_height, img_width], dtype=np.uint16) * 20
    blue_arr = np.ones([img_height, img_width], dtype=np.uint16) * 30

    img_arr = np.dstack((red_arr, green_arr, blue_arr))

    bayer_arr = np.ones([img_height, img_width], dtype=np.uint16)
    bayer_arr[0::2,0::2] = 0 # Red entries in bater patter
                             # Green entries are already set by np.ones intialisation
    bayer_arr[1::2,1::2] = 2 # blue entries in bayer patter
    print("bayer\n",bayer_arr[:8,:12], "\n")

    mosaiced_arr = np.zeros([img_height, img_width], dtype=np.uint16)
    Y_ind = np.repeat(np.arange(0, img_width).reshape(1, img_width), img_height, 0)
    X_ind = np.repeat(np.arange(0, img_height).reshape(img_height, 1), img_width, 1)

    start_time = time.time()
    demos_arr = img_arr[X_ind, Y_ind, bayer_arr]
    end_time = time.time()

    print(demos_arr.shape)
    print("demos\n",demos_arr[:8,:12], "\n")
    print("Mosaic took {:.3f}s".format(end_time - start_time))

编辑：正如@Georgy所指出的，这个问题与我在搜索字词中找不到的this one相似，因此，该帖子可能会充当该帖子的路标。尽管扁平索引算法不同，但我的答案仍然适用，因为我的维度顺序不同。上面的答案等同于另一个问题中的ogrid版本。实际上，ogrid可以通过替换以下代码来使用：

# Y_ind = np.repeat(np.arange(0, img_width).reshape(1, img_width), img_height, 0)
# X_ind = np.repeat(np.arange(0, img_height).reshape(img_height, 1), img_width, 1)
X_ind, Y_ind = np.ogrid[0:img_height, 0:img_width]

您可以像这样实现选择选项（仅限于32个选项之间的选择）：

start_time = time.time()
demos_arr = bayer_arr.choose((img_arr[...,0], img_arr[...,1], img_arr[...,2]))
end_time = time.time()

ogrid解决方案在我的计算机上运行的时间为12ms，choice解决方案在34ms内运行

Answer 1

您要numpy.take_along_axis：

output_arr = numpy.take_along_axis(data_arr, index_arr[:, :, numpy.newaxis], axis=2)
output_arr = output_arr[:,:,0]  # Since take_along_axis keeps the same number of dimensions

此功能是numpy 1.15.0中的新增功能。

https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html

请注意，data_arr和index_arr需要具有相同的维数。因此，您需要将index_array重塑为3维，然后再将结果重塑为2维。即：

start_time = time.time()
demos_arr = np.take_along_axis(img_arr, bayer_arr.reshape([img_height, img_width, 1]), axis=2).reshape([img_height, img_width])
end_time = time.time()

沿轴获取的计时结果与ogrid实现相同。

使用第二个矩阵将3D数组展平为2D数组以选择第三维元素

1 个答案: