Vectorizing nested for loop

user_10891160 0
Viewed 124

I need to vectorize two nested for loops and dont know how to do it. One is for gray scale images and one for color images. I want to filter an image with the kuwahara filter. The code you see below is the last step I need to vectorize to get a fast function.

  • the array img_kuwahara is in shape of mxn or mxnx3 (color image)
  • the array index_min is in shape of mxn
  • the array mean is in shape of 4xmxn (gray scale) or 3x4xmxn (color)

I need to get the right value out of the mean array into the img_kuwahara array.

as sample data you can use the following arrays:

index_min = np.array([[0, 1, 1, 2, 3],[3, 3, 2, 2, 2],[2, 3, 3, 0, 2],[0, 1, 1, 0, 3],[2, 1, 3, 0, 0]])

mean = np.random.randint(0, 256, size=(4,5,5)) (gray scale images)

mean = np.random.randint(0, 256, size=(3,4,5,5)) (color images)

row = 5, columns = 5

Thank you for your help

# Edit gray scale image
if len(image.shape) == 2:

    # Set result image
    img_kuwahara = np.zeros((row, columns), dtype=imgtyp)

    for k in range(0, row):
        for i in range(0, columns):
            img_kuwahara[k, i] = mean[index_min[k, i], k, i]



# Edit color image
if len(image.shape) == 3:

    # Set result image
    img_kuwahara = np.zeros((row, columns, 3), dtype=imgtyp)

    for k in range(0, row):
        for i in range(0, columns):
            img_kuwahara[k, i, 0] = mean[0][index_min[k, i], k, i]
            img_kuwahara[k, i, 1] = mean[1][index_min[k, i], k, i]
            img_kuwahara[k, i, 2] = mean[2][index_min[k, i], k, i]
vectorization for-loop python
1 Answers

The first loop can be vectorized using np.meshgrid:

j, i = np.meshgrid(range(columns), range(rows))
img_kuwahara = mean[index_min[i, j], i, j]

The second loop can be vectorized by using an additional np.moveaxis (assuming that mean is actually a 4D array in that case, not a list of 3D arrays; otherwise just convert it):

j, i = np.meshgrid(range(columns), range(rows))
img_kuwahara = np.moveaxis(mean, 0, -1)[index_min[i, j], i, j]

Alternatively to np.meshgrid you can also use np.mgrid (which supports a more natural syntax):

i, j = np.mgrid[:rows, :columns]
user_37672390
Related