Now I have a 3-dimension np.array [height, weight, 3]. (It is a image) And I want to implement an RGB -> YUV algorithm myself RGB2YUV. However, iterating from each pixel and applying the transform is too slow.
for x in xrange(height):
for y in xrange(weight):
img[x,y] = mat_1 * img[x,y]
Is there any way to use some built_in method implementing this?
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 the indices, once you've written it down on a piece of paper.
Example to show value equality 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
