Numpy 3d array map operation
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