NumPy 三维数组的映射操作

Question

现在我有一个三维的np.array [height, weight, 3]，它是一张图片。我想要自己实现RGB -> YUV算法RGB2YUV，但是针对每个像素点进行转换的迭代太慢了。

for x in xrange(height):
    for y in xrange(weight):
          img[x,y] = mat_1 * img[x,y]

有没有一种使用内置方法来实现这个的方式？

- Lyken

Answer 1

这似乎是np.einsum的一个很好的应用案例：

yuv_image = np.einsum('kl,ijl->ijk', transformation_matrix, rgb_image)

一旦你在纸上写下来，就很容易得出索引。

以下示例展示了两种方法的值相等性：

>>> rgb_image = np.random.rand(2*4*3).reshape(2,4,3)
>>> transformation_matrix = np.random.rand(9).reshape(3,3)
>>> z = np.empty_like(rgb_image)
>>> for x in range(rgb_image.shape[0]):
...     for y in range(rgb_image.shape[1]):
...         z[x,y] = np.dot(transformation_matrix, rgb_image[x,y,:])
...
>>> np.allclose(z, np.einsum('kl,ijl->ijk', transformation_matrix, rgb_image))
True

- Oliver W.