Working with a map 3d map

Question

Now I have a 3-dimensional np.array [height, weight, 3]. (This image) And I want to implement the RGB → YUV algorithm of RGB2YUV itself . However, repeating from each pixel and applying the conversion is too slow.

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

Is there any built_in method that implements this?

Answer 1

This seems like a good use case for np.einsum :

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

It's easy to come up with indexes as soon as you write them down on a piece of paper.

An example to display the equality of values ​​of both approaches:

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