在 numpy 中向量化非平凡的 for 循环

Question

我正试图通过矢量化来使这段代码运行得更快，因为我相信 python 中的 for 循环很慢。我不完全理解矢量化，所以 for 循环内的切片给我带来了麻烦。

注意：这是针对不允许任何非 numpy 库的分配。

self.gx = numpy.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
self.gy = numpy.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

def create_vector(self, image):
    """Creates a gradient vector for each pixel in the image

    Returns: vec_data: [mag, angle, avg_value]"""
    vec_data = numpy.zeros((image.shape[0], image.shape[1], 3), dtype='float')

    for y in xrange(1, image.shape[0]-1):
        for x in xrange(1, image.shape[1]-1):

           #Get 3x3 matrix around the pixel
           subimage = image[y-1:y+2,x-1:x+2]

           #Apply sobel operator
           dx = (self.gx*subimage).sum()
           dy = (self.gy*subimage).sum()

           vec_data[y,x,0] = abs(dx) + abs(dy)
           vec_data[y,x,1] = abs(math.atan2(dx,dy))
           vec_data[y,x,2] = subimage.sum()/9 #Average of 3x3 pixels around x, y

    return vec_data

Answer 1

代码的直接矢量化是通过窗口视图完成的：

import numpy as np

image = np.arange(25).reshape(5, 5)
gx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
gm = np.ones((3, 3))/9
rows, cols = image.shape
k_rows, k_cols = gx.shape

from numpy.lib.stride_tricks import as_strided
image_view = as_strided(image, shape=(rows - k_rows + 1, cols - k_cols + 1,
                                      k_rows, k_cols),
                        strides=image.strides*2)
dx = np.einsum('ijkl,kl->ij',image_view, gx)
dy = np.einsum('ijkl,kl->ij',image_view, gy)
dm = np.einsum('ijkl,kl->ij',image_view, gm)

>>> dx
array([[-8, -8, -8],
       [-8, -8, -8],
       [-8, -8, -8]])
>>> dy
array([[-40, -40, -40],
       [-40, -40, -40],
       [-40, -40, -40]])
>>> dm
array([[  6.,   7.,   8.],
       [ 11.,  12.,  13.],
       [ 16.,  17.,  18.]])

从那些你可以构建你所追求的输出。

如果您遇到性能问题，您的卷积核是可分离的，即 2D 卷积可以沿正交轴拆分为 2 个 1D 卷积，这应该比上述解决方案运行得更快。

Answer 2

因此，实际上，您通过在移动窗口上进行乘法和求和所做的就是所谓的卷积。 Numpy/Scipy 在ndimage 模块中有这个功能。因此，您可以获得dx 和dy 的数组，而不仅仅是一个窗口。然后，您可以获得mag、ang 和avg 层以及适用于整个dx 和dy 数组的numpy 函数。因此，分别计算这些层中的每一层，然后返回三个事物的dstack，如果您希望将它们全部放在一个数组中：

import numpy as np
from scipy import ndimage

gx = np.array([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])

def create_vector(image):
    """Creates a gradient vector for each pixel in the image

    Returns: vec_data: [mag, angle, avg_value]"""

    #Apply sobel operator using convolution
    dx = ndimage.convolve(gx, image)
    dy = ndimage.convolve(gy, image)

    vec_data_mag = np.abs(dx) + np.abs(dy)
    vec_data_ang = np.abs(np.arctan2(dy, dx))    # are you sure you want abs here?
    vec_data_avg = ndimage.convolve(np.ones(3,3), image)

    return np.dstack([vec_data_mag, vec_data_angl, vec_data_avg])