Нужно ускорить очень медленный цикл для манипулирования изображениями на Python

Question

В настоящее время я заканчиваю программу в Pyhton (3.6) согласно внутренним требованиям. Как часть этого, мне приходится проходить через цветное изображение (3 байта на пиксель, R, G & B) и искажать пиксель изображения на пиксель.

У меня такой же код на других языках (C ++, C #), и неоптимизированный код выполняется примерно за две секунды, тогда как оптимизированный код выполняется менее чем за секунду. Под неоптимизированным кодом я подразумеваю, что умножение матриц выполняется 10-строчной функцией, которую я реализовал. Оптимизированная версия просто использует внешние библиотеки для умножения.

В Python этот код занимает около 300 секунд. Я не могу придумать, как векторизовать эту логику или ускорить ее, так как внутри вложенного цикла есть пара «если». Любая помощь будет принята с благодарностью.

import numpy as np

#for test purposes:
#roi = rect.rect(0, 0, 1200, 1200)
#input = DCImage.DCImage(1200, 1200, 3)
#correctionImage = DCImage.DCImage(1200,1200,3)
#siteToImage= np.zeros((3,3), np.float32)
#worldToSite= np.zeros ((4, 4))
#r11 = r12 = r13 = r21 = r22 = r23 = r31 = r32 = r33 = 0.0
#xMean = yMean = zMean = 0 
#tx = ty = tz = 0
#epsilon = np.finfo(float).eps
#fx = fy = cx = cy = k1 = k2 = p1 = p2 = 0


for i in range (roi.x, roi.x + roi.width):
    for j in range (roi.y , roi.y + roi.height):
        if ( (input.pixels [i] [j] == [255, 0, 0]).all()):
            #Coordinates conversion
            siteMat = np.matmul(siteToImage, [i, j, 1])
            world =np.matmul(worldToSite, [siteMat[0], siteMat[1], 0.0, 1.0])

            xLocal = world[0] - xMean
            yLocal = world[1] - yMean
            zLocal = z_ortho - zMean
            #From World to camera
            xCam = r11*xLocal + r12*yLocal + r13*zLocal + tx
            yCam = r21*xLocal + r22*yLocal + r23*zLocal + ty
            zCam = r31*xLocal + r32*yLocal + r33*zLocal + tz
            if (zCam > epsilon or zCam < -epsilon):
                xCam = xCam / zCam
                yCam = yCam / zCam

            #// DISTORTIONS                           
            r2 = xCam*xCam + yCam*yCam
            a1 = 2*xCam*yCam
            a2 = r2 + 2*xCam*xCam
            a3 = r2 + 2*yCam*yCam
            cdist = 1 + k1*r2 + k2*r2*r2
            u = int((xCam * cdist + p1 * a1 + p2 * a2) * fx + cx + 0.5)
            v = int((yCam * cdist + p1 * a3 + p2 * a1) * fy + cy + 0.5)

            if (u>=0 and u<correctionImage.width and v>=0 and v < correctionImage.height):
                input.pixels [i] [j] =  correctionImage.pixels [u][v]

Answer 1

Вы обычно векторизуете подобные вещи, создавая карту смещения.

Создайте сложное изображение, в котором каждый пиксель имеет значение своей собственной координаты, примените обычные математические операции для вычисления любого преобразования, которое вы хотите, а затем примените карту к исходному изображению.

Например,в pyvips вы можете написать:

import sys
import pyvips

image = pyvips.Image.new_from_file(sys.argv[1])

# this makes an image where pixel (0, 0) (at the top-left) has value [0, 0],
# and pixel (image.width, image.height) at the bottom-right has value
# [image.width, image.height]
index = pyvips.Image.xyz(image.width, image.height)

# make a version with (0, 0) at the centre, negative values up and left,
# positive down and right
centre = index - [image.width / 2, image.height / 2]

# to polar space, so each pixel is now distance and angle in degrees
polar = centre.polar()

# scale sin(distance) by 1/distance to make a wavey pattern
d = 10000 * (polar[0] * 3).sin() / (1 + polar[0])

# and back to rectangular coordinates again to make a set of vectors we can
# apply to the original index image
distort = index + d.bandjoin(polar[1]).rect()

# distort the image
distorted = image.mapim(distort)

# pick pixels from either the distorted image or the original, depending on some
# condition
result = (d.abs() > 10 or image[2] > 100).ifthenelse(distorted, image)

result.write_to_file(sys.argv[2])

Это просто глупый шаблон колебания, но вы можете поменять его на любое искажение.Затем запустите на этом двухъядерном ноутбуке 2015 года:

$ /usr/bin/time -f %M:%e ./wobble.py ~/pics/horse1920x1080.jpg x.jpg
54572:0.31

300 мс и 55 МБ памяти, чтобы сделать:

Answer 2

После долгих испытаний единственным способом ускорить выполнение функции без ее написания на C ++ была ее разборка и векторизация.Способ сделать это в этом конкретном случае - создать массив с действительными индексами в начале функции и использовать их в качестве кортежей для индексации окончательного решения.

subArray[roi.y:roi.y+roi.height,roi.x:roi.x+roi.width,] = input.pixels[roi.y:roi.y+roi.height,roi.x:roi.x+roi.width,]

#Calculate valid XY indexes
y_index, x_index = np.where(np.all(subArray== np.array([255,0,0]), axis=-1))

#....
#do stuff
#....

#Join result values with XY indexes
ij_xy = np.column_stack((i, j, y_index, x_index))

#Only keep valid ij values
valids_ij_xy = ij_xy [(ij_xy [:,0] >= 0) & (ij_xy [:,0] < correctionImage.height) & (ij_xy [:,1] >= 0) & (ij_xy [:,1] < correctionImage.width)]

#Assign values
input.pixels [tuple(np.array(valids_ij_xy [:,2:]).T)] = correctionImage.pixels[tuple(np.array(valids_ij_xy [:,:2]).T)]