Correcting coordinates of an image with polynomial root in python

Question

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)
dst[::]=((255,0,0))

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])
        else:
            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]]

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.