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

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) 的二维矩阵。

下面是一个简单的 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指出的，这个问题类似于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 解决方案的运行时间为 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 reshape 为 3 维，然后将结果再次 reshape 为 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实现相同。