I want to optimize this for loop for correcting coordinates of an image, it takes too long which is not suited for my system. I have done some profiling, the numpy roots is taking most of the time (near to 90%). Could someone suggest some optimization or vectorization of the code? Or a better alternative?

src = cv2.imread('distorted_JJ.bmp')
dist_center = np.array([512, 224])     
k1 = 0.15           
k2 = 0.52

h,w,_ = src.shape
xc = dist_center[0]
yc = dist_center[1]

dst = np.zeros([h,w,3],dtype=np.uint8)

for i in range(h):
    for j in range(w):
        ru = np.array([j-xc, yc-i])/w            
        p = [k2 , 0, k1, 0, 1, ru]
        abs_rd = np.roots(p)
        if i == yc and j == xc:
            rd = np.array([0,0])
            rd = ru * (p/abs_rd)

        v = np.array([xc/w + rd[0], yc/w - rd[1]])
        v = v*w
        v = v.astype(int)
        dst[i][j] = src[v[1],v[0]]

1 Answer 1


This is more of a suggestion as I don't have the time to code it, but it is perhaps too long for a comment.

From your code it seems that it is assumed that this fifth degree polynomial has only a single real root. It must have one, because the complex ones must come in conjugate pairs, but I don't see why it has only one.

I will assume it has an unique real root. Here's the idea:

  1. Calculate abs_ru in a fast, vectorized manner.
  2. Order these real values.
  3. Take the smallest, use numpy.roots to find the corresponding unique real root.
  4. Starting from this, consider the next abs_ru value and the corresponding polynomial. If this next abs_ru is far, then consider using numpy.roots again, else a few iteration steps using a root finding iteration should suffice. Consider a gradient descent on the square of the polynomial, or a Newton iteration.

A few things on the code as it is.

Root finding is generally done numerically, therefore it is perhaps better to consider the numeric roots that have a real value of an absolute value below say 1e-5 instead of ~np.iscomplex(abs_rd).

Root finding is expensive, and it is redundant if i == yc and j == xc. Move it to the else branch.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.