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?](http://stackoverflow.com/questions/14448763/is-there-a-convenient-way-to-apply-a-lookup-table-to-a-large-array-in-numpy) — Sergei Lebedev, Mar 11 '16 at 13:13

score 1 · Accepted Answer · answered Mar 11 '16 at 13:26

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 '16 at 13:26

unutbu

842,883
184
1,785
1,677

I missed the obvious! – Divakar Mar 11 '16 at 13:27
@Divakar: Bully for me this time, but I often learn from your inventive (and fast) solutions. Thanks! – unutbu Mar 11 '16 at 13:33
@unutbu That's news for me!! Glad to hear that really, because it happens the other way too :) – Divakar Mar 11 '16 at 13:36

Divakar · Answer 2 · 2016-03-11T13:24:23.733

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 '16 at 13:24

answered Mar 11 '16 at 13:19

Divakar

218,885
19
262
358

Thank's @Divakar for your time and solution! – hetsch Mar 11 '16 at 13:30

Optimizing or replacing array iteration with python loop by numpy functionality

2 Answers2