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

Question

我有两个输入数组：data_arr 维度 (i,j,k) 和 index_arr 维度 (i,j)。 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 循环以极快的速度完成此操作。我更喜欢使用智能索引的更快捷的东西。目前我能设计的最好的方法是创建两个额外的 2D 大小为 (i,j) 的矩阵。

下面是一个简单的 MWE，它使用标准的拜耳模式从 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 所指出的，这个问题类似于，我没有在我的搜索词中找到它，所以也许这个 post 将作为那个 post 的标志。另一个 post 中的答案是适用的，尽管扁平化索引算法不同，因为我的维度顺序不同。上面的答案相当于另一个问题中的 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 解决方案运行时间为 12 毫秒，选择解决方案运行时间为 34 毫秒

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()

take along axis的计时结果与ogrid实现相同。

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

Flatten a 3D array to a 2D array using a second matrix to choose elements in third dimension

python

arrays

numpy

matrix-indexing

demosaicing