Simple python animate for loops

advertisements

I am trying a simple nested for loop in python to scan a threshold-ed image to detect the white pixels and store their location. The problem is that although the array it is reading from is only 160*120 (19200) it still takes about 6s to execute, my code is as follows and any help or guidance would be greatly appreciated:

im = Image.open('PYGAMEPIC')

r, g, b = np.array(im).T
x = np.zeros_like(b)

height = len(x[0])
width = len(x)

x[r > 120] = 255
x[g > 100] = 0
x[b > 100] = 0              

row_array = np.zeros(shape = (19200,1))
col_array = np.zeros(shape = (19200,1))

z = 0
for i in range (0,width-1):
    for j in range (0,height-1):
        if x[i][j] == 255:
            z = z+1
            row_array[z] = i

            col_array[z] = j


First, it shouldn't take 6 seconds. Trying your code on a 160x120 image takes ~0.2 s for me.

That said, for good numpy performance, you generally want to avoid loops. Sometimes it's simpler to vectorize along all except the smallest axis and loop along that, but when possible you should try to do everything at once. This usually makes things both faster (pushing the loops down to C) and easier.

Your for loop itself seems a little strange to me-- you seem to have an off-by-one error both in terms of where you're starting storing the results (your first value is placed in z=1, not z=0) and in terms of how far you're looking (range(0, x-1) doesn't include x-1, so you're missing the last row/column-- probably you want range(x).)

If all you want is the indices where r > 120 but neither g > 100 nor b > 100, there are much simpler approaches. We can create boolean arrays. For example, first we can make some dummy data:

>>> r = np.random.randint(0, 255, size=(8,8))
>>> g = np.random.randint(0, 255, size=(8,8))
>>> b = np.random.randint(0, 255, size=(8,8))

Then we can find the places where our condition is met:

>>> (r > 120) & ~(g > 100) & ~(b > 100)
array([[False,  True, False, False, False, False, False, False],
       [False, False,  True, False, False, False, False, False],
       [False,  True, False, False, False, False, False, False],
       [False, False, False,  True, False,  True, False, False],
       [False, False, False, False, False, False, False, False],
       [False,  True, False, False, False, False, False, False],
       [False, False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False, False]], dtype=bool)

Then we can use np.where to get the coordinates:

>>> r_idx, c_idx = np.where((r > 120) & ~(g > 100) & ~(b > 100))
>>> r_idx
array([0, 1, 2, 3, 3, 5])
>>> c_idx
array([1, 2, 1, 3, 5, 1])

And we can sanity-check these by indexing back into r, g, and b:

>>> r[r_idx, c_idx]
array([166, 175, 155, 150, 241, 222])
>>> g[r_idx, c_idx]
array([ 6, 29, 19, 62, 85, 31])
>>> b[r_idx, c_idx]
array([67, 97, 30,  4, 50, 71])