Python - More efficient way to loop through matrices

Question

I have

• an (height x width x 3) RGB image img
• a (256 x 256 x 256) matrix color_mapping mapping colors
• a (height x width) matrix pixel_result which I want to fill with the elements mapped with color_mapping for each pixel of the image

How can I fill that last matrix efficiently?

So far, I simply run through all the pixels in for loops, but I'm wondering if there is a more efficient way to do it

for i in range(width):
    for j in range(height):
        pixel_result[j,i] = color_mapping[img[j,i,0],img[j,i,1],img[j,i,2]]

Answer 1

You can implement this by subscripting the color_mapping:

pixel_result = color_mapping[img[:,:,0], img[:,:,1], img[:,:,2]]

For example if test the performance with:

>>> def f():
...   for i in range(width):
...     for j in range(height):
...         pixel_result[j,i] = color_mapping[img[j,i,0],img[j,i,1],img[j,i,2]]
... 
>>> def g():
...     pixel_result = color_mapping[img[:,:,0], img[:,:,1], img[:,:,2]]
>>> color_mapping = np.random.randn(256,256,256)

Then for a small image: img = np.random.randint(0, 256, (12, 7, 3)), we get:

>>> timeit(f, number=10000)
0.5247259399620816
>>> timeit(g, number=10000)
0.032307324931025505

For larger images, like img = np.random.randint(0, 256, (1920, 1080, 3)), we get:

>>> timeit(f, number=10)
18.934733690926805
>>> timeit(g, number=10)
0.5807857210747898

or for a 72×128 image:

>>> timeit(f, number=100)
0.6014469779329374
>>> timeit(g, number=100)
0.011570235947147012