CODEWITHSUNDEEP
pythonnumpy
cjm2671
cjm2671

Reputation: 19466

Vectorize a numpy routine

I've got the following routine I've written that takes two arbitrary curves and warps the space between them so it fits between two straight lines. For the loop, it process it per column as np.linspace doesn't operate on vectors AFAIK. Is there way to get rid of this loop and hit the whole thing at once?

def warp(img, iris_curve, pupil_curve):
  height, width = img.shape[:2]
  iris_height = np.uint8(np.max(np.int8(iris_curve) - pupil_curve))
  out = np.zeros((iris_height, width))

  for r in range(0,width):
    map_theta = np.linspace(pupil_curve[r], iris_curve[r], iris_height)
    map_theta = np.uint8(np.rint(map_theta))
    out[:, r] = img[map_theta, r]

  return np.uint8(out)

Thanks!

Upvotes: 1

Views: 81

Answers (2)

Jaime
Jaime

Reputation: 67427

If you peek into the source code of np.linspace, you can use that as a guide to vectorize your code. Your loop would then be replaced by something like:

steps = (iris_curve - pupil_curve) / iris_height
map_theta = np.arange(iris_height)[:, None] * steps + pupil_curve
map_theta = np.rint(map_theta).astype(np.uint8)
out = img[map_theta, np.arange(width)]

You may have to transpose the output, it's hard to debug this kind of code without an example.

Upvotes: 1

John Greenall
John Greenall

Reputation: 1690

My first thoughts looking at this code is that the for loop is unlikely to be the most significant thing hampering performace. There is an awful lot of copying of data around in this code (calling linspace creates a new array then this is copied back into the larger output array. casting to different types also involves copying data around). For example, can you not initiate your output as

np.empty((h,w),dtype=np.uint8)

Moreover do you really need to explicitly calculate all these values? Unless you reference all of them you might be better off just using the 2D linear interpolator from scipy.

If you really want to produce the output as-is I think you'll have to write something custom in Cython or similar.

Upvotes: 0

Related Questions