Optimizing or replacing array iteration with python loop by numpy functionality

Question

I have the following code that works like expected, but I'm curious if the loop can be replaced by a native numpy function/method for better performance. What I have is one array holding RGB values that I use as a lookup table and two 2d arrays holding greyscale values (0-255). Each value of these two arrays corresponds to the value of one axis of the lookup table.

As mentioned, what would be really nice is getting rid of the (slow) loop in python and using a faster numpy method.

#!/usr/bin/env python3
from PIL import Image
import numpy as np

dim = (2000, 2000)
rows, cols = dim

# holding a 256x256 RGB color lookup table
color_map = np.random.random_integers(0, 255, (256,256,3))
# image 1 greyscale values
color_map_idx_row = np.random.randint(0, 255, dim)
# image 2 greyscale values
color_map_idx_col = np.random.randint(0, 255, dim)

# output image data
result_data = np.zeros((rows, cols, 3), dtype=np.uint8)

# is there any built in function in numpy that could
# replace this loop?
# -------------------------------------------------------

for i in range(rows):
    for j in range(cols):
        row_idx = color_map_idx_row.item(i, j)
        col_idx = color_map_idx_col.item(i, j)
        rgb_color = color_map[row_idx,col_idx]
        result_data[i,j] = rgb_color


img = Image.fromarray(result_data, 'RGB')
img.save('result.png')

Possible duplicate of Is there a convenient way to apply a lookup table to a large array in numpy? — Sergei Lebedev
– Sergei Lebedev, Commented Mar 11, 2016 at 13:13

unutbu · Accepted Answer · 2016-03-11 13:26:38Z

1

You can replace the double-for loop with fancy-indexing:

In [33]: result_alt = color_map[color_map_idx_row, color_map_idx_col]

This confirms the result is the same:

In [36]: np.allclose(result_data, result_alt)
Out[36]: True

answered Mar 11, 2016 at 13:26

unutbu

886k197 gold badges1.9k silver badges1.7k bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Divakar Over a year ago

I missed the obvious!

unutbu Over a year ago

@Divakar: Bully for me this time, but I often learn from your inventive (and fast) solutions. Thanks!

Divakar Over a year ago

@unutbu That's news for me!! Glad to hear that really, because it happens the other way too :)

Divakar · Accepted Answer · 2016-03-11 13:24:23Z

0

You can reshape the 3D array into a 2D array with the axis=1 holding the three channels. Then, use row-slicing with the row indices being calculated as linear indices from the row and column indices arrays. Please note that the reshaped array being a view only, won't burden any of the workspace memory. Thus, we would have -

m = color_map.shape[0]
out = color_map.reshape(-1,3)[color_map_idx_row*m + color_map_idx_col]

edited Mar 11, 2016 at 13:24

answered Mar 11, 2016 at 13:19

Divakar

222k19 gold badges273 silver badges374 bronze badges

1 Comment

hetsch Over a year ago

Thank's @Divakar for your time and solution!

Collectives™ on Stack Overflow

Optimizing or replacing array iteration with python loop by numpy functionality

2 Answers 2

3 Comments

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related