CODEWITHSUNDEEP
pythonnumpy
hetsch
hetsch

Reputation: 1568

Optimizing or replacing array iteration with python loop by numpy functionality

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')

Upvotes: 0

Views: 655

Answers (2)

unutbu
unutbu

Reputation: 879551

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

Upvotes: 1

Divakar
Divakar

Reputation: 221564

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]

Upvotes: 0

Related Questions