CODEWITHSUNDEEP
pythonoptimizationnumpy
Luca
Luca

Reputation: 10996

Can I vectorise this python code?

I have written this python code to get neighbours of a label (a set of pixels sharing some common properties). The neighbours for a label are defined as the other labels that lie on the other side of the boundary (the neighbouring labels share a boundary). So, the code I wrote works but is extremely slow:

# segments: It is a 2-dimensional numpy array (an image really)
# where segments[x, y] = label_index. So each entry defines the
# label associated with a pixel.

# i: The label whose neighbours we want.

def get_boundaries(segments, i):
    neighbors = []
    for y in range(1, segments.shape[1]):
        for x in range(1, segments.shape[0]):
            # Check if current index has the label we want 
            if segments[x-1, y] == i:
                # Check if neighbour in the x direction has
                # a different label
                if segments[x-1, y] != segments[x, y]:
                    neighbors.append(segments[x,y])

            # Check if neighbour in the y direction has
            # a different label
            if segments[x, y-1] == i:
                if segments[x, y-1] != segments[x, y]:
                    neighbors.append(segments[x, y])

    return np.unique(np.asarray(neighbors))

As you can imagine, I have probably completely misused python here. I was wondering if there is a way to optimize this code to make it more pythonic.

Upvotes: 2

Views: 135

Answers (2)

John Zwinck
John Zwinck

Reputation: 249153

Here you go:

def get_boundaries2(segments, i):
    x, y = np.where(segments == i) # where i is
    right = x + 1
    rightMask = right < segments.shape[0] # keep in bounds
    down = y + 1
    downMask = down < segments.shape[1]
    rightNeighbors = segments[right[rightMask], y[rightMask]]
    downNeighbors = segments[x[downMask], down[downMask]]
    neighbors = np.union1d(rightNeighbors, downNeighbors)
    return neighbors

As you can see, there are no Python loops at all; I also tried to minimize copies (the first attempt made a copy of segments with a NAN border, but then I devised the "keep in bounds" check).

Note that I did not filter out i itself from the "neighbors" here; you can add that easily at the end if you want. Some timings:

Input 2000x3000: original takes 13 seconds, mine takes 370 milliseconds (35x speedup).

Input 1000x300: original takes 643 ms, mine takes 17.5 ms (36x speedup).

Upvotes: 5

parchment
parchment

Reputation: 4002

You need to replace your for loops with numpy's implicit looping.

I don't know enough about your code to convert it directly, but I can give an example.

Suppose you have an array of 100000 random integers, and you need to get an array of each element divided by its neighbor.

import random, numpy as np 
a = np.fromiter((random.randint(1, 100) for i in range(100000)), int)

One way to do this would be:

[a[i] / a[i+1] for i in range(len(a)-1)]

Or this, which is much faster:

a / np.roll(a, -1)

Timeit:

initcode = 'import random, numpy as np; a = np.fromiter((random.randint(1, 100) for i in range(100000)), int)'
timeit.timeit('[a[i] / a[i+1] for i in range(len(a)-1)]', initcode, number=100)
5.822079309000401
timeit.timeit('(a / np.roll(a, -1))', initcode, number=100)
0.1392055350006558

Upvotes: 0

Related Questions