Исправление координат изображения с корнем полинома в Python

Я хочу оптимизировать этот цикл for для исправления координат изображения, это занимает слишком много времени, что не подходит для моей системы. Я провел некоторое профилирование, большую часть времени занимают корни numpy (около 90%). Может кто-нибудь предложить оптимизацию или векторизацию кода? Или лучшая альтернатива?

src = cv2.imread('distorted_JJ.bmp')
dist_center = np.array([822, 519])     
k1 = 0.45           
k2 = 0.82

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
        abs_ru = np.sqrt(ru.dot(ru))
        
        p = [k2 , 0, k1, 0, 1, -abs_ru]
        abs_rd = np.roots(p)
        abs_rd = abs_rd[~np.iscomplex(abs_rd)].real
        if i == yc and j == xc:
            rd = np.array([0,0])
        else:
            rd = ru * (abs_rd/abs_ru)

        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. Рассчитать abs_ru быстро, векторизованно.
2. Закажите эти реальные ценности.
3. Возьмите наименьшее, используйте numpy.roots найти соответствующий единственный действительный корень.
4. Исходя из этого, рассмотрим следующий abs_ru значение и соответствующий многочлен. Если это следующий abs_ru далеко, тогда подумайте об использовании numpy.roots опять же, в противном случае будет достаточно нескольких итерационных шагов с использованием итерации поиска корня. Рассмотрим градиентный спуск по квадрату многочлена или итерацию Ньютона.

Несколько вещей о коде как есть.

Поиск корня обычно выполняется численно, поэтому, возможно, лучше рассмотреть числовые корни, которые имеют реальное значение абсолютного значения ниже, скажем, 1e-5, а не ~np.iscomplex(abs_rd).

Поиск корня стоит дорого, и он излишен, если i == yc and j == xc. Переместите его в ветку else.